package mathkit

import (
	"math"
	"strings"
	"testing"
)

func almostEqual(a, b, tol float64) bool { return math.Abs(a-b) <= tol }

func mustExpr(t *testing.T, src string) *Expr {
	t.Helper()
	e, err := ParseExpr(src)
	if err != nil {
		t.Fatalf("ParseExpr(%q): %v", src, err)
	}
	return e
}

func TestExprEval(t *testing.T) {
	cases := []struct {
		src  string
		vars map[string]float64
		want float64
	}{
		{"2 + 3 * 4", nil, 14},
		{"(2 + 3) * 4", nil, 20},
		{"2 ^ 3 ^ 2", nil, 512}, // right-associative: 2^(3^2) = 2^9
		{"-2 ^ 2", nil, -4},     // unary minus binds looser than ^: -(2^2)
		{"x * x + 2*x + 1", map[string]float64{"x": 3}, 16},
		{"sigmoid(0)", nil, 0.5},
		{"relu(-5)", nil, 0},
		{"relu(5)", nil, 5},
		{"max(3, 7)", nil, 7},
		{"sqrt(16)", nil, 4},
		{"pi", nil, math.Pi},
	}
	for _, c := range cases {
		e := mustExpr(t, c.src)
		got, err := e.Eval(c.vars)
		if err != nil {
			t.Fatalf("Eval(%q): %v", c.src, err)
		}
		if !almostEqual(got, c.want, 1e-9) {
			t.Errorf("Eval(%q) = %v, want %v", c.src, got, c.want)
		}
	}
}

func TestExprUnboundVariable(t *testing.T) {
	e := mustExpr(t, "x + 1")
	if _, err := e.Eval(nil); err == nil {
		t.Fatal("expected an error for an unbound variable, got nil")
	}
}

func TestMatMul(t *testing.T) {
	a := Matrix{{1, 2}, {3, 4}}
	b := Matrix{{5, 6}, {7, 8}}
	got, err := MatMul(a, b)
	if err != nil {
		t.Fatal(err)
	}
	want := Matrix{{19, 22}, {43, 50}}
	for i := range want {
		for j := range want[i] {
			if got[i][j] != want[i][j] {
				t.Errorf("MatMul[%d][%d] = %v, want %v", i, j, got[i][j], want[i][j])
			}
		}
	}
}

func TestMatMulShapeMismatch(t *testing.T) {
	a := Matrix{{1, 2, 3}}
	b := Matrix{{1, 2}}
	if _, err := MatMul(a, b); err == nil {
		t.Fatal("expected a shape-mismatch error")
	}
}

func TestDeterminantAndInverse(t *testing.T) {
	a := Matrix{{4, 7}, {2, 6}}
	det, err := Determinant(a)
	if err != nil {
		t.Fatal(err)
	}
	if !almostEqual(det, 10, 1e-9) {
		t.Errorf("det = %v, want 10", det)
	}
	inv, err := Inverse(a)
	if err != nil {
		t.Fatal(err)
	}
	// A * A^-1 should be the identity.
	prod, err := MatMul(a, inv)
	if err != nil {
		t.Fatal(err)
	}
	for i := 0; i < 2; i++ {
		for j := 0; j < 2; j++ {
			want := 0.0
			if i == j {
				want = 1
			}
			if !almostEqual(prod[i][j], want, 1e-9) {
				t.Errorf("A*A^-1[%d][%d] = %v, want %v", i, j, prod[i][j], want)
			}
		}
	}
}

func TestInverseSingular(t *testing.T) {
	a := Matrix{{1, 2}, {2, 4}} // row 2 = 2 * row 1
	if _, err := Inverse(a); err == nil {
		t.Fatal("expected a singular-matrix error")
	}
}

func TestDominantEigen(t *testing.T) {
	// [[2,0],[0,1]] has eigenvalues 2 and 1 — dominant is 2, eigenvector (1,0).
	a := Matrix{{2, 0}, {0, 1}}
	lambda, v, err := DominantEigen(a, 0)
	if err != nil {
		t.Fatal(err)
	}
	if !almostEqual(lambda, 2, 1e-6) {
		t.Errorf("eigenvalue = %v, want 2", lambda)
	}
	if !almostEqual(math.Abs(v[0]), 1, 1e-6) || !almostEqual(v[1], 0, 1e-6) {
		t.Errorf("eigenvector = %v, want ~(1,0)", v)
	}
}

func TestDerivative(t *testing.T) {
	// d/dx[x^2] at x=3 is 6.
	e := mustExpr(t, "x^2")
	d, err := Derivative(e, "x", map[string]float64{"x": 3}, 0)
	if err != nil {
		t.Fatal(err)
	}
	if !almostEqual(d, 6, 1e-4) {
		t.Errorf("derivative = %v, want 6", d)
	}
}

func TestGradient(t *testing.T) {
	// f(x,y) = x^2 + y^2 ; gradient at (3,4) is (6,8).
	e := mustExpr(t, "x^2 + y^2")
	g, err := Gradient(e, []string{"x", "y"}, map[string]float64{"x": 3, "y": 4}, 0)
	if err != nil {
		t.Fatal(err)
	}
	if !almostEqual(g[0], 6, 1e-4) || !almostEqual(g[1], 8, 1e-4) {
		t.Errorf("gradient = %v, want [6 8]", g)
	}
}

func TestIntegrate(t *testing.T) {
	// ∫0..1 x^2 dx = 1/3.
	e := mustExpr(t, "x^2")
	v, err := Integrate(e, "x", 0, 1, 100, nil)
	if err != nil {
		t.Fatal(err)
	}
	if !almostEqual(v, 1.0/3.0, 1e-9) {
		t.Errorf("integral = %v, want 1/3", v)
	}
}

func TestGradientDescentConvergesToMinimum(t *testing.T) {
	// f(x) = (x-3)^2 has its minimum at x=3.
	e := mustExpr(t, "(x - 3)^2")
	res, err := GradientDescent(e, []string{"x"}, []float64{0}, 0.1, 200, 1e-10)
	if err != nil {
		t.Fatal(err)
	}
	final := res.Path[len(res.Path)-1][0]
	if !almostEqual(final, 3, 1e-3) {
		t.Errorf("gradient descent converged to x=%v, want ~3", final)
	}
}

func TestNewtonRootFindsSquareRoot(t *testing.T) {
	// Root of x^2 - 2 = 0 is sqrt(2).
	e := mustExpr(t, "x^2 - 2")
	res, err := NewtonRoot(e, "x", 1, 0, 0, nil)
	if err != nil {
		t.Fatal(err)
	}
	root := res.Path[len(res.Path)-1]
	if !almostEqual(root, math.Sqrt2, 1e-6) {
		t.Errorf("newton root = %v, want %v", root, math.Sqrt2)
	}
}

func TestStatistics(t *testing.T) {
	data := []float64{2, 4, 4, 4, 5, 5, 7, 9}
	mean, _ := Mean(data)
	if !almostEqual(mean, 5, 1e-9) {
		t.Errorf("mean = %v, want 5", mean)
	}
	popVar, _ := Variance(data, true)
	if !almostEqual(popVar, 4, 1e-9) {
		t.Errorf("population variance = %v, want 4", popVar)
	}
	std, _ := StdDev(data, true)
	if !almostEqual(std, 2, 1e-9) {
		t.Errorf("population std_dev = %v, want 2", std)
	}
}

func TestCorrelationPerfect(t *testing.T) {
	x := []float64{1, 2, 3, 4, 5}
	y := []float64{2, 4, 6, 8, 10}
	c, err := Correlation(x, y)
	if err != nil {
		t.Fatal(err)
	}
	if !almostEqual(c, 1, 1e-9) {
		t.Errorf("correlation = %v, want 1", c)
	}
}

func TestNormalDistribution(t *testing.T) {
	// Standard normal at x=0: pdf = 1/sqrt(2*pi), cdf = 0.5.
	pdf, _ := NormalPDF(0, 0, 1)
	if !almostEqual(pdf, 1/math.Sqrt(2*math.Pi), 1e-9) {
		t.Errorf("pdf(0) = %v", pdf)
	}
	cdf, _ := NormalCDF(0, 0, 1)
	if !almostEqual(cdf, 0.5, 1e-9) {
		t.Errorf("cdf(0) = %v, want 0.5", cdf)
	}
}

func TestSigmoidNumericallyStable(t *testing.T) {
	if v := Sigmoid(-1000); v != 0 {
		t.Errorf("sigmoid(-1000) = %v, want 0", v)
	}
	if v := Sigmoid(1000); !almostEqual(v, 1, 1e-12) {
		t.Errorf("sigmoid(1000) = %v, want ~1", v)
	}
	if v := Sigmoid(0); v != 0.5 {
		t.Errorf("sigmoid(0) = %v, want 0.5", v)
	}
}

func TestSoftmaxSumsToOne(t *testing.T) {
	out := Softmax([]float64{1, 2, 3})
	var sum float64
	for _, v := range out {
		sum += v
	}
	if !almostEqual(sum, 1, 1e-9) {
		t.Errorf("softmax sums to %v, want 1", sum)
	}
	// Largest logit should get the largest probability.
	if out[2] <= out[1] || out[1] <= out[0] {
		t.Errorf("softmax ordering wrong: %v", out)
	}
}

func TestMSELoss(t *testing.T) {
	v, err := MSELoss([]float64{1, 2, 3}, []float64{1, 2, 4})
	if err != nil {
		t.Fatal(err)
	}
	// errors: 0, 0, -1 -> squared: 0,0,1 -> mean = 1/3
	if !almostEqual(v, 1.0/3.0, 1e-9) {
		t.Errorf("mse = %v, want 1/3", v)
	}
}

func TestLinearRegressionPerfectFit(t *testing.T) {
	// y = 2 + 3x, exactly.
	X := Matrix{{1}, {2}, {3}, {4}, {5}}
	y := []float64{5, 8, 11, 14, 17}
	res, err := LinearRegression(X, y)
	if err != nil {
		t.Fatal(err)
	}
	if !almostEqual(res.Coefficients[0], 2, 1e-6) {
		t.Errorf("intercept = %v, want 2", res.Coefficients[0])
	}
	if !almostEqual(res.Coefficients[1], 3, 1e-6) {
		t.Errorf("slope = %v, want 3", res.Coefficients[1])
	}
	if !almostEqual(res.RSquared, 1, 1e-6) {
		t.Errorf("R^2 = %v, want 1", res.RSquared)
	}
}

// --- Solve dispatcher: the JSON-boundary layer the agent tool calls ---

func TestSolveEvaluate(t *testing.T) {
	out, err := Solve("evaluate", map[string]any{
		"expression": "x^2 + y^2",
		"variables":  map[string]any{"x": 3.0, "y": 4.0},
	})
	if err != nil {
		t.Fatal(err)
	}
	if !strings.Contains(out, "25") {
		t.Errorf("Solve(evaluate) = %q, want it to contain 25", out)
	}
}

func TestSolveMatrixMul(t *testing.T) {
	out, err := Solve("matrix_mul", map[string]any{
		"a": []any{[]any{1.0, 2.0}, []any{3.0, 4.0}},
		"b": []any{[]any{5.0, 6.0}, []any{7.0, 8.0}},
	})
	if err != nil {
		t.Fatal(err)
	}
	if !strings.Contains(out, "19") || !strings.Contains(out, "50") {
		t.Errorf("Solve(matrix_mul) = %q, missing expected entries", out)
	}
}

func TestSolveUnknownOperation(t *testing.T) {
	if _, err := Solve("not_a_real_operation", nil); err == nil {
		t.Fatal("expected an error for an unknown operation")
	}
}

func TestSolveGradientDescent(t *testing.T) {
	out, err := Solve("gradient_descent", map[string]any{
		"expression":    "(x - 3)^2",
		"variables":     []any{"x"},
		"start":         []any{0.0},
		"learning_rate": 0.1,
		"iterations":    200.0,
	})
	if err != nil {
		t.Fatal(err)
	}
	if !strings.Contains(out, "converged=true") {
		t.Errorf("Solve(gradient_descent) = %q, expected convergence", out)
	}
}

func TestSolveMissingField(t *testing.T) {
	_, err := Solve("mean", map[string]any{})
	if err == nil {
		t.Fatal("expected a missing-field error")
	}
	if !strings.Contains(err.Error(), "data") {
		t.Errorf("error %q should mention the missing field", err.Error())
	}
}

func TestOperationsListIsSorted(t *testing.T) {
	ops := Operations()
	if len(ops) == 0 {
		t.Fatal("Operations() returned nothing")
	}
	for i := 1; i < len(ops); i++ {
		if ops[i-1] > ops[i] {
			t.Fatalf("Operations() not sorted: %q before %q", ops[i-1], ops[i])
		}
	}
}