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

package integrate

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

import "math"

// Differential-algebraic equations in the semi-implicit mass-matrix
// form M·y' = f(t, y) with a singular M whose rank deficiency sits in
// whole zero rows (and, by the same counts, whole zero columns). The
// differential rows carry the implicit Euler scheme, the algebraic
// rows are constraints the implicit relation enforces at every step,
// and the coupled nonlinear system of each step is solved by Newton
// over the numerical Jacobian and the library's dense LU, the
// machinery the ODE stiff solvers already carry.
//
// The honest contract. The solver is first order and takes equal
// steps, the DAE twin of IntegrateBackwardEuler. Consistent initial
// values are the caller's contract: the solver verifies the initial
// residual on the algebraic rows and refuses when it sits beyond the
// tolerance, naming the row, but it does not project a general start
// onto the constraint manifold (a consistent initialiser is a
// documented follow-up). A second-order formula is not offered yet
// either: BDF2 keeps index 1 only when the algebraic start satisfies
// the constraint to second order, which again needs the missing
// projection step. And the index is certified at t0: the solver factors
// the Jacobian of the algebraic rows against the algebraic variables
// and refuses when it is singular, which is the textbook refutation of
// the index-1 contract and is exactly what catches the Cartesian
// pendulum with multipliers, an index-3 system whose per-step solves
// would otherwise converge while the trajectory drifted.

// DAEOptions tunes IntegrateDAE. RelTol ≤ 0 means 1e-6, AbsTol ≤ 0
// means 1e-9, the ODEOptions defaults; they size the Newton tolerance
// of the per-step solves.
type DAEOptions struct {
	RelTol float64
	AbsTol float64
}

// IntegrateDAE integrates the mass-matrix differential-algebraic
// system M·y' = f(t, y) from t0 to t1 over the given number of equal
// implicit Euler steps and returns y(t1). The mass matrix must be
// square and singular, with its rank deficiency carried by whole zero
// rows matched by whole zero columns; the zero rows are the algebraic
// constraints and the zero columns the algebraic variables. Backward
// integration works: a t1 < t0 simply integrates in the negative
// direction. A wrong-shaped state or matrix, a nonsingular matrix, a
// rank deficiency outside whole zero rows, an inconsistent initial
// residual on an algebraic row, a singular algebraic block at t0 (the
// index-1 refutation), or a Newton iteration that cannot converge is
// an error, never a silently truncated trajectory.
func IntegrateDAE(f func(t float64, y *core.Array) (*core.Array, error),
	m *core.Array, t0, t1 float64, y0 *core.Array, steps int, opts DAEOptions) (*core.Array, error) {
	const name = "IntegrateDAE"
	if steps <= 0 {
		return nil, base.Errf("%s: steps must be ≥ 1, got %d", name, steps)
	}
	y, err := odeCheck(name, y0, nil)
	if err != nil {
		return nil, err
	}
	n := len(y)
	mat, err := daeMassMatrix(name, m, n)
	if err != nil {
		return nil, err
	}
	relTol, absTol := opts.RelTol, opts.AbsTol
	if relTol <= 0 {
		relTol = 1e-6
	}
	if absTol <= 0 {
		absTol = 1e-9
	}
	rows, cols, err := daeAlgebraicSets(name, mat)
	if err != nil {
		return nil, err
	}
	// The initial residual on the algebraic rows is the constraint the
	// steps will enforce; a start that violates it has no trajectory.
	fy0, err := odeEval(name, f, t0, y, n, nil)
	if err != nil {
		return nil, err
	}
	for _, row := range rows {
		residual := fy0[row]
		limit := 100 * (absTol + relTol*math.Abs(residual))
		if math.Abs(residual) > limit {
			return nil, base.Errf("%s: the initial residual on algebraic row %d is %g, beyond the consistency tolerance %g: consistent initial values are the caller's contract",
				name, row, residual, limit)
		}
	}
	// The index-1 certificate: the algebraic rows must determine the
	// algebraic variables, which is the nonsingularity of their
	// Jacobian block. A singular block means the mass-matrix rank
	// alone admitted index 1 but the system behaves above it.
	w := &odeWork{}
	jac, err := odeJacobian(name, f, t0, y, w)
	if err != nil {
		return nil, err
	}
	block := make([][]float64, len(rows))
	for i, r := range rows {
		block[i] = make([]float64, len(cols))
		for j, c := range cols {
			block[i][j] = jac[r*n+c]
		}
	}
	base.Factor(block)
	if err := base.CheckSingular(name, block); err != nil {
		return nil, base.Errf("%s: %w: the Jacobian of the algebraic rows against the algebraic variables is singular at t=%g, the system behaves above index 1",
			name, err, t0)
	}
	h := (t1 - t0) / float64(steps)
	myn := make([]float64, n)
	// One Newton result buffer serves every step: daeNewton overwrites
	// it fully before the step copies it into the state, so no step
	// allocates its own.
	zbuf := make([]float64, n)
	t := t0
	if odeArrived(t, t1) {
		return arrayFromVector(y), nil
	}
	// The step times come from the exact grid t0 + i·h, never from an
	// accumulated t += h: the addition's rounding walks over a long
	// run, while each grid point carries a single rounding that stays
	// put.
	for i := range steps {
		tNext := t0 + float64(i+1)*h
		daeMatVec(myn, mat, y)
		if nerr := daeNewton(name, f, w, mat, tNext, h, myn, y, zbuf, absTol, relTol); nerr != nil {
			return nil, nerr
		}
		copy(y, zbuf)
	}
	return arrayFromVector(y), nil
}

// daeMassMatrix reads the mass matrix into a dense float64 row-major
// form and validates its shape against the state length.
func daeMassMatrix(name string, m *core.Array, n int) ([][]float64, error) {
	if m.Dtype() == core.Complex {
		return nil, base.Errf("%s: complex mass matrices are not supported", name)
	}
	shape := m.Shape()
	if m.NDim() != 2 || shape[0] != n || shape[1] != n {
		return nil, base.Errf("%s: the mass matrix must be square of the state's length %d, got shape %s",
			name, n, base.ShapeText(shape))
	}
	mat := make([][]float64, n)
	for i := range n {
		mat[i] = make([]float64, n)
		for j := range n {
			v := m.FloatAt(i*n + j)
			if math.IsNaN(v) || math.IsInf(v, 0) {
				return nil, base.Errf("%s: the mass matrix carries the non-finite entry %g at row %d, column %d", name, v, i, j)
			}
			mat[i][j] = v
		}
	}
	return mat, nil
}

// daeAlgebraicSets returns the algebraic row and column indices: the
// zero rows and the zero columns of the mass matrix, verified to be
// equally numerous and to carry the whole rank deficiency.
func daeAlgebraicSets(name string, mat [][]float64) ([]int, []int, error) {
	n := len(mat)
	worst := 0.0
	for i := range n {
		for j := range n {
			worst = math.Max(worst, math.Abs(mat[i][j]))
		}
	}
	tiny := worst * 1e-14
	var rows, cols []int
	for i := range n {
		zero := true
		for j := range n {
			if math.Abs(mat[i][j]) > tiny {
				zero = false
				break
			}
		}
		if zero {
			rows = append(rows, i)
		}
	}
	for j := range n {
		zero := true
		for i := range n {
			if math.Abs(mat[i][j]) > tiny {
				zero = false
				break
			}
		}
		if zero {
			cols = append(cols, j)
		}
	}
	if len(rows) == 0 {
		return nil, nil, base.Errf("%s: the mass matrix has no zero rows, want a singular mass matrix", name)
	}
	if len(rows) != len(cols) {
		return nil, nil, base.Errf("%s: the mass matrix carries %d zero rows and %d zero columns, want equal counts for the semi-explicit index-1 form",
			name, len(rows), len(cols))
	}
	if rank := daeRank(mat, worst*1e-12); rank != n-len(rows) {
		return nil, nil, base.Errf("%s: the mass matrix carries rank deficiency outside whole zero rows (rank %d with %d zero rows), which the index-1 contract refutes",
			name, rank, len(rows))
	}
	return rows, cols, nil
}

// daeRank counts the pivots Gaussian elimination with partial pivoting
// leaves above the tolerance.
func daeRank(mat [][]float64, tol float64) int {
	n := len(mat)
	a := make([][]float64, n)
	for i := range n {
		a[i] = cloneDenseSlice(mat[i])
	}
	rank, row := 0, 0
	for col := 0; col < n && row < n; col++ {
		piv := row
		for i := row + 1; i < n; i++ {
			if math.Abs(a[i][col]) > math.Abs(a[piv][col]) {
				piv = i
			}
		}
		if math.Abs(a[piv][col]) <= tol {
			continue
		}
		a[piv], a[row] = a[row], a[piv]
		for i := row + 1; i < n; i++ {
			f := a[i][col] / a[row][col]
			for j := col; j < n; j++ {
				a[i][j] -= f * a[row][j]
			}
		}
		rank++
		row++
	}
	return rank
}

// daeMatVec writes the product m·x into dst.
func daeMatVec(dst []float64, m [][]float64, x []float64) {
	for i := range dst {
		s := 0.0
		for j, v := range m[i] {
			s += v * x[j]
		}
		dst[i] = s
	}
}

// daeNewton solves the implicit Euler step equation M·z − M·y_n =
// h·f(tNext, z) for z by Newton over a numerical Jacobian and the
// library's LU solver, the mass-matrix twin of odeNewton. The
// converged state lands in dst, which must not alias seed: a driver
// that hands one buffer down every step allocates it once per run.
// The Jacobian is frozen from the seed and rebuilt twice when
// convergence drags; convergence is measured on the residual against
// the error scale the caller integrates to, an order of magnitude
// below it, but never below the floating-point floor of the
// residual's own terms. An iteration that outlives twenty rounds, or
// a singular Newton matrix, surfaces as errNewtonStalled; an f that
// fails is the fatal error it is.
func daeNewton(name string, f func(t float64, y *core.Array) (*core.Array, error),
	w *odeWork, m [][]float64, tNext, h float64, myn, seed, dst []float64,
	absTol, relTol float64) error {
	n := len(seed)
	w.use(n)
	z := dst
	copy(z, seed)
	mz := w.mz
	for iteration := range 20 {
		daeMatVec(mz, m, z)
		out, err := odeCall(name, f, tNext, z, n, &w.views)
		if err != nil {
			return err
		}
		readVector(w.fzs, out)
		worst, terms := 0.0, 0.0
		for i := range n {
			w.g[i] = mz[i] - h*w.fzs[i] - myn[i]
			worst = math.Max(worst, math.Abs(w.g[i]))
			terms = math.Max(terms, math.Abs(mz[i])+math.Abs(h*w.fzs[i])+math.Abs(myn[i]))
		}
		limit := math.Max(0.1*(absTol+relTol*normInfOfStep(z)), 8*base.EpsF*terms)
		if worst <= limit {
			return nil
		}
		if iteration == 0 || iteration == 4 || iteration == 10 {
			jac, jerr := odeJacobian(name, f, tNext, z, w)
			if jerr != nil {
				return jerr
			}
			// Newton matrix M − h·J, a fresh LU for the frozen
			// Jacobian; the iterations that follow only substitute.
			// Every row is rebuilt entry by entry before the
			// factorisation reads it.
			for i := range n {
				row := w.mat[i]
				for j := range n {
					row[j] = m[i][j] - h*jac[i*n+j]
				}
			}
			w.perm, _ = base.Factor(w.mat)
			if err := base.CheckSingular(name, w.mat); err != nil {
				return base.Errf("%s: %w, singular Newton matrix at t=%g",
					name, errNewtonStalled, tNext)
			}
		}
		for i := range n {
			w.col[i] = -w.g[i]
		}
		odePermuteColumn(w.col, w.perm, w.visited)
		base.SolveColumn(w.mat, w.col)
		for i := range n {
			z[i] += w.col[i]
		}
	}
	return base.Errf("%s: %w at t=%g", name, errNewtonStalled, tNext)
}