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

package stats

import (
	"math"
	"math/big"
	"testing"

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

// The hidden Markov machinery against exact referents. The model
// probabilities are rational, so a short sequence's likelihood and
// its most likely path are computed exactly by enumerating every
// path, and the recursions are compared against that enumeration, not
// against quoted figures.

// exactModel holds one small model's parameters as rationals.
type exactModel struct {
	initial    []float64
	transition []float64 // 2×2
	emission   []float64 // 2×2
}

// ratRow converts a float row to rationals.
func ratRow(vals []float64) []*big.Rat {
	out := make([]*big.Rat, len(vals))
	for i, v := range vals {
		out[i] = big.NewRat(int64(v*1e15), 1e15)
	}
	return out
}

// exactPathProbability evaluates one path's joint probability with
// the observations it explains.
func exactPathProbability(m *exactModel, observations, path []int) *big.Rat {
	initial := ratRow(m.initial)
	trans := make([][]*big.Rat, 2)
	for k := range 2 {
		trans[k] = ratRow(m.transition[k*2 : (k+1)*2])
	}
	emis := make([][]*big.Rat, 2)
	for k := range 2 {
		emis[k] = ratRow(m.emission[k*2 : (k+1)*2])
	}
	total := new(big.Rat).Set(initial[path[0]])
	total.Mul(total, emis[path[0]][observations[0]])
	for t := 1; t < len(path); t++ {
		total.Mul(total, trans[path[t-1]][path[t]])
		total.Mul(total, emis[path[t]][observations[t]])
	}
	return total
}

func TestHiddenMarkovAgainstExactEnumeration(t *testing.T) {
	m := &exactModel{
		initial:    []float64{0.6, 0.4},
		transition: []float64{0.7, 0.3, 0.2, 0.8},
		emission:   []float64{0.9, 0.1, 0.25, 0.75},
	}
	model, err := NewHiddenMarkovModel(m.initial, m.transition, m.emission)
	if err != nil {
		t.Fatal(err)
	}
	observations := []int{0, 0, 1, 0, 1, 1, 0, 1}
	// Every path over the sequence, exactly.
	bestProb := new(big.Rat)
	totalProb := new(big.Rat)
	var bestPath []int
	path := make([]int, len(observations))
	var walk func(t int)
	walk = func(t int) {
		if t == len(path) {
			p := exactPathProbability(m, observations, path)
			totalProb.Add(totalProb, p)
			if p.Cmp(bestProb) > 0 {
				bestProb.Set(p)
				bestPath = append([]int(nil), path...)
			}
			return
		}
		for _, s := range []int{0, 1} {
			path[t] = s
			walk(t + 1)
		}
	}
	walk(0)
	// The forward recursion's likelihood against the exact total.
	_, ll, err := model.Forward(observations)
	if err != nil {
		t.Fatal(err)
	}
	wantLL, _ := new(big.Float).SetRat(totalProb).Float64()
	if math.Abs(ll-math.Log(wantLL)) > 1e-12 {
		t.Fatalf("the forward likelihood = %.16f, want the exact %.16f", ll, math.Log(wantLL))
	}
	// Viterbi against the exact best path.
	states, pathProb, err := model.Viterbi(observations)
	if err != nil {
		t.Fatal(err)
	}
	wantBest, _ := new(big.Float).SetRat(bestProb).Float64()
	if math.Abs(pathProb-math.Log(wantBest)) > 1e-12 {
		t.Fatalf("Viterbi's log probability = %.16f, want the exact %.16f", pathProb, math.Log(wantBest))
	}
	for step := range states {
		if states[step] != bestPath[step] {
			t.Fatalf("Viterbi's path %v leaves the exact best path %v at step %d", states, bestPath, step)
		}
	}
	// The smoothing's last row must coincide with filtering's: both
	// see the whole sequence by then.
	filtered, _, err := model.Forward(observations)
	if err != nil {
		t.Fatal(err)
	}
	smoothed, _, err := model.Smooth(observations)
	if err != nil {
		t.Fatal(err)
	}
	for k := range filtered[len(observations)-1] {
		if filtered[len(observations)-1][k] != smoothed[len(observations)-1][k] {
			t.Fatalf("the last filtered posterior %v disagrees with the smoothed one %v",
				filtered[len(observations)-1], smoothed[len(observations)-1])
		}
		if smoothed[len(observations)-1][k] < 0 || smoothed[len(observations)-1][k] > 1 {
			t.Fatalf("the smoothed posterior %g is not a probability", smoothed[len(observations)-1][k])
		}
	}
}

func TestHiddenMarkovSingleStateLikelihood(t *testing.T) {
	// One state: the model is an independent-symbol law and the
	// likelihood is the product of the emissions, by hand.
	model, err := NewHiddenMarkovModel([]float64{1}, []float64{1}, []float64{0.25, 0.75})
	if err != nil {
		t.Fatal(err)
	}
	observations := []int{1, 1, 0, 1, 1, 1, 0, 1, 1}
	_, ll, err := model.Forward(observations)
	if err != nil {
		t.Fatal(err)
	}
	want := 0.0
	for _, o := range observations {
		if o == 0 {
			want += math.Log(0.25)
		} else {
			want += math.Log(0.75)
		}
	}
	if math.Abs(ll-want) > 1e-12 {
		t.Fatalf("the single-state likelihood = %.16f, want %.16f", ll, want)
	}
	path, _, err := model.Viterbi(observations)
	if err != nil {
		t.Fatal(err)
	}
	for _, s := range path {
		if s != 0 {
			t.Fatalf("the single-state decode left the only state: %v", path)
		}
	}
}

func TestHiddenMarkovScalingSurvivesLongSequences(t *testing.T) {
	// A thousand-step sequence under a sticky model: the raw forward
	// probabilities underflow long before the end, and the scaled
	// recursions must answer a finite likelihood and rows that stay
	// distributions.
	model, err := NewHiddenMarkovModel(
		[]float64{0.5, 0.5},
		[]float64{0.99, 0.01, 0.02, 0.98},
		[]float64{0.8, 0.2, 0.3, 0.7})
	if err != nil {
		t.Fatal(err)
	}
	observations := make([]int, 1000)
	for t := range observations {
		observations[t] = (t * 7 / 3) % 2
	}
	filtered, ll, err := model.Forward(observations)
	if err != nil {
		t.Fatal(err)
	}
	if math.IsInf(ll, 0) || math.IsNaN(ll) {
		t.Fatalf("the long sequence's likelihood = %g", ll)
	}
	for step, row := range filtered {
		total := 0.0
		for _, v := range row {
			total += v
		}
		if math.Abs(total-1) > 1e-9 {
			t.Fatalf("the filtered row at %d sums to %g", step, total)
		}
	}
}

func TestHiddenMarkovFitRecovery(t *testing.T) {
	// Baum-Welch on a sequence drawn from a known model: the fitted
	// likelihood must clear the known model's own, the fit must
	// converge, and a repeat run must be bit-identical.
	truth, err := NewHiddenMarkovModel(
		[]float64{0.65, 0.35},
		[]float64{0.8, 0.2, 0.15, 0.85},
		[]float64{0.85, 0.15, 0.3, 0.7})
	if err != nil {
		t.Fatal(err)
	}
	observations := hmmSimulate(t, truth, 400, 20260925)
	_, truthLL, err := truth.Forward(observations)
	if err != nil {
		t.Fatal(err)
	}
	fit, err := FitHiddenMarkovModel(core.NewGenerator(7), observations, 2, 2)
	if err != nil {
		t.Fatal(err)
	}
	if !fit.Converged {
		t.Fatalf("the fit did not converge (%d iterations)", fit.Iterations)
	}
	if fit.LogLikelihood < truthLL {
		t.Fatalf("the fitted likelihood %.6f sits below the truth's %.6f, expectation maximisation failed to climb",
			fit.LogLikelihood, truthLL)
	}
	again, err := FitHiddenMarkovModel(core.NewGenerator(7), observations, 2, 2)
	if err != nil {
		t.Fatal(err)
	}
	if again.LogLikelihood != fit.LogLikelihood || again.Iterations != fit.Iterations {
		t.Fatal("two identical fits disagreed")
	}
	for i := range fit.Model.Initial {
		if fit.Model.Initial[i] != again.Model.Initial[i] ||
			fit.Model.Transition[i] != again.Model.Transition[i] ||
			fit.Model.Emission[i] != again.Model.Emission[i] {
			t.Fatal("two identical fits returned different parameters")
		}
	}
	// The fitted model answers the same likelihood through Forward as
	// the fit reported.
	_, ll, err := fit.Model.Forward(observations)
	if err != nil {
		t.Fatal(err)
	}
	if ll != fit.LogLikelihood {
		t.Fatalf("the fitted model's likelihood %.16f differs from the fit's report %.16f", ll, fit.LogLikelihood)
	}
}

// hmmSimulate draws an observation sequence from a model through the
// house generator, by inverse transform on the cumulative rows.
func hmmSimulate(t *testing.T, model *HiddenMarkovModel, length int, seed int64) []int {
	t.Helper()
	g := core.NewGenerator(seed)
	states := len(model.Initial)
	symbols := len(model.Emission) / states
	draw := func(row []float64) int {
		u := g.Unit()
		total := 0.0
		for i, v := range row {
			total += v
			if u < total {
				return i
			}
		}
		return len(row) - 1
	}
	state := draw(model.Initial)
	out := make([]int, length)
	for t := range length {
		out[t] = draw(model.Emission[state*symbols : (state+1)*symbols])
		state = draw(model.Transition[state*states : (state+1)*states])
	}
	return out
}

func TestHiddenMarkovRefusals(t *testing.T) {
	model, err := NewHiddenMarkovModel([]float64{0.6, 0.4}, []float64{0.7, 0.3, 0.2, 0.8}, []float64{0.9, 0.1, 0.25, 0.75})
	if err != nil {
		t.Fatal(err)
	}
	if _, _, err := model.Forward(nil); err == nil {
		t.Fatal("an empty sequence was accepted")
	}
	if _, _, err := model.Forward([]int{0, 2, 1}); err == nil {
		t.Fatal("an out-of-range symbol was accepted")
	}
	if _, err := NewHiddenMarkovModel([]float64{0.6, 0.5}, []float64{0.7, 0.3, 0.2, 0.8}, []float64{0.9, 0.1, 0.25, 0.75}); err == nil {
		t.Fatal("an initial distribution summing over 1 was accepted")
	}
	if _, err := NewHiddenMarkovModel([]float64{0.6, 0.4}, []float64{0.7, 0.3, 0.2}, []float64{0.9, 0.1, 0.25, 0.75}); err == nil {
		t.Fatal("a short transition matrix was accepted")
	}
	if _, err := NewHiddenMarkovModel([]float64{0.6, 0.4}, []float64{0.7, 0.3, 0.2, 0.8}, []float64{0.9, 0.1, 0.25}); err == nil {
		t.Fatal("a ragged emission matrix was accepted")
	}
	if _, err := NewHiddenMarkovModel([]float64{1.2, -0.2}, []float64{0.7, 0.3, 0.2, 0.8}, []float64{0.9, 0.1, 0.25, 0.75}); err == nil {
		t.Fatal("a probability outside [0, 1] was accepted")
	}
	if _, err := FitHiddenMarkovModel(nil, []int{0, 1}, 2, 2); err == nil {
		t.Fatal("a nil generator was accepted")
	}
	if _, err := FitHiddenMarkovModel(core.NewGenerator(1), []int{0, 1, 5}, 2, 2); err == nil {
		t.Fatal("an out-of-range training symbol was accepted")
	}
	var nilModel *HiddenMarkovModel
	if _, _, err := nilModel.Forward([]int{0}); err == nil {
		t.Fatal("a nil model was accepted")
	}
}

func TestHiddenMarkovZeroProbabilitySequence(t *testing.T) {
	// A structural zero in the emissions: the only state never emits
	// symbol 1, so a sequence holding it has probability zero and no
	// posterior exists. The recursions must refuse rather than divide
	// the zero normaliser into NaN posteriors.
	model, err := NewHiddenMarkovModel([]float64{1}, []float64{1}, []float64{1, 0})
	if err != nil {
		t.Fatal(err)
	}
	if _, _, err := model.Forward([]int{1}); err == nil {
		t.Fatal("a zero-probability sequence was accepted by Forward")
	}
	if _, _, err := model.Smooth([]int{1}); err == nil {
		t.Fatal("a zero-probability sequence was accepted by Smooth")
	}
	// A sequence that dies one step in: the first symbol is live, the
	// second is not, and the refusal is the same.
	if _, _, err := model.Forward([]int{0, 1}); err == nil {
		t.Fatal("a sequence that dies at step 1 was accepted by Forward")
	}
	// The live half of the same model still answers: the zero emission
	// silences nothing a legal sequence needs.
	filtered, ll, err := model.Forward([]int{0, 0})
	if err != nil {
		t.Fatal(err)
	}
	if ll != 0 || filtered[0][0] != 1 {
		t.Fatalf("the certain sequence answered (%g, %v), want (0, [1])", ll, filtered[0])
	}
}

func TestHiddenMarkovFitSingleObservation(t *testing.T) {
	// One observation carries emission and initial evidence but no
	// transition evidence: the re-estimation divided the zero count into
	// NaN rows, the constructor refused them, and the sweep then called
	// a method on the nil model and crashed the process. The fit must
	// answer with a valid model whose transitions keep their starting
	// estimate.
	g := core.NewGenerator(11)
	res, err := FitHiddenMarkovModel(g, []int{0}, 2, 2)
	if err != nil {
		t.Fatalf("FitHiddenMarkovModel on one observation: %v", err)
	}
	if res.Model == nil {
		t.Fatal("FitHiddenMarkovModel on one observation returned no model")
	}
	if math.IsNaN(res.LogLikelihood) || math.IsInf(res.LogLikelihood, 0) {
		t.Fatalf("log likelihood = %g, want a finite value", res.LogLikelihood)
	}
	for k := range 2 {
		sum := 0.0
		for _, v := range res.Model.Transition[k*2 : k*2+2] {
			if math.IsNaN(v) || v < 0 {
				t.Fatalf("transition row %d holds %g, want probabilities", k, v)
			}
			sum += v
		}
		if math.Abs(sum-1) > 1e-9 {
			t.Fatalf("transition row %d sums to %g, want 1", k, sum)
		}
	}
}

func TestViterbiZeroProbabilitySequence(t *testing.T) {
	// Symbol 2 is emitted by no state, so every path through the sequence
	// has probability zero. Forward and Smooth refuse such a sequence, and
	// Viterbi documents the same refusals: a meaningless path beside a
	// −Inf log probability is not an answer.
	model, err := NewHiddenMarkovModel(
		[]float64{0.5, 0.5},
		[]float64{0.5, 0.5, 0.5, 0.5},
		[]float64{0.5, 0.5, 0, 0.5, 0.5, 0},
	)
	if err != nil {
		t.Fatal(err)
	}
	if _, _, err := model.Viterbi([]int{0, 2, 1}); err == nil {
		t.Fatal("a zero-probability sequence was accepted by Viterbi")
	}
	// The live symbols of the same model still decode.
	path, lp, err := model.Viterbi([]int{0, 1})
	if err != nil {
		t.Fatal(err)
	}
	if path[0] != path[1] || lp >= 0 {
		t.Fatalf("Viterbi([0, 1]) = (%v, %g), want a coherent path with a negative log probability", path, lp)
	}
}