// Copyright (c) 2026 Petr Balvín <opensource@petrbalvin.org> (https://petrbalvin.org)
// SPDX-License-Identifier: MIT

package stats

import (
	"math"
	"testing"

	"sourcedock.dev/petrbalvin/tensor/internal/core"
)

// The mixed model against referents. The balanced one-way random
// effects model has a closed-form REML answer, the analysis of
// variance estimators, so the sweep's optimum is compared against
// figures computed from the raw data rather than quoted; the rest
// pins the recovery of known effects, the refusal surface and the
// determinism of the whole pipeline.

// mixedNoise is a deterministic stand-in for measurement noise: a
// bounded, aperiodic wiggle no sweep can mistake for structure.
func mixedNoise(i int) float64 {
	return 0.3*math.Sin(7.3*float64(i)+1.1)*math.Cos(2.1*float64(i)) +
		0.1*math.Sin(0.7*float64(i))
}

// mixedJitter is deterministic white jitter on [−1, 1): the xorshift
// finaliser of the house generator's mixing constants, run on the row
// index. Unlike the smooth wiggle above it cannot be absorbed by a
// within-group linear span, which is what the random-slope fit needs
// its residual scale to be.
func mixedJitter(i int) float64 {
	z := uint64(i)*2685821657736338717 + 1
	z ^= z >> 13
	z ^= z << 7
	z ^= z >> 17
	return float64(z>>11)/(1<<52)*2 - 1
}

func mixedVec(t *testing.T, vals []float64, shape ...int) *core.Array {
	t.Helper()
	a, err := core.FromFloats(vals, shape...)
	if err != nil {
		t.Fatal(err)
	}
	return a
}

func TestMixedModelBalancedOneWayANOVAReferent(t *testing.T) {
	// The balanced one-way random effects model: y_ig = μ + b_g + ε_ig
	// with k observations in each of m groups. REML's optimum is the
	// ANOVA answer: σ̂²_e = MSW and σ̂²_b = (MSB − MSW)/k, and the
	// intercept's GLS variance is (k·σ²_b + σ²_e)/(mk).
	const (
		mGroups = 8
		k       = 6
	)
	y := make([]float64, 0, mGroups*k)
	for g := range mGroups {
		effect := 2.0 * math.Sin(1.7*float64(g)+0.4) // the drawn b_g
		for i := range k {
			y = append(y, 5+effect+mixedNoise(g*k+i))
		}
	}
	groups := make([]int, mGroups*k)
	for g := range mGroups {
		for i := range k {
			groups[g*k+i] = g
		}
	}
	res, err := LinearMixedModel(
		mixedVec(t, y, len(y)),
		mixedVec(t, ones(len(y)), len(y), 1),
		mixedVec(t, ones(len(y)), len(y), 1),
		groups)
	if err != nil {
		t.Fatal(err)
	}
	if !res.Converged {
		t.Fatalf("the balanced fit did not converge (%d iterations)", res.Iterations)
	}
	// MSW and MSB from the raw data.
	grand := 0.0
	groupMeans := make([]float64, mGroups)
	for g := range mGroups {
		s := 0.0
		for i := range k {
			s += y[g*k+i]
		}
		groupMeans[g] = s / float64(k)
		grand += groupMeans[g]
	}
	grand /= float64(mGroups)
	msw := 0.0
	for g := range mGroups {
		for i := range k {
			d := y[g*k+i] - groupMeans[g]
			msw += d * d
		}
	}
	msw /= float64(mGroups * (k - 1))
	msb := 0.0
	for g := range mGroups {
		d := groupMeans[g] - grand
		msb += d * d
	}
	msb *= float64(k) / float64(mGroups-1)
	wantWithin := math.Max((msb-msw)/float64(k), 0)
	if math.Abs(res.ResidualVariance-msw) > 0.02*msw {
		t.Fatalf("σ̂²_e = %.6f, want the MSW %.6f", res.ResidualVariance, msw)
	}
	if math.Abs(res.RandomCovariance[0]-wantWithin) > 0.05*wantWithin {
		t.Fatalf("σ̂²_b = %.6f, want the ANOVA answer %.6f", res.RandomCovariance[0], wantWithin)
	}
	if math.Abs(res.Coefficients[0]-grand) > 1e-6 {
		t.Fatalf("μ̂ = %.8f, want the grand mean %.8f", res.Coefficients[0], grand)
	}
	wantVar := (float64(k)*res.RandomCovariance[0] + res.ResidualVariance) / float64(mGroups*k)
	if se := res.StandardErrors[0]; math.Abs(se*se-wantVar) > 1e-9*math.Max(1, wantVar) {
		t.Fatalf("SE² = %.10f, want the GLS variance %.10f", se*se, wantVar)
	}
	// The conditional fitted values reproduce the group means plus the
	// shrinkage the model applies; the residuals must complement them
	// to the response.
	for i := range len(y) {
		if math.Abs(res.Fitted[i]+res.Residuals[i]-y[i]) > 1e-9 {
			t.Fatalf("row %d: fitted + residuals = %g, want %g", i, res.Fitted[i]+res.Residuals[i], y[i])
		}
	}
}

// ones returns n constant 1 values, the intercept column.
func ones(n int) []float64 {
	out := make([]float64, n)
	for i := range out {
		out[i] = 1
	}
	return out
}

func TestMixedModelRandomSlopeRecovery(t *testing.T) {
	// Fixed effects of 1 and 2 with a per-group random slope, built by
	// hand so the truth is known exactly. The random design carries
	// the covariate alone: with the intercept column beside it the
	// random span covers the fixed design and the REML surface loses
	// its interior optimum to a ridge of singular covariance (pinned
	// by the divergence test below). The slope values are centred, so
	// the fixed part of the truth is exactly (1, 2).
	slopeValues := []float64{0.9, -1.1, 1.9, -0.3, -1.8, 0.4}
	groups := make([]int, 0, 36)
	xs := make([]float64, 0, 36)
	y := make([]float64, 0, 36)
	jitter := make([]float64, 0, 36)
	for g := range slopeValues {
		for _, x := range []float64{-1, -0.6, -0.2, 0.2, 0.6, 1} {
			groups = append(groups, g)
			xs = append(xs, x)
			jitter = append(jitter, mixedJitter(len(xs)-1))
		}
	}
	// The jitter is centred on its own sample: the fixed part of the
	// truth must stay exactly (1, 2), and a noise vector with a mean
	// would tilt the intercept instead of testing the recovery.
	mean := 0.0
	for _, j := range jitter {
		mean += j
	}
	mean /= float64(len(jitter))
	for g, s := range slopeValues {
		for k := range 6 {
			i := g*6 + k
			x := xs[i]
			y = append(y, 1+2*x+0.5*s*x+0.1*(jitter[i]-mean))
		}
	}
	design := make([]float64, 0, 2*len(xs))
	for _, x := range xs {
		design = append(design, 1, x)
	}
	res, err := LinearMixedModel(
		mixedVec(t, y, len(y)),
		mixedVec(t, design, len(xs), 2),
		mixedVec(t, xs, len(xs), 1),
		groups)
	if err != nil {
		t.Fatal(err)
	}
	if !res.Converged {
		t.Fatalf("the slope fit did not converge (%d iterations)", res.Iterations)
	}
	if math.Abs(res.Coefficients[0]-1) > 0.05 || math.Abs(res.Coefficients[1]-2) > 0.05 {
		t.Fatalf("β̂ = (%.4f, %.4f), want (1, 2)", res.Coefficients[0], res.Coefficients[1])
	}
	// The random slopes must rank with the true ones, and the group
	// labels must come back in first-appearance order.
	if len(res.GroupLabels) != 6 {
		t.Fatalf("group labels %v, want six groups", res.GroupLabels)
	}
	strongest := 0
	weakest := 0
	for g, s := range slopeValues {
		if s > slopeValues[strongest] {
			strongest = g
		}
		if s < slopeValues[weakest] {
			weakest = g
		}
	}
	if !(res.RandomEffects[weakest][0] < res.RandomEffects[strongest][0]) {
		t.Fatalf("the random slopes do not rank with the truth (%v against %v)",
			res.RandomEffects[weakest][0], slopeValues[strongest])
	}
	for _, se := range res.StandardErrors {
		if !(se > 0) || math.IsInf(se, 0) {
			t.Fatalf("the standard error %g is not finite and positive", se)
		}
	}
}

func TestMixedModelDivergenceRefusal(t *testing.T) {
	// A random design that spans the fixed one under an unstructured
	// covariance: the intercept and slope of every group absorb what
	// the fixed effects name, and the REML surface climbs a ridge of
	// singular Σ without a summit. The fit refuses with the condition
	// named instead of publishing the climb.
	groups := []int{0, 0, 0, 0, 1, 1, 1, 1, 2, 2, 2, 2}
	xs := []float64{-1, -0.5, 0.5, 1, -1, -0.5, 0.5, 1, -1, -0.5, 0.5, 1}
	y := make([]float64, len(groups))
	for i, g := range groups {
		y[i] = 1 + 2*xs[i] + 0.4*float64(g)*xs[i] + float64(g) + 0.1*mixedNoise(i)
	}
	design := make([]float64, 0, 2*len(xs))
	for _, x := range xs {
		design = append(design, 1, x)
	}
	_, err := LinearMixedModel(
		mixedVec(t, y, len(y)),
		mixedVec(t, design, len(xs), 2),
		mixedVec(t, design, len(xs), 2),
		groups)
	if err == nil {
		t.Fatal("a saturated random design was accepted")
	}
}

func TestMixedModelMatchesOLSWithoutRandomEffects(t *testing.T) {
	// When the groups carry no shared signal the components collapse
	// towards zero and the fit must land on the plain least squares
	// answer.
	xs := make([]float64, 18)
	y := make([]float64, 18)
	for i := range 18 {
		x := -1 + 2*float64(i)/17
		xs[i] = x
		y[i] = 3 - x + 0.4*mixedNoise(i)
	}
	groups := make([]int, 18)
	for i := range 18 {
		groups[i] = i % 6
	}
	design := make([]float64, 0, 36)
	for _, x := range xs {
		design = append(design, 1, x)
	}
	res, err := LinearMixedModel(
		mixedVec(t, y, len(y)),
		mixedVec(t, design, len(xs), 2),
		mixedVec(t, design, len(xs), 2),
		groups)
	if err != nil {
		t.Fatal(err)
	}
	ref, rerr := LinearRegression(mixedVec(t, design, len(xs), 2), mixedVec(t, y, len(y)))
	if rerr != nil {
		t.Fatal(rerr)
	}
	for j := range 2 {
		if math.Abs(res.Coefficients[j]-ref.Coefficients[j]) > 0.01 {
			t.Fatalf("coefficient %d = %.6f, want the OLS %.6f", j, res.Coefficients[j], ref.Coefficients[j])
		}
	}
	if res.RandomCovariance[0] > 0.01 {
		t.Fatalf("Σ̂ = %g on a group-free sample, want a collapsed component", res.RandomCovariance[0])
	}
}

func TestMixedModelDeterministic(t *testing.T) {
	y := make([]float64, 16)
	groups := make([]int, 16)
	for i := range 16 {
		y[i] = 2 + 0.9*math.Sin(float64(i%4)) + mixedNoise(i)
		groups[i] = i / 4
	}
	run := func() *LinearMixedModelResult {
		res, err := LinearMixedModel(
			mixedVec(t, y, len(y)),
			mixedVec(t, ones(len(y)), len(y), 1),
			mixedVec(t, ones(len(y)), len(y), 1),
			groups)
		if err != nil {
			t.Fatal(err)
		}
		return res
	}
	a, b := run(), run()
	if a.LogLikelihood != b.LogLikelihood || a.ResidualVariance != b.ResidualVariance ||
		a.RandomCovariance[0] != b.RandomCovariance[0] || a.Iterations != b.Iterations {
		t.Fatal("two identical fits disagreed")
	}
	for i := range a.Fitted {
		if a.Fitted[i] != b.Fitted[i] {
			t.Fatalf("row %d: fitted %g against %g", i, a.Fitted[i], b.Fitted[i])
		}
	}
}

func TestMixedModelShuffledLabels(t *testing.T) {
	// Labels out of order and with gaps must canonicalise by first
	// appearance, and permuting the rows with their labels must leave
	// the fitted components where they started.
	y := []float64{1.0, 1.2, 3.0, 3.1, 5.2, 5.1, 1.1, 3.2, 5.0, 1.3}
	groups := []int{7, 7, 3, 3, 5, 5, 7, 3, 5, 7}
	res, err := LinearMixedModel(
		mixedVec(t, y, len(y)),
		mixedVec(t, ones(len(y)), len(y), 1),
		mixedVec(t, ones(len(y)), len(y), 1),
		groups)
	if err != nil {
		t.Fatal(err)
	}
	want := []int{7, 3, 5}
	for i, label := range res.GroupLabels {
		if label != want[i] {
			t.Fatalf("group labels %v, want %v", res.GroupLabels, want)
		}
	}
	if !(res.RandomEffects[0][0] < res.RandomEffects[1][0] && res.RandomEffects[1][0] < res.RandomEffects[2][0]) {
		t.Fatalf("the random effects %v do not rank with the group means", res.RandomEffects)
	}
}

func TestMixedModelRefusals(t *testing.T) {
	y := mixedVec(t, []float64{1, 2, 3, 4}, 4)
	x := mixedVec(t, []float64{1, 1, 1, 1}, 4, 1)
	z := mixedVec(t, []float64{1, 1, 1, 1}, 4, 1)
	groups := []int{0, 0, 1, 1}
	if _, err := LinearMixedModel(mixedVec(t, []float64{1, 2}, 2, 1), x, z, groups); err == nil {
		t.Fatal("a rank-2 response was accepted")
	}
	if _, err := LinearMixedModel(mixedVec(t, []float64{1, 2, 3, 4, 5}, 5), x, z, groups); err == nil {
		t.Fatal("a row count mismatch was accepted")
	}
	if _, err := LinearMixedModel(y, mixedVec(t, []float64{1, 1, 1, 1}, 4, 1), mixedVec(t, []float64{1, 1, 1, 1}, 4, 1), []int{0, 1, 2}); err == nil {
		t.Fatal("a short label vector was accepted")
	}
	if _, err := LinearMixedModel(y, x, z, []int{0, 0, 1, -2}); err == nil {
		t.Fatal("a negative label was accepted")
	}
	// More coefficients than observations.
	big := mixedVec(t, []float64{1, 0, 1, 0, 1, 0, 1, 0, 1, 1, 1, 1}, 4, 3)
	if _, err := LinearMixedModel(y, big, z, groups); err == nil {
		t.Fatal("a saturated design was accepted")
	}
	// A singular fixed design: two identical columns.
	sing := mixedVec(t, []float64{1, 1, 1, 1, 2, 2, 2, 2}, 4, 2)
	if _, err := LinearMixedModel(y, sing, z, groups); err == nil {
		t.Fatal("a collinear fixed design was accepted")
	}
	// A single observation carries no fit.
	if _, err := LinearMixedModel(mixedVec(t, []float64{1}, 1), mixedVec(t, []float64{1}, 1, 1), mixedVec(t, []float64{1}, 1, 1), []int{0}); err == nil {
		t.Fatal("a one-row fit was accepted")
	}
}

// cholSolveInTest factors a symmetric positive definite matrix by its
// own Cholesky and solves against one right-hand side, an independent
// route the REML referent below evaluates its pieces through.
func cholSolveInTest(v []float64, rhs []float64, m int) (logDet float64, solution []float64) {
	l := make([]float64, m*m)
	for i := range m {
		for j := range i + 1 {
			s := v[i*m+j]
			for k := range j {
				s -= l[i*m+k] * l[j*m+k]
			}
			if i == j {
				l[i*m+j] = math.Sqrt(s)
			} else {
				l[i*m+j] = s / l[j*m+j]
			}
		}
		logDet += 2 * math.Log(l[i*m+i])
	}
	x := append([]float64(nil), rhs...)
	for i := range m {
		s := x[i]
		for k := range i {
			s -= l[i*m+k] * x[k]
		}
		x[i] = s / l[i*m+i]
	}
	for i := m - 1; i >= 0; i-- {
		s := x[i]
		for k := i + 1; k < m; k++ {
			s -= l[k*m+i] * x[k]
		}
		x[i] = s / l[i*m+i]
	}
	return logDet, x
}

func TestMixedModelREMLLogLikelihoodReferent(t *testing.T) {
	// The reported LogLikelihood is checked against an independent
	// evaluation of the REML criterion at the fitted components,
	// assembled from first principles in this test:
	// -2·logL = Σ_g log|V_g| + rᵀV⁻¹r + log|XᵀV⁻¹X| + (n−p)·ln 2π,
	// with V_g = Z_gΣZ_gᵀ + σ²I and r the fixed-part residual. The
	// one-way balanced design keeps V_g compound symmetric, so the
	// pieces are small and the route shares no arithmetic with the fit.
	const (
		mGroups = 6
		k       = 4
	)
	y := make([]float64, 0, mGroups*k)
	groups := make([]int, mGroups*k)
	for g := range mGroups {
		effect := 1.5 * math.Sin(0.9*float64(g)+0.2)
		for i := range k {
			y = append(y, 3+effect+mixedNoise(g*k+i))
			groups[g*k+i] = g
		}
	}
	res, err := LinearMixedModel(
		mixedVec(t, y, len(y)),
		mixedVec(t, ones(len(y)), len(y), 1),
		mixedVec(t, ones(len(y)), len(y), 1),
		groups)
	if err != nil {
		t.Fatal(err)
	}
	if !res.Converged {
		t.Fatalf("the fit did not converge (%d iterations)", res.Iterations)
	}
	sigma2 := res.ResidualVariance
	tau2 := res.RandomCovariance[0]
	beta := res.Coefficients[0]
	n := len(y)
	p := 1
	logDetV := 0.0
	quad := 0.0
	xtvix := 0.0
	for g := range mGroups {
		m := k
		v := make([]float64, m*m)
		for i := range m {
			for j := range m {
				v[i*m+j] = tau2
			}
			v[i*m+i] += sigma2
		}
		r := make([]float64, m)
		for i := range m {
			r[i] = y[g*k+i] - beta
		}
		ld, u := cholSolveInTest(v, r, m)
		logDetV += ld
		for i := range m {
			quad += r[i] * u[i]
		}
		onesRHS := make([]float64, m)
		for i := range onesRHS {
			onesRHS[i] = 1
		}
		_, w := cholSolveInTest(v, onesRHS, m)
		for i := range m {
			xtvix += w[i]
		}
	}
	want := -0.5 * (logDetV + quad + math.Log(xtvix) + float64(n-p)*math.Log(2*math.Pi))
	if math.Abs(res.LogLikelihood-want) > 1e-8*(1+math.Abs(want)) {
		t.Fatalf("LogLikelihood = %.10f, want the independent REML %.10f (difference %.3e)",
			res.LogLikelihood, want, res.LogLikelihood-want)
	}
}