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

package stats

import (
	"math"

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

// Hidden Markov models over discrete observations: the filtered and
// smoothed state posteriors, the most likely state path and the
// Baum-Welch fit of the parameters. Every computation runs on the
// scaled recursions, where each sweep's normaliser absorbs the
// likelihood of what has been seen, so a long sequence cannot
// underflow the way the raw forward probabilities would; the log
// likelihood itself is the sum of the normalisers' logarithms. The
// decoding walks in log space with the lowest index winning ties, and
// the fit draws its starting point from the house generator, so the
// whole surface is deterministic for a given generator state and
// sequence.

// HiddenMarkovModel is a hidden Markov model over a set of hidden
// states and a set of observation symbols. The fields hold the
// distribution parameters in row-major order: the rows of Transition
// are the conditionals P(next = j | current = k) and the rows of
// Emission the conditionals P(symbol = m | state = k). The
// constructors validate every row as a distribution, so a model value
// is always safe to evaluate.
type HiddenMarkovModel struct {
	// Initial is the distribution over states at the first step.
	Initial []float64
	// Transition is the states×states matrix of one-step
	// probabilities, row-major.
	Transition []float64
	// Emission is the states×symbols matrix of observation
	// probabilities, row-major.
	Emission []float64
}

// HiddenMarkovFitResult carries the Baum-Welch fit of a hidden Markov
// model.
type HiddenMarkovFitResult struct {
	// Model is the fitted model.
	Model *HiddenMarkovModel
	// LogLikelihood is the log likelihood of the training sequence
	// under the fitted model.
	LogLikelihood float64
	// Iterations counts the Baum-Welch sweeps taken; Converged reports
	// whether the log likelihood settled under the tolerance before
	// the budget ran out.
	Iterations int
	Converged  bool
}

// hmmMaxIterations caps the Baum-Welch sweeps, with the same
// reporting convention the other fits carry.
const hmmMaxIterations = 500

// hmmTolerance is the convergence tolerance on the log likelihood:
// the sweeps stop once it moves by less than hmmTolerance scaled by
// 1 + |log likelihood|. Baum-Welch's ascent crawls along a flat ridge
// long before its parameters settle to the last digit, and the
// tolerance stops the sweep where the objective's remaining climb is
// immaterial rather than where the parameters stop moving.
const hmmTolerance = 1e-6

// hmmFloor is the probability floor the re-estimation applies before
// renormalising: a parameter driven to exactly zero would silence a
// state or a symbol forever, and the floor keeps every path alive at
// a width no plausible parameter sits under.
const hmmFloor = 1e-12

// NewHiddenMarkovModel validates and copies the parameters of a
// hidden Markov model: Initial of the length of the state set,
// Transition with one row per state over the states, Emission with
// one row per state over the symbols. Every entry must be finite and
// in [0, 1] and every row must sum to 1 within 1e-9; a model that
// failed any of these would silently weight impossible paths later.
func NewHiddenMarkovModel(initial, transition, emission []float64) (*HiddenMarkovModel, error) {
	const name = "NewHiddenMarkovModel"
	states := len(initial)
	if states < 1 {
		return nil, base.Errf("%s: the model needs at least one state", name)
	}
	if len(transition) != states*states {
		return nil, base.Errf("%s: the transition matrix holds %d entries, want %d for %d states",
			name, len(transition), states*states, states)
	}
	if len(emission) == 0 || len(emission)%states != 0 {
		return nil, base.Errf("%s: the emission matrix holds %d entries, not a whole number of rows for %d states",
			name, len(emission), states)
	}
	if err := hmmCheckRow(name, "the initial distribution", initial); err != nil {
		return nil, err
	}
	for k := range states {
		if err := hmmCheckRow(name, "the transition rows", transition[k*states:(k+1)*states]); err != nil {
			return nil, err
		}
	}
	symbols := len(emission) / states
	for k := range states {
		if err := hmmCheckRow(name, "the emission rows", emission[k*symbols:(k+1)*symbols]); err != nil {
			return nil, err
		}
	}
	return &HiddenMarkovModel{
		Initial:    append([]float64(nil), initial...),
		Transition: append([]float64(nil), transition...),
		Emission:   append([]float64(nil), emission...),
	}, nil
}

// hmmCheckRow validates one distribution row: finite, inside [0, 1],
// summing to 1 within 1e-9.
func hmmCheckRow(name, what string, row []float64) error {
	total := 0.0
	for _, v := range row {
		if math.IsNaN(v) || math.IsInf(v, 0) {
			return base.Errf("%s: %s hold the non-finite value %g", name, what, v)
		}
		if v < 0 || v > 1 {
			return base.Errf("%s: %s hold the probability %g outside [0, 1]", name, what, v)
		}
		total += v
	}
	if math.Abs(total-1) > 1e-9 {
		return base.Errf("%s: %s sum to %g, not 1", name, what, total)
	}
	return nil
}

// hmmStates returns the model's state count.
func (m *HiddenMarkovModel) hmmStates() int { return len(m.Initial) }

// hmmSymbols returns the model's symbol count.
func (m *HiddenMarkovModel) hmmSymbols() int { return len(m.Emission) / max(1, len(m.Initial)) }

// hmmWidth returns the model's symbol count from its own fields, the
// shape the recursions read the emission rows with.
func (m *HiddenMarkovModel) hmmWidth() int { return len(m.Emission) / len(m.Initial) }

// hmmCheckSequence validates an observation sequence against the
// model's symbol set.
func hmmCheckSequence(name string, observations []int, symbols int) error {
	if len(observations) == 0 {
		return base.Errf("%s: the observation sequence is empty", name)
	}
	for t, o := range observations {
		if o < 0 || o >= symbols {
			return base.Errf("%s: observation %d is symbol %d, outside the model's %d symbols", name, t, o, symbols)
		}
	}
	return nil
}

// hmmForward runs the scaled forward recursion: alpha[t] is the
// P(state, observations 0..t) vector normalised by its own total, and
// the totals accumulate into the log likelihood. The returned rows
// are fresh. A normaliser of exactly zero means the sequence has
// probability zero under the model, the structural zeros the
// parameters carry having silenced every state at some step: the
// posteriors are undefined there and the recursion refuses rather
// than divide the zero into NaNs.
func (m *HiddenMarkovModel) hmmForward(observations []int) ([][]float64, []float64, error) {
	const name = "HiddenMarkovModel.Forward"
	states := m.hmmStates()
	width := m.hmmWidth()
	alpha := make([][]float64, len(observations))
	scales := make([]float64, len(observations))
	row := make([]float64, states)
	for k := range states {
		row[k] = m.Initial[k] * m.Emission[k*width+observations[0]]
	}
	total := 0.0
	for _, v := range row {
		total += v
	}
	if total == 0 {
		return nil, nil, base.Errf("%s: the observation sequence has probability zero under the model", name)
	}
	scales[0] = total
	for k := range states {
		row[k] /= total
	}
	alpha[0] = row
	for t := 1; t < len(observations); t++ {
		next := make([]float64, states)
		sym := observations[t]
		for j := range states {
			total := 0.0
			for k := range states {
				total += row[k] * m.Transition[k*states+j]
			}
			next[j] = total * m.Emission[j*width+sym]
		}
		total := 0.0
		for _, v := range next {
			total += v
		}
		if total == 0 {
			return nil, nil, base.Errf("%s: the observation sequence has probability zero under the model at step %d", name, t)
		}
		scales[t] = total
		for j := range states {
			next[j] /= total
		}
		alpha[t] = next
		row = next
	}
	return alpha, scales, nil
}

// hmmBackward runs the scaled backward recursion against the forward
// scales, so the product alpha·beta normalises directly into the
// smoothed posteriors.
func (m *HiddenMarkovModel) hmmBackward(observations []int, scales []float64) [][]float64 {
	states := m.hmmStates()
	width := m.hmmWidth()
	beta := make([][]float64, len(observations))
	last := make([]float64, states)
	for k := range states {
		last[k] = 1
	}
	beta[len(observations)-1] = last
	for t := len(observations) - 2; t >= 0; t-- {
		row := make([]float64, states)
		sym := observations[t+1]
		for k := range states {
			total := 0.0
			for j := range states {
				total += m.Transition[k*states+j] * m.Emission[j*width+sym] * beta[t+1][j]
			}
			row[k] = total / scales[t+1]
		}
		beta[t] = row
	}
	return beta
}

// Forward runs the scaled forward recursion over the observation
// sequence and returns the filtered state posteriors, one row of
// state probabilities per step, with the sequence's log likelihood
// P(observations | model). Refuses a nil model, an empty sequence,
// any observation outside the model's symbol set, and a sequence of
// probability zero under the model, whose structural zero an emission
// or transition carries leaves the filtered posteriors undefined.
func (m *HiddenMarkovModel) Forward(observations []int) (filtered [][]float64, logLikelihood float64, err error) {
	const name = "HiddenMarkovModel.Forward"
	if m == nil {
		return nil, 0, base.Errf("%s: the model is nil", name)
	}
	if err := hmmCheckSequence(name, observations, m.hmmSymbols()); err != nil {
		return nil, 0, err
	}
	alpha, scales, ferr := m.hmmForward(observations)
	if ferr != nil {
		return nil, 0, ferr
	}
	return alpha, hmmLogLikelihood(scales), nil
}

// Smooth runs the forward-backward recursions over the observation
// sequence and returns the smoothed state posteriors P(state at t |
// all observations), one row per step, with the sequence's log
// likelihood. The refusals are Forward's own.
func (m *HiddenMarkovModel) Smooth(observations []int) (smoothed [][]float64, logLikelihood float64, err error) {
	const name = "HiddenMarkovModel.Smooth"
	if m == nil {
		return nil, 0, base.Errf("%s: the model is nil", name)
	}
	if err := hmmCheckSequence(name, observations, m.hmmSymbols()); err != nil {
		return nil, 0, err
	}
	alpha, scales, ferr := m.hmmForward(observations)
	if ferr != nil {
		return nil, 0, ferr
	}
	beta := m.hmmBackward(observations, scales)
	gamma := make([][]float64, len(observations))
	for t := range observations {
		row := make([]float64, len(m.Initial))
		total := 0.0
		for k := range row {
			row[k] = alpha[t][k] * beta[t][k]
			total += row[k]
		}
		for k := range row {
			row[k] /= total
		}
		gamma[t] = row
	}
	return gamma, hmmLogLikelihood(scales), nil
}

// hmmLogLikelihood folds the forward scales into the sequence's log
// likelihood.
func hmmLogLikelihood(scales []float64) float64 {
	total := 0.0
	for _, s := range scales {
		total += math.Log(s)
	}
	return total
}

// Viterbi decodes the most likely state path through the observation
// sequence in log space, breaking ties toward the lowest state index,
// and returns the path with its log probability
// log P(path, observations | model). The refusals are Forward's own.
func (m *HiddenMarkovModel) Viterbi(observations []int) (states []int, logProbability float64, err error) {
	const name = "HiddenMarkovModel.Viterbi"
	if m == nil {
		return nil, 0, base.Errf("%s: the model is nil", name)
	}
	if err := hmmCheckSequence(name, observations, m.hmmSymbols()); err != nil {
		return nil, 0, err
	}
	n := m.hmmStates()
	width := m.hmmWidth()
	scores := make([]float64, n)
	for k := range n {
		scores[k] = math.Log(m.Initial[k]) + math.Log(m.Emission[k*width+observations[0]])
	}
	back := make([][]int, len(observations))
	back[0] = make([]int, n)
	for t := 1; t < len(observations); t++ {
		next := make([]float64, n)
		prev := make([]int, n)
		sym := observations[t]
		for j := range n {
			best := math.Inf(-1)
			arg := 0
			for k := range n {
				candidate := scores[k] + math.Log(m.Transition[k*n+j])
				if candidate > best {
					best = candidate
					arg = k
				}
			}
			next[j] = best + math.Log(m.Emission[j*width+sym])
			prev[j] = arg
		}
		scores = next
		back[t] = prev
	}
	best := math.Inf(-1)
	arg := 0
	for k := range n {
		if scores[k] > best {
			best = scores[k]
			arg = k
		}
	}
	if math.IsInf(best, -1) {
		// Every path has probability zero, the structural zero an
		// emission or transition carries: the same refusal Forward makes,
		// not a meaningless path beside a −Inf log probability.
		return nil, 0, base.Errf("%s: the observation sequence has probability zero under the model", name)
	}
	path := make([]int, len(observations))
	path[len(observations)-1] = arg
	for t := len(observations) - 1; t > 0; t-- {
		path[t-1] = back[t][path[t]]
	}
	return path, best, nil
}

// FitHiddenMarkovModel fits a hidden Markov model of the given state
// and symbol counts to the observation sequence by Baum-Welch
// expectation maximisation. The starting point draws every
// distribution row from the flat Dirichlet through the house
// generator, so the fit is deterministic for the generator state;
// expectation maximisation climbs to a local maximum of the
// likelihood, and a different generator state may land on a different
// one.
//
// The sweeps stop once the log likelihood moves by less than
// hmmTolerance relative, within hmmMaxIterations sweeps; Converged
// names which happened. The re-estimation floors every parameter at
// hmmFloor and renormalises the rows, so no symbol or transition is
// silenced outright by one sweep. A nil generator, a state or symbol
// count below one, an empty sequence or an observation outside the
// symbol set is an error.
func FitHiddenMarkovModel(g *core.Generator, observations []int, states, symbols int) (*HiddenMarkovFitResult, error) {
	const name = "FitHiddenMarkovModel"
	if g == nil {
		return nil, base.Errf("%s: the generator is nil", name)
	}
	if states < 1 {
		return nil, base.Errf("%s: the model needs at least one state, got %d", name, states)
	}
	if symbols < 1 {
		return nil, base.Errf("%s: the model needs at least one symbol, got %d", name, symbols)
	}
	if err := hmmCheckSequence(name, observations, symbols); err != nil {
		return nil, err
	}
	initial := hmmDrawRow(g, states)
	transition := make([]float64, 0, states*states)
	for range states {
		transition = append(transition, hmmDrawRow(g, states)...)
	}
	emission := make([]float64, 0, states*symbols)
	for range states {
		emission = append(emission, hmmDrawRow(g, symbols)...)
	}
	model, err := NewHiddenMarkovModel(initial, transition, emission)
	if err != nil {
		return nil, err
	}
	prev := math.Inf(-1)
	converged := false
	iterations := hmmMaxIterations
	for iter := 1; iter <= hmmMaxIterations; iter++ {
		logLik, gamma, xi, serr := model.hmmSweep(observations)
		if serr != nil {
			return nil, serr
		}
		move := logLik - prev
		prev = logLik
		if iter > 1 && math.Abs(move) <= hmmTolerance*(1+math.Abs(logLik)) {
			// The model has not moved since the sweep above, so the
			// log likelihood in hand is the fitted model's own.
			converged = true
			iterations = iter
			break
		}
		model = hmmReestimate(model, observations, gamma, xi)
	}
	if !converged {
		// The budget ran out after a re-estimation: the reported
		// likelihood is the final model's own, not its predecessor's.
		logLik, _, _, serr := model.hmmSweep(observations)
		if serr != nil {
			return nil, serr
		}
		prev = logLik
	}
	return &HiddenMarkovFitResult{
		Model:         model,
		LogLikelihood: prev,
		Iterations:    iterations,
		Converged:     converged,
	}, nil
}

// hmmDrawRow draws one distribution row of the given width from the
// flat Dirichlet through the generator: the normalised exponentials
// of uniform draws, the classic construction of the simplex's uniform
// distribution.
func hmmDrawRow(g *core.Generator, width int) []float64 {
	row := make([]float64, width)
	total := 0.0
	for i := range width {
		weight := -math.Log(max(g.Unit(), 1e-300))
		row[i] = weight
		total += weight
	}
	for i := range row {
		row[i] /= total
	}
	return row
}

// hmmSweep runs the forward-backward pass at the model's current
// parameters: the log likelihood, the smoothed posteriors gamma and
// the pairwise posteriors xi the M step reads. xi[t] is the
// states×states row-major table P(state at t, state at t+1 |
// observations) for the step from t to t+1.
func (m *HiddenMarkovModel) hmmSweep(observations []int) (float64, [][]float64, [][]float64, error) {
	states := m.hmmStates()
	width := m.hmmWidth()
	alpha, scales, ferr := m.hmmForward(observations)
	if ferr != nil {
		return 0, nil, nil, ferr
	}
	beta := m.hmmBackward(observations, scales)
	logLik := hmmLogLikelihood(scales)
	gamma := make([][]float64, len(observations))
	for t := range observations {
		row := make([]float64, states)
		total := 0.0
		for k := range states {
			row[k] = alpha[t][k] * beta[t][k]
			total += row[k]
		}
		for k := range states {
			row[k] /= total
		}
		gamma[t] = row
	}
	xi := make([][]float64, len(observations)-1)
	for t := range xi {
		table := make([]float64, states*states)
		total := 0.0
		sym := observations[t+1]
		for k := range states {
			for j := range states {
				v := alpha[t][k] * m.Transition[k*states+j] * m.Emission[j*width+sym] * beta[t+1][j]
				table[k*states+j] = v
				total += v
			}
		}
		for i := range table {
			table[i] /= total
		}
		xi[t] = table
	}
	return logLik, gamma, xi, nil
}

// hmmReestimate applies one Baum-Welch M step: the counts the
// posteriors carry become the new rows, floored and renormalised.
func hmmReestimate(model *HiddenMarkovModel, observations []int, gamma, xi [][]float64) *HiddenMarkovModel {
	states := model.hmmStates()
	symbols := model.hmmSymbols()
	initial := make([]float64, states)
	copy(initial, gamma[0])
	transition := make([]float64, states*states)
	for t := range xi {
		for i, v := range xi[t] {
			transition[i] += v
		}
	}
	for k := range states {
		den := 0.0
		for j := range states {
			den += transition[k*states+j]
		}
		if den == 0 {
			// No transition evidence reached this row, the single-
			// observation sequence being the case in point: dividing the
			// zero fills the row with NaN the constructor refuses, and
			// the sweep after it crashed on the nil model that refusal
			// left. The row keeps its previous estimate instead.
			copy(transition[k*states:(k+1)*states], model.Transition[k*states:(k+1)*states])
			continue
		}
		for j := range states {
			transition[k*states+j] = math.Max(transition[k*states+j]/den, hmmFloor)
		}
		hmmRenormalise(transition[k*states : (k+1)*states])
	}
	emission := make([]float64, states*symbols)
	for t, row := range gamma {
		sym := observations[t]
		for k := range states {
			emission[k*symbols+sym] += row[k]
		}
	}
	for k := range states {
		den := 0.0
		for s := range symbols {
			den += emission[k*symbols+s]
		}
		for s := range symbols {
			emission[k*symbols+s] = math.Max(emission[k*symbols+s]/den, hmmFloor)
		}
		hmmRenormalise(emission[k*symbols : (k+1)*symbols])
	}
	// The floors and renormalisations above keep every row a valid
	// distribution, so the constructor's refusals cannot fire here.
	fitted, _ := NewHiddenMarkovModel(initial, transition, emission)
	return fitted
}

// hmmRenormalise scales one distribution row back to sum 1.
func hmmRenormalise(row []float64) []float64 {
	total := 0.0
	for _, v := range row {
		total += v
	}
	for i := range row {
		row[i] /= total
	}
	return row
}