package mathkit

import (
	"fmt"
	"math"
)

// Matrix is a dense matrix stored row-major: Matrix[i][j] is row i,
// column j. Every row must have the same length — functions that take
// a Matrix validate that with Matrix.Dims before doing any math, so a
// ragged input is reported as an error rather than panicking.
type Matrix [][]float64

// Dims returns the matrix's row and column counts and an error if the
// rows aren't all the same length (a ragged "matrix" isn't one).
func (m Matrix) Dims() (rows, cols int, err error) {
	rows = len(m)
	if rows == 0 {
		return 0, 0, fmt.Errorf("empty matrix")
	}
	cols = len(m[0])
	for i, row := range m {
		if len(row) != cols {
			return 0, 0, fmt.Errorf("ragged matrix: row 0 has %d columns, row %d has %d", cols, i, len(row))
		}
	}
	return rows, cols, nil
}

// Identity returns the n×n identity matrix.
func Identity(n int) Matrix {
	m := make(Matrix, n)
	for i := range m {
		m[i] = make([]float64, n)
		m[i][i] = 1
	}
	return m
}

func sameDims(a, b Matrix) (rows, cols int, err error) {
	ra, ca, err := a.Dims()
	if err != nil {
		return 0, 0, err
	}
	rb, cb, err := b.Dims()
	if err != nil {
		return 0, 0, err
	}
	if ra != rb || ca != cb {
		return 0, 0, fmt.Errorf("shape mismatch: %dx%d vs %dx%d", ra, ca, rb, cb)
	}
	return ra, ca, nil
}

// MatAdd returns a+b, element-wise.
func MatAdd(a, b Matrix) (Matrix, error) {
	rows, cols, err := sameDims(a, b)
	if err != nil {
		return nil, fmt.Errorf("matrix_add: %w", err)
	}
	out := make(Matrix, rows)
	for i := 0; i < rows; i++ {
		out[i] = make([]float64, cols)
		for j := 0; j < cols; j++ {
			out[i][j] = a[i][j] + b[i][j]
		}
	}
	return out, nil
}

// MatSub returns a-b, element-wise.
func MatSub(a, b Matrix) (Matrix, error) {
	rows, cols, err := sameDims(a, b)
	if err != nil {
		return nil, fmt.Errorf("matrix_sub: %w", err)
	}
	out := make(Matrix, rows)
	for i := 0; i < rows; i++ {
		out[i] = make([]float64, cols)
		for j := 0; j < cols; j++ {
			out[i][j] = a[i][j] - b[i][j]
		}
	}
	return out, nil
}

// MatMul returns the matrix product a×b (standard inner-dimension
// matching: a is m×n, b must be n×p).
func MatMul(a, b Matrix) (Matrix, error) {
	ra, ca, err := a.Dims()
	if err != nil {
		return nil, fmt.Errorf("matrix_mul: left operand: %w", err)
	}
	rb, cb, err := b.Dims()
	if err != nil {
		return nil, fmt.Errorf("matrix_mul: right operand: %w", err)
	}
	if ca != rb {
		return nil, fmt.Errorf("matrix_mul: inner dimensions must match: %dx%d times %dx%d", ra, ca, rb, cb)
	}
	out := make(Matrix, ra)
	for i := 0; i < ra; i++ {
		out[i] = make([]float64, cb)
		for j := 0; j < cb; j++ {
			var sum float64
			for k := 0; k < ca; k++ {
				sum += a[i][k] * b[k][j]
			}
			out[i][j] = sum
		}
	}
	return out, nil
}

// Scale multiplies every element of m by k.
func Scale(m Matrix, k float64) Matrix {
	out := make(Matrix, len(m))
	for i, row := range m {
		out[i] = make([]float64, len(row))
		for j, v := range row {
			out[i][j] = v * k
		}
	}
	return out
}

// Transpose returns m^T.
func Transpose(m Matrix) (Matrix, error) {
	rows, cols, err := m.Dims()
	if err != nil {
		return nil, fmt.Errorf("matrix_transpose: %w", err)
	}
	out := make(Matrix, cols)
	for j := 0; j < cols; j++ {
		out[j] = make([]float64, rows)
		for i := 0; i < rows; i++ {
			out[j][i] = m[i][j]
		}
	}
	return out, nil
}

// Trace returns the sum of the diagonal of a square matrix.
func Trace(m Matrix) (float64, error) {
	rows, cols, err := m.Dims()
	if err != nil {
		return 0, fmt.Errorf("matrix_trace: %w", err)
	}
	if rows != cols {
		return 0, fmt.Errorf("matrix_trace: matrix must be square, got %dx%d", rows, cols)
	}
	var sum float64
	for i := 0; i < rows; i++ {
		sum += m[i][i]
	}
	return sum, nil
}

// clone makes a deep copy so Gaussian-elimination routines never
// mutate the caller's matrix.
func (m Matrix) clone() Matrix {
	out := make(Matrix, len(m))
	for i, row := range m {
		out[i] = append([]float64(nil), row...)
	}
	return out
}

// luDecompose runs Gaussian elimination with partial pivoting on a
// clone of m, returning the resulting upper-triangular-ish matrix and
// the sign of the permutation (needed to get the determinant's sign
// right) and the pivots used. Both Determinant and Inverse build on
// this one routine rather than duplicating elimination logic.
func luDecompose(m Matrix) (a Matrix, sign float64, err error) {
	rows, cols, err := m.Dims()
	if err != nil {
		return nil, 0, err
	}
	if rows != cols {
		return nil, 0, fmt.Errorf("matrix must be square, got %dx%d", rows, cols)
	}
	n := rows
	a = m.clone()
	sign = 1
	for col := 0; col < n; col++ {
		// Partial pivot: swap in the row with the largest absolute
		// value in this column, for numerical stability.
		pivotRow := col
		best := math.Abs(a[col][col])
		for r := col + 1; r < n; r++ {
			if v := math.Abs(a[r][col]); v > best {
				best, pivotRow = v, r
			}
		}
		if best < 1e-12 {
			return a, 0, nil // singular; determinant is 0, caller decides
		}
		if pivotRow != col {
			a[col], a[pivotRow] = a[pivotRow], a[col]
			sign = -sign
		}
		for r := col + 1; r < n; r++ {
			factor := a[r][col] / a[col][col]
			for c := col; c < n; c++ {
				a[r][c] -= factor * a[col][c]
			}
		}
	}
	return a, sign, nil
}

// Determinant computes det(m) via LU (Gaussian elimination with
// partial pivoting), O(n^3) — the standard numerically-stable
// approach, versus the O(n!) cofactor expansion a direct reading of
// the definition would suggest.
func Determinant(m Matrix) (float64, error) {
	a, sign, err := luDecompose(m)
	if err != nil {
		return 0, fmt.Errorf("matrix_determinant: %w", err)
	}
	n := len(a)
	det := sign
	for i := 0; i < n; i++ {
		det *= a[i][i]
	}
	return det, nil
}

// Inverse computes m^-1 via Gauss-Jordan elimination on [m | I].
// Returns an error for a non-square or singular (determinant ≈ 0)
// matrix — ML code hitting a singular matrix (e.g. a rank-deficient
// design matrix in least squares) usually wants to know that, not a
// silently garbage answer.
func Inverse(m Matrix) (Matrix, error) {
	rows, cols, err := m.Dims()
	if err != nil {
		return nil, fmt.Errorf("matrix_inverse: %w", err)
	}
	if rows != cols {
		return nil, fmt.Errorf("matrix_inverse: matrix must be square, got %dx%d", rows, cols)
	}
	n := rows
	aug := make(Matrix, n)
	for i := 0; i < n; i++ {
		aug[i] = make([]float64, 2*n)
		copy(aug[i], m[i])
		aug[i][n+i] = 1
	}
	for col := 0; col < n; col++ {
		pivotRow := col
		best := math.Abs(aug[col][col])
		for r := col + 1; r < n; r++ {
			if v := math.Abs(aug[r][col]); v > best {
				best, pivotRow = v, r
			}
		}
		if best < 1e-12 {
			return nil, fmt.Errorf("matrix_inverse: matrix is singular (not invertible)")
		}
		aug[col], aug[pivotRow] = aug[pivotRow], aug[col]
		pivot := aug[col][col]
		for c := 0; c < 2*n; c++ {
			aug[col][c] /= pivot
		}
		for r := 0; r < n; r++ {
			if r == col {
				continue
			}
			factor := aug[r][col]
			for c := 0; c < 2*n; c++ {
				aug[r][c] -= factor * aug[col][c]
			}
		}
	}
	out := make(Matrix, n)
	for i := 0; i < n; i++ {
		out[i] = append([]float64(nil), aug[i][n:]...)
	}
	return out, nil
}

// Rank computes the matrix rank as the number of non-negligible
// pivots found by Gaussian elimination (works for non-square
// matrices too, unlike Determinant/Inverse).
func Rank(m Matrix) (int, error) {
	rows, cols, err := m.Dims()
	if err != nil {
		return 0, fmt.Errorf("matrix_rank: %w", err)
	}
	a := m.clone()
	rank := 0
	for col := 0; col < cols && rank < rows; col++ {
		pivotRow := -1
		best := 1e-9
		for r := rank; r < rows; r++ {
			if v := math.Abs(a[r][col]); v > best {
				best, pivotRow = v, r
			}
		}
		if pivotRow == -1 {
			continue
		}
		a[rank], a[pivotRow] = a[pivotRow], a[rank]
		for r := rank + 1; r < rows; r++ {
			factor := a[r][col] / a[rank][col]
			for c := col; c < cols; c++ {
				a[r][c] -= factor * a[rank][c]
			}
		}
		rank++
	}
	return rank, nil
}

// MatVecMul returns m×v for a column vector v.
func MatVecMul(m Matrix, v []float64) ([]float64, error) {
	rows, cols, err := m.Dims()
	if err != nil {
		return nil, fmt.Errorf("matrix-vector multiply: %w", err)
	}
	if cols != len(v) {
		return nil, fmt.Errorf("matrix-vector multiply: matrix has %d columns but vector has %d entries", cols, len(v))
	}
	out := make([]float64, rows)
	for i := 0; i < rows; i++ {
		var sum float64
		for j := 0; j < cols; j++ {
			sum += m[i][j] * v[j]
		}
		out[i] = sum
	}
	return out, nil
}

// Dot returns the dot (inner) product of two equal-length vectors.
func Dot(a, b []float64) (float64, error) {
	if len(a) != len(b) {
		return 0, fmt.Errorf("dot_product: vectors have different lengths (%d vs %d)", len(a), len(b))
	}
	var sum float64
	for i := range a {
		sum += a[i] * b[i]
	}
	return sum, nil
}

// Norm computes the Lp norm of a vector. p=2 is Euclidean (the
// default any ML codebase means when it just says "norm"); p=1 is
// Manhattan; p=math.Inf(1) is the max-absolute-value norm.
func Norm(v []float64, p float64) float64 {
	if math.IsInf(p, 1) {
		var m float64
		for _, x := range v {
			if a := math.Abs(x); a > m {
				m = a
			}
		}
		return m
	}
	var sum float64
	for _, x := range v {
		sum += math.Pow(math.Abs(x), p)
	}
	return math.Pow(sum, 1/p)
}

// DominantEigen estimates the eigenvalue of largest absolute value
// and its eigenvector via power iteration — the workhorse behind
// things like PageRank and quick PCA-component estimates, and much
// simpler/more numerically robust to implement from scratch than a
// full QR eigendecomposition, which is the right trade-off for a
// hand-rolled numerical toolkit. iterations<=0 defaults to 200, which
// converges for any matrix with a well-separated dominant eigenvalue
// to within float64 precision.
func DominantEigen(m Matrix, iterations int) (eigenvalue float64, eigenvector []float64, err error) {
	rows, cols, err := m.Dims()
	if err != nil {
		return 0, nil, fmt.Errorf("dominant_eigen: %w", err)
	}
	if rows != cols {
		return 0, nil, fmt.Errorf("dominant_eigen: matrix must be square, got %dx%d", rows, cols)
	}
	if iterations <= 0 {
		iterations = 200
	}
	n := rows
	v := make([]float64, n)
	for i := range v {
		v[i] = 1 / math.Sqrt(float64(n)) // deterministic, reproducible start
	}
	var lambda float64
	for iter := 0; iter < iterations; iter++ {
		next, err := MatVecMul(m, v)
		if err != nil {
			return 0, nil, err
		}
		norm := Norm(next, 2)
		if norm < 1e-15 {
			return 0, nil, fmt.Errorf("dominant_eigen: did not converge (matrix may be nilpotent or zero)")
		}
		for i := range next {
			next[i] /= norm
		}
		v = next
		mv, _ := MatVecMul(m, v)
		lambda, _ = Dot(v, mv)
	}
	return lambda, v, nil
}