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

package integrate

import (
	"math"
	"slices"

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

// The finite element surface for second-order problems on general
// two-dimensional domains: piecewise-linear (P1) elements on a
// conforming triangular mesh, the stiffness matrix assembled straight
// into the sparse triple format, Dirichlet values eliminated by
// lifting, Neumann boundaries free of charge, and the reduced system
// handed to the sparse Cholesky factorisation the direct-solvers
// surface provides.

// TriangleMesh2D carries a conforming triangular mesh: vertex
// coordinates as x,y pairs and triangles as triples of vertex
// indices. The orientation of a triangle does not matter; a triangle
// with zero area does and is refused at construction.
type TriangleMesh2D struct {
	// Vertices holds x,y for every vertex: two entries per vertex.
	Vertices []float64
	// Triangles holds three vertex indices per triangle.
	Triangles []int64
}

// NewTriangleMesh2D builds a mesh from a vertex table with two
// columns and a triangle table with three columns of vertex indices.
// Indices must lie in range and a degenerate triangle (three
// collinear vertices) is an error: its stiffness contribution is
// undefined.
func NewTriangleMesh2D(vertices *core.Array, triangles *core.Array) (*TriangleMesh2D, error) {
	const name = "NewTriangleMesh2D"
	if vertices.Dtype() == core.Complex || triangles.Dtype() == core.Complex {
		return nil, base.Errf("%s: complex mesh data is not supported", name)
	}
	if vertices.NDim() != 2 || vertices.Shape()[1] != 2 {
		return nil, base.Errf("%s: the vertex table must be rank 2 with two columns, got shape %s", name, base.ShapeText(vertices.Shape()))
	}
	if triangles.Dtype() != core.Int {
		return nil, base.Errf("%s: the triangle table must hold integers, got %s", name, triangles.Dtype())
	}
	if triangles.NDim() != 2 || triangles.Shape()[1] != 3 {
		return nil, base.Errf("%s: the triangle table must be rank 2 with three columns, got shape %s", name, base.ShapeText(triangles.Shape()))
	}
	n := vertices.Shape()[0]
	m := triangles.Shape()[0]
	if n < 3 {
		return nil, base.Errf("%s: a mesh needs at least three vertices, got %d", name, n)
	}
	if m == 0 {
		// An empty triangle table would surface deep in the sparse
		// factorisation on the zero rows of the free nodes, far from
		// the mesh that caused it.
		return nil, base.Errf("%s: the triangle table must not be empty", name)
	}
	mesh := &TriangleMesh2D{Vertices: make([]float64, 2*n), Triangles: make([]int64, 3*m)}
	for i := range 2 * n {
		v := vertices.FloatAt(i)
		if math.IsNaN(v) || math.IsInf(v, 0) {
			return nil, base.Errf("%s: vertex coordinate %d is not finite", name, i)
		}
		mesh.Vertices[i] = v
	}
	for p := range 3 * m {
		idx := triangles.RawInts()[p]
		if idx < 0 || idx >= int64(n) {
			return nil, base.Errf("%s: triangle vertex index %d out of range for %d vertices", name, idx, n)
		}
		mesh.Triangles[p] = idx
	}
	// A triangle with zero area carries no stiffness: refuse it here
	// where the caller can name the triangle, not mid-assembly.
	for t := range m {
		a, b, c := mesh.Triangles[3*t], mesh.Triangles[3*t+1], mesh.Triangles[3*t+2]
		ax, ay := mesh.Vertices[2*a], mesh.Vertices[2*a+1]
		bx, by := mesh.Vertices[2*b], mesh.Vertices[2*b+1]
		cx, cy := mesh.Vertices[2*c], mesh.Vertices[2*c+1]
		if area := math.Abs((bx-ax)*(cy-ay)-(cx-ax)*(by-ay)) / 2; area == 0 {
			return nil, base.Errf("%s: triangle %d is degenerate (zero area)", name, t)
		}
	}
	return mesh, nil
}

// Vertices2 returns the vertex count.
func (m *TriangleMesh2D) Vertices2() int { return len(m.Vertices) / 2 }

// Triangles3 returns the triangle count.
func (m *TriangleMesh2D) Triangles3() int { return len(m.Triangles) / 3 }

// BoundaryEdges returns the mesh's boundary edges as flat pairs of
// vertex indices: an edge belongs to the boundary when exactly one
// triangle carries it. The pairs are sorted, so the result is a pure
// function of the mesh.
func (m *TriangleMesh2D) BoundaryEdges() []int {
	count := make(map[[2]int]int, len(m.Triangles))
	key := func(a, b int) [2]int {
		if a < b {
			return [2]int{a, b}
		}
		return [2]int{b, a}
	}
	for t := 0; t < m.Triangles3(); t++ {
		a, b, c := int(m.Triangles[3*t]), int(m.Triangles[3*t+1]), int(m.Triangles[3*t+2])
		count[key(a, b)]++
		count[key(b, c)]++
		count[key(c, a)]++
	}
	edges := make([]int, 0, 8)
	for e, n := range count {
		if n == 1 {
			edges = append(edges, e[0], e[1])
		}
	}
	slices.Sort(edges)
	return edges
}

// GridTriangleMesh2D builds the structured triangulation of the
// axis-aligned rectangle [x0, x0+width] × [y0, y0+height] with m by n
// cells, two triangles per cell. m and n must both be positive.
func GridTriangleMesh2D(x0, y0, width, height float64, m, n int) (*TriangleMesh2D, error) {
	const name = "GridTriangleMesh2D"
	if m <= 0 || n <= 0 {
		return nil, base.Errf("%s: the cell counts must be positive, got %d by %d", name, m, n)
	}
	// The same guard NewTriangleMesh2D applies to its vertex table: a
	// non-finite extent or origin would lay out vertices at NaN or Inf
	// and only surface mid-factorisation, far from the cause.
	if !(width > 0) || !(height > 0) || math.IsInf(width, 0) || math.IsInf(height, 0) ||
		math.IsNaN(x0) || math.IsInf(x0, 0) || math.IsNaN(y0) || math.IsInf(y0, 0) {
		return nil, base.Errf("%s: the extents must be finite and positive and the origin finite, got origin (%g, %g), extents %g by %g",
			name, x0, y0, width, height)
	}
	vertices := make([]float64, 2*(m+1)*(n+1))
	for j := range n + 1 {
		for i := range m + 1 {
			vertices[2*(j*(m+1)+i)] = x0 + width*float64(i)/float64(m)
			vertices[2*(j*(m+1)+i)+1] = y0 + height*float64(j)/float64(n)
		}
	}
	at := func(i, j int) int64 { return int64(j*(m+1) + i) }
	triangles := make([]int64, 0, 6*m*n)
	for j := range n {
		for i := range m {
			triangles = append(triangles,
				at(i, j), at(i+1, j), at(i+1, j+1),
				at(i, j), at(i+1, j+1), at(i, j+1))
		}
	}
	return &TriangleMesh2D{Vertices: vertices, Triangles: triangles}, nil
}

// FEMPoissonOptions carries the data SolvePoissonFEM2D needs beside
// the mesh and the source: the conductivity, the prescribed boundary
// values, and the optional flux boundary.
type FEMPoissonOptions struct {
	// Kappa is the constant conductivity when KappaFunc is nil. It
	// must be positive.
	Kappa float64
	// KappaFunc, when set, gives the conductivity at a point. It is
	// evaluated at the triangle centroids and must be positive there
	// for every triangle; a non-positive value names the triangle.
	KappaFunc func(x, y float64) float64
	// DirichletNodes lists the vertices with prescribed values and
	// DirichletValues the values in the same order. The nodes leave
	// the system with their rows and columns; at least one is
	// required, because a purely Neumann problem has no unique
	// solution.
	DirichletNodes  []int
	DirichletValues []float64
	// NeumannEdges lists boundary edges as flat pairs of vertex
	// indices and NeumannFlux gives the flux κ∂u/∂n along each edge's
	// outward normal: each edge receives half of length·flux at its
	// midpoint into both endpoints. A nil flux means zero.
	NeumannEdges []int
	NeumannFlux  func(x, y float64) float64
	// Ordering selects the fill-reducing permutation for the sparse
	// Cholesky factorisation. The zero value is the natural order;
	// meshes usually want SparseOrderingReverseCuthillMcKee.
	Ordering linalg.SparseOrdering
}

// SolvePoissonFEM2D solves −∇·(κ∇u) = f on the mesh with
// piecewise-linear elements: the stiffness matrix is assembled per
// triangle (the conductivity evaluated at the centroids when it
// varies), the load is lumped at the vertices from f at the
// centroids, Neumann fluxes are integrated along their edges, and
// Dirichlet values are eliminated by lifting. f may be nil for the
// homogeneous equation.
func SolvePoissonFEM2D(mesh *TriangleMesh2D, f func(x, y float64) float64, opts FEMPoissonOptions) (*core.Array, error) {
	const name = "SolvePoissonFEM2D"
	if mesh == nil {
		return nil, base.Errf("%s: the mesh must not be nil", name)
	}
	n := mesh.Vertices2()
	// With KappaFunc nil the constant conductivity is the value used,
	// so it must be positive and finite; with the field set the
	// constant is a placeholder, but a non-finite one is still refused
	// rather than silently ignored.
	if opts.KappaFunc == nil {
		if !(opts.Kappa > 0) || math.IsInf(opts.Kappa, 0) {
			return nil, base.Errf("%s: the conductivity must be positive, got %g", name, opts.Kappa)
		}
	} else if math.IsNaN(opts.Kappa) || math.IsInf(opts.Kappa, 0) {
		return nil, base.Errf("%s: the conductivity must be positive, got %g", name, opts.Kappa)
	}
	if len(opts.DirichletNodes) != len(opts.DirichletValues) {
		return nil, base.Errf("%s: %d Dirichlet nodes but %d values", name, len(opts.DirichletNodes), len(opts.DirichletValues))
	}
	if len(opts.DirichletNodes) == 0 {
		return nil, base.Errf("%s: a purely Neumann problem has no unique solution; prescribe at least one Dirichlet value", name)
	}
	// The Dirichlet nodes as a dense marker with their prescribed
	// values: the lifting and the unit rows below each visit every
	// assembled entry, and a marker answers those visits in constant
	// time where a set of nodes answered with a hash. A node listed
	// twice keeps its last value and appears once, as it did in the
	// set; the appended order does not reach the assembled system,
	// whose coordinate entries the sparse conversion sorts and merges
	// by coordinate.
	dirichletMark := make([]bool, n)
	dirichletVal := make([]float64, n)
	dirichletNodes := make([]int, 0, len(opts.DirichletNodes))
	for p, d := range opts.DirichletNodes {
		if d < 0 || d >= n {
			return nil, base.Errf("%s: Dirichlet node %d out of range for %d vertices", name, d, n)
		}
		v := opts.DirichletValues[p]
		if math.IsNaN(v) || math.IsInf(v, 0) {
			return nil, base.Errf("%s: Dirichlet value at node %d is not finite", name, d)
		}
		if !dirichletMark[d] {
			dirichletNodes = append(dirichletNodes, d)
		}
		dirichletMark[d] = true
		dirichletVal[d] = v
	}
	if len(opts.NeumannEdges)%2 != 0 {
		return nil, base.Errf("%s: %d Neumann edge indices, want pairs", name, len(opts.NeumannEdges))
	}
	for p := 0; p < len(opts.NeumannEdges); p += 2 {
		a, b := opts.NeumannEdges[p], opts.NeumannEdges[p+1]
		if a < 0 || a >= n || b < 0 || b >= n || a == b {
			return nil, base.Errf("%s: Neumann edge [%d,%d] is not a valid vertex pair", name, a, b)
		}
	}
	// Assembly: nine entries per triangle, symmetric by construction;
	// the load is lumped one third of the triangle area to each of
	// its vertices, with the conductivity evaluated at the centroid
	// when it varies.
	entries := make([]float64, 0, 9*mesh.Triangles3())
	rows := make([]int, 0, 9*mesh.Triangles3())
	cols := make([]int, 0, 9*mesh.Triangles3())
	load := make([]float64, n)
	for t := 0; t < mesh.Triangles3(); t++ {
		a, b, c := int(mesh.Triangles[3*t]), int(mesh.Triangles[3*t+1]), int(mesh.Triangles[3*t+2])
		ax, ay := mesh.Vertices[2*a], mesh.Vertices[2*a+1]
		bx, by := mesh.Vertices[2*b], mesh.Vertices[2*b+1]
		cx, cy := mesh.Vertices[2*c], mesh.Vertices[2*c+1]
		area := math.Abs((bx-ax)*(cy-ay)-(cx-ax)*(by-ay)) / 2
		kappa := opts.Kappa
		if opts.KappaFunc != nil {
			kappa = opts.KappaFunc((ax+bx+cx)/3, (ay+by+cy)/3)
			if !(kappa > 0) || math.IsNaN(kappa) || math.IsInf(kappa, 0) {
				return nil, base.Errf("%s: the conductivity at triangle %d is %g, want positive", name, t, kappa)
			}
		}
		// The gradient basis: b are the y differences, c the x
		// differences, and K = κ/(4A)·(b⊗b + c⊗c).
		bb := [3]float64{by - cy, cy - ay, ay - by}
		cc := [3]float64{cx - bx, ax - cx, bx - ax}
		nodes := [3]int{a, b, c}
		for i := range 3 {
			for j := range 3 {
				v := kappa * (bb[i]*bb[j] + cc[i]*cc[j]) / (4 * area)
				rows = append(rows, nodes[i])
				cols = append(cols, nodes[j])
				entries = append(entries, v)
			}
		}
		if f != nil {
			fv := f((ax+bx+cx)/3, (ay+by+cy)/3)
			// A non-finite source value would flow into the load and the
			// solve would publish an all-NaN solution with a nil error,
			// the breach every other integrator here refuses up front.
			if math.IsNaN(fv) || math.IsInf(fv, 0) {
				return nil, base.Errf("%s: the source returned the non-finite value %g at triangle %d", name, fv, t)
			}
			contribution := area / 3 * fv
			load[a] += contribution
			load[b] += contribution
			load[c] += contribution
		}
	}
	// Neumann fluxes: half of length·flux into each endpoint of every
	// listed edge, the flux evaluated at the edge midpoint.
	if len(opts.NeumannEdges) > 0 && opts.NeumannFlux != nil {
		for p := 0; p < len(opts.NeumannEdges); p += 2 {
			a, b := opts.NeumannEdges[p], opts.NeumannEdges[p+1]
			ax, ay := mesh.Vertices[2*a], mesh.Vertices[2*a+1]
			bx, by := mesh.Vertices[2*b], mesh.Vertices[2*b+1]
			length := math.Hypot(bx-ax, by-ay)
			fv := opts.NeumannFlux((ax+bx)/2, (ay+by)/2)
			// A non-finite flux lands in the load like a non-finite
			// source, so the same refusal answers it.
			if math.IsNaN(fv) || math.IsInf(fv, 0) {
				return nil, base.Errf("%s: the Neumann flux returned the non-finite value %g on edge [%d, %d]", name, fv, a, b)
			}
			flux := length / 2 * fv
			load[a] += flux
			load[b] += flux
		}
	}
	// Dirichlet lifting: the known boundary values move to the right
	// hand side, then their rows and columns leave the system as
	// unit rows.
	for p, i := range rows {
		if j := cols[p]; dirichletMark[j] {
			load[i] -= entries[p] * dirichletVal[j]
		}
	}
	keptRows := make([]int64, 0, len(rows))
	keptCols := make([]int64, 0, len(rows))
	keptVals := make([]float64, 0, len(rows))
	for p := range rows {
		i, j := rows[p], cols[p]
		if dirichletMark[i] || dirichletMark[j] {
			continue
		}
		keptRows = append(keptRows, int64(i))
		keptCols = append(keptCols, int64(j))
		keptVals = append(keptVals, entries[p])
	}
	for _, d := range dirichletNodes {
		keptRows = append(keptRows, int64(d))
		keptCols = append(keptCols, int64(d))
		keptVals = append(keptVals, 1)
		load[d] = dirichletVal[d]
	}
	indices, err := core.FromInts(pairInts(keptRows, keptCols), len(keptVals), 2)
	if err != nil {
		return nil, base.Errf("%s: %w", name, err)
	}
	coo, err := core.NewSparseCOO(indices, fromSlice(keptVals, len(keptVals)), []int{n, n})
	if err != nil {
		return nil, base.Errf("%s: %w", name, err)
	}
	order := opts.Ordering
	factor, err := linalg.NewSparseCholesky(coo, order)
	if err != nil {
		return nil, base.Errf("%s: %w", name, err)
	}
	rhs := core.New(core.Float, []int{n}...)
	copy(rhs.RawFloats(), load)
	return factor.Solve(rhs)
}

// pairInts interleaves row and column indices into the index table
// the sparse coordinate format expects.
func pairInts(rows, cols []int64) []int64 {
	out := make([]int64, 2*len(rows))
	for p := range rows {
		out[2*p] = rows[p]
		out[2*p+1] = cols[p]
	}
	return out
}