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

package grad

import (
	"math"
	"testing"

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

// gaussianLogDensity builds the log density of independent standard
// normals: log π(q) = −‖q‖²/2.
func gaussianLogDensity(q *Tensor) (*Tensor, error) {
	sq, err := q.Mul(q)
	if err != nil {
		return nil, err
	}
	total, err := sq.Sum()
	if err != nil {
		return nil, err
	}
	return total.Scale(-0.5)
}

// gammaLogDensity builds log π(x) = log x − x, the unnormalised log
// density of a Gamma(2, 1) distribution.
func gammaLogDensity(q *Tensor) (*Tensor, error) {
	logq, err := q.Log()
	if err != nil {
		return nil, err
	}
	shifted, err := logq.Sub(q)
	if err != nil {
		return nil, err
	}
	return shifted.Sum()
}

// TestSampleHMCNormal runs the chain on a two-dimensional standard
// normal and pins the sample moments: the target's mean is zero, its
// variance one and its components independent. A fixed seed makes the
// draw deterministic, so the bounds are checked facts about this run,
// not hopes about a random one.
func TestSampleHMCNormal(t *testing.T) {
	q0, err := core.FromFloats([]float64{2, -2}, 2)
	if err != nil {
		t.Fatalf("FromFloats: %v", err)
	}
	samples, err := SampleHMC(gaussianLogDensity, q0,
		HMCOptions{Step: 0.3, Steps: 20, BurnIn: 500, Samples: 6000, Thin: 1, Seed: 42})
	if err != nil {
		t.Fatalf("SampleHMC: %v", err)
	}
	if got := samples.Shape(); got[0] != 6000 || got[1] != 2 {
		t.Fatalf("samples shape %v, want [6000 2]", got)
	}
	rows := samples.Shape()[0]
	mean := []float64{0, 0}
	variance := []float64{0, 0}
	for r := range rows {
		for c := range 2 {
			v := samples.FloatAt(r*2 + c)
			mean[c] += v / float64(rows)
		}
	}
	for r := range rows {
		for c := range 2 {
			d := samples.FloatAt(r*2+c) - mean[c]
			variance[c] += d * d / float64(rows)
		}
	}
	for c := range 2 {
		if math.Abs(mean[c]) > 0.15 {
			t.Fatalf("mean[%d] = %.4g, want |mean| ≤ 0.15", c, mean[c])
		}
		if math.Abs(variance[c]-1) > 0.2 {
			t.Fatalf("variance[%d] = %.4g, want 1 ± 0.2", c, variance[c])
		}
	}
	// Cross moment of the independent components.
	cov := 0.0
	for r := range rows {
		cov += (samples.FloatAt(r*2) - mean[0]) * (samples.FloatAt(r*2+1) - mean[1]) / float64(rows)
	}
	if math.Abs(cov) > 0.15 {
		t.Fatalf("cross moment = %.4g, want |cov| ≤ 0.15", cov)
	}
}

// TestSampleHMCDeterministic checks the seed contract: the same seed
// replays bit-identically, a different seed does not.
func TestSampleHMCDeterministic(t *testing.T) {
	q0, _ := core.FromFloats([]float64{1}, 1)
	run := func(seed int64) *core.Array {
		samples, err := SampleHMC(gaussianLogDensity, q0,
			HMCOptions{Step: 0.4, Steps: 16, BurnIn: 100, Samples: 200, Seed: seed})
		if err != nil {
			t.Fatalf("SampleHMC: %v", err)
		}
		return samples
	}
	a, b := run(7), run(7)
	c := run(8)
	for i := range a.Len() {
		if a.FloatAt(i) != b.FloatAt(i) {
			t.Fatalf("the same seed produced different samples at %d", i)
		}
	}
	same := true
	for i := range a.Len() {
		if a.FloatAt(i) != c.FloatAt(i) {
			same = false
			break
		}
	}
	if same {
		t.Fatal("different seeds produced identical samples")
	}
}

// TestSampleHMCSupportedDensity samples a Gamma(2, 1) target,
// log π(x) = log x − x on x > 0. Proposals that overshoot into the
// forbidden half-line yield a NaN density and are rejected, so every
// kept sample stays positive and the mean approaches 2.
func TestSampleHMCSupportedDensity(t *testing.T) {
	q0, _ := core.FromFloats([]float64{1}, 1)
	samples, err := SampleHMC(gammaLogDensity, q0,
		HMCOptions{Step: 0.3, Steps: 10, BurnIn: 500, Samples: 6000, Seed: 3})
	if err != nil {
		t.Fatalf("SampleHMC: %v", err)
	}
	sum := 0.0
	for i := range samples.Len() {
		x := samples.FloatAt(i)
		if x <= 0 {
			t.Fatalf("sample %d = %g left the support", i, x)
		}
		sum += x
	}
	mean := sum / float64(samples.Len())
	if math.Abs(mean-2) > 0.15 {
		t.Fatalf("Gamma(2) mean = %.4g, want 2 ± 0.15", mean)
	}
}

// TestSampleHMCRejectedByGuard exercises the explicit rejection path:
// the density errors outside its support instead of returning NaN,
// and the chain still stays inside it.
func TestSampleHMCRejectedByGuard(t *testing.T) {
	q0, _ := core.FromFloats([]float64{0.5}, 1)
	samples, err := SampleHMC(func(q *Tensor) (*Tensor, error) {
		x := q.Data().FloatAt(0)
		if x <= 0 {
			return nil, errf("outside the support")
		}
		return gammaLogDensity(q)
	}, q0, HMCOptions{Step: 0.5, Steps: 20, BurnIn: 300, Samples: 2000, Seed: 5})
	if err != nil {
		t.Fatalf("SampleHMC: %v", err)
	}
	for i := range samples.Len() {
		if samples.FloatAt(i) <= 0 {
			t.Fatalf("sample %d = %g left the support", i, samples.FloatAt(i))
		}
	}
}

// TestSampleHMCErrors pins the validation contract, including a
// density that fails at the start, returns a non-scalar, or never
// touches the leaf core.
func TestSampleHMCErrors(t *testing.T) {
	q0, _ := core.FromFloats([]float64{1}, 1)
	if _, err := SampleHMC(nil, q0, HMCOptions{Step: 0.1, Steps: 5, Samples: 1}); err == nil {
		t.Fatal("expected an error for a nil density")
	}
	rank2, _ := core.FromFloats([]float64{1, 1}, 1, 2)
	if _, err := SampleHMC(gaussianLogDensity, rank2, HMCOptions{Step: 0.1, Steps: 5, Samples: 1}); err == nil {
		t.Fatal("expected an error for a rank-2 state")
	}
	if _, err := SampleHMC(gaussianLogDensity, nil, HMCOptions{Step: 0.1, Steps: 5, Samples: 1}); err == nil {
		t.Fatal("expected an error for a nil state")
	}
	if _, err := SampleHMC(gaussianLogDensity, q0, HMCOptions{Steps: 5, Samples: 1}); err == nil {
		t.Fatal("expected an error for a non-positive step")
	}
	if _, err := SampleHMC(gaussianLogDensity, q0, HMCOptions{Step: 0.1, Samples: 1}); err == nil {
		t.Fatal("expected an error for a non-positive step count")
	}
	if _, err := SampleHMC(gaussianLogDensity, q0, HMCOptions{Step: 0.1, Steps: 5}); err == nil {
		t.Fatal("expected an error for a non-positive sample count")
	}
	if _, err := SampleHMC(gaussianLogDensity, q0, HMCOptions{Step: 0.1, Steps: 5, Samples: 1, BurnIn: -3}); err == nil {
		t.Fatal("expected an error for a negative BurnIn")
	}
	nonScalar := func(q *Tensor) (*Tensor, error) {
		two, _ := core.FromFloats([]float64{1, 2}, 2)
		return FromArray(two, false), nil
	}
	if _, err := SampleHMC(nonScalar, q0, HMCOptions{Step: 0.1, Steps: 5, Samples: 1}); err == nil {
		t.Fatal("expected an error for a non-scalar density")
	}
	detached := func(q *Tensor) (*Tensor, error) {
		one, _ := core.FromFloats([]float64{1}, 1)
		return FromArray(one, false), nil
	}
	if _, err := SampleHMC(detached, q0, HMCOptions{Step: 0.1, Steps: 5, Samples: 1}); err == nil {
		t.Fatal("expected an error for a density disconnected from the leaf")
	}
	failsAtStart := func(q *Tensor) (*Tensor, error) {
		return nil, errf("no density at the start")
	}
	if _, err := SampleHMC(failsAtStart, q0, HMCOptions{Step: 0.1, Steps: 5, Samples: 1}); err == nil {
		t.Fatal("expected the start-time density error to be fatal")
	}
}