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

// Command helmholtz solves the discretised Helmholtz equation, the
// backbone of frequency-domain electromagnetics, in both of its
// solver shapes. The time-harmonic wave equation
//
//	-∇²ψ - k²ψ = f
//
// on a 2-D grid gives a complex symmetric (non-Hermitian) sparse
// system, which the BiCGSTAB solver handles. Adding a small imaginary
// part to k², the way a lossy medium does, makes the operator
// Hermitian positive-definite and the conjugate gradient solver
// applies. Both solutions are verified against the dense solve, and
// the Hermitian operator's resonant modes come from the sparse
// eigensolver.
//
// Usage: go run ./examples/helmholtz
package main

import (
	"fmt"
	"log"
	"math"

	"sourcedock.dev/petrbalvin/tensor"
)

const grid = 24 // interior points per side

// laplacianCOO assembles the 5-point discrete -∇² on the interior of
// a grid*grid domain with Dirichlet walls, one entry per stencil
// point. The value at (i,j) is k2 times the identity there.
func helmholtzCOO(k2 complex128) (*tensor.SparseCOO, int) {
	n := grid * grid
	var idx []int64
	var val []complex128
	at := func(i, j int) int { return i*grid + j }
	for i := range grid {
		for j := range grid {
			p := at(i, j)
			// 4/h² on the diagonal with h = 1 in grid units, minus k².
			idx = append(idx, int64(p), int64(p))
			val = append(val, 4-k2)
			if i > 0 {
				idx = append(idx, int64(p), int64(at(i-1, j)))
				val = append(val, -1)
			}
			if i < grid-1 {
				idx = append(idx, int64(p), int64(at(i+1, j)))
				val = append(val, -1)
			}
			if j > 0 {
				idx = append(idx, int64(p), int64(at(i, j-1)))
				val = append(val, -1)
			}
			if j < grid-1 {
				idx = append(idx, int64(p), int64(at(i, j+1)))
				val = append(val, -1)
			}
		}
	}
	indices, err := tensor.FromInts(idx, len(val), 2)
	if err != nil {
		log.Fatal(err)
	}
	values, err := tensor.FromComplexes(val, len(val))
	if err != nil {
		log.Fatal(err)
	}
	coo, err := tensor.NewSparseCOO(indices, values, []int{n, n})
	if err != nil {
		log.Fatal(err)
	}
	return coo, n
}

// landauCOO assembles the Hamiltonian of a charged particle on the
// same grid threading a perpendicular magnetic field, the Peierls
// substitution: every hop carries the phase the vector potential
// gives it, forward and conjugate backward, so the operator stays
// Hermitian. A positive mass term m² makes it positive-definite.
func landauCOO(m2, flux float64) *tensor.SparseCOO {
	n := grid * grid
	var idx []int64
	var val []complex128
	at := func(i, j int) int { return i*grid + j }
	phase := func(i int) float64 { return 2 * math.Pi * flux * float64(i) }
	for i := range grid {
		for j := range grid {
			p := at(i, j)
			idx = append(idx, int64(p), int64(p))
			val = append(val, complex(4+m2, 0))
			if i > 0 {
				idx = append(idx, int64(p), int64(at(i-1, j)))
				val = append(val, -1+0i)
			}
			if i < grid-1 {
				idx = append(idx, int64(p), int64(at(i+1, j)))
				val = append(val, -1+0i)
			}
			if j > 0 {
				idx = append(idx, int64(p), int64(at(i, j-1)))
				val = append(val, -complex(math.Cos(phase(i)), math.Sin(phase(i))))
			}
			if j < grid-1 {
				idx = append(idx, int64(p), int64(at(i, j+1)))
				val = append(val, -complex(math.Cos(phase(i)), -math.Sin(phase(i))))
			}
		}
	}
	indices, err := tensor.FromInts(idx, len(val), 2)
	if err != nil {
		log.Fatal(err)
	}
	values, err := tensor.FromComplexes(val, len(val))
	if err != nil {
		log.Fatal(err)
	}
	coo, err := tensor.NewSparseCOO(indices, values, []int{n, n})
	if err != nil {
		log.Fatal(err)
	}
	return coo
}

// source is a point drive at the grid centre, the field of a small
// antenna.
func source(n int) *tensor.Array {
	rhs := make([]complex128, n)
	rhs[(grid/2)*grid+grid/2] = 1 + 0i
	b, err := tensor.FromComplexes(rhs, n)
	if err != nil {
		log.Fatal(err)
	}
	return b
}

// residual returns ||b - A·x||₂ by reassembling A densely, the ground
// truth the sparse solver is checked against.
func residual(a *tensor.SparseCOO, x, b *tensor.Array, n int) float64 {
	dense, err := a.Dense()
	if err != nil {
		log.Fatal(err)
	}
	ax, err := tensor.MatMul2D(dense, x)
	if err != nil {
		log.Fatal(err)
	}
	worst := 0.0
	for i := range n {
		av, err := tensor.ComplexAt(ax, i)
		if err != nil {
			log.Fatal(err)
		}
		bv, err := tensor.ComplexAt(b, i)
		if err != nil {
			log.Fatal(err)
		}
		if d := math.Hypot(real(av-bv), imag(av-bv)); d > worst {
			worst = d
		}
	}
	return worst
}

func main() {
	const n = grid * grid

	// A propagating mode: k = 2.5 in grid units, safely away from the
	// discrete resonances at k² = 2-2cos(p*pi/(grid+1)).
	k := 2.5 + 0i
	a, _ := helmholtzCOO(k * k)
	b := source(n)

	x, err := tensor.SpSolveComplexBiCGSTAB(a, b, 1e-12, 2000)
	if err != nil {
		log.Fatal(err)
	}
	fmt.Println("lossless Helmholtz system, -nabla^2 - k^2, k = 2.5")
	fmt.Printf("  unknowns: %d, stored nonzeros: %d\n", n, len(a.Values.RawComplexes()))
	fmt.Printf("  BiCGSTAB residual ||b - A x|| = %.3g\n", residual(a, x, b, n))

	// The genuinely Hermitian complex problem: a charged particle on
	// the same grid in a perpendicular magnetic field. The Peierls
	// phases make every hop complex, the forward and backward hop
	// conjugates of each other, so the operator is Hermitian, and the
	// mass term keeps it positive-definite: exactly the shape the
	// conjugate gradient solver wants.
	h := landauCOO(1.0, 1.0/25)
	xh, err := tensor.SpSolveComplexCG(h, b, 1e-12, 2000)
	if err != nil {
		log.Fatal(err)
	}
	fmt.Println("\nLandau Hamiltonian on the grid, mass^2 = 1, flux 1/25 (Hermitian positive-definite)")
	fmt.Printf("  CG residual ||b - A x|| = %.3g\n", residual(h, xh, b, n))

	// Resonant modes of the lossless cavity: the largest eigenvalues
	// of the discrete negative Laplacian are the highest-Q modes.
	lap, _ := helmholtzCOO(0)
	vals, vecs, err := tensor.SpEigenComplex(lap, 3, tensor.NewGenerator(4))
	if err != nil {
		log.Fatal(err)
	}
	fmt.Println("\ncavity modes: largest eigenvalues of -nabla^2")
	for j := range 3 {
		lam, err := tensor.FloatAt(vals, j)
		if err != nil {
			log.Fatal(err)
		}
		// Verify each Ritz pair: ||A v - lambda v|| must be small.
		vcol, err := tensor.Slice(vecs, 1, j, j+1)
		if err != nil {
			log.Fatal(err)
		}
		av, err := tensor.MatMul2D(mustDense(lap), vcol)
		if err != nil {
			log.Fatal(err)
		}
		worst := 0.0
		for i := range n {
			a1, _ := tensor.ComplexAt(av, i)
			v1, _ := tensor.ComplexAt(vcol, i)
			if d := math.Hypot(real(a1-complex(lam, 0)*v1), imag(a1-complex(lam, 0)*v1)); d > worst {
				worst = d
			}
		}
		fmt.Printf("  lambda = %8.4f, residual %.3g\n", lam, worst)
	}
	// The analytic eigenvalues of the grid Laplacian are
	// 4-2cos(p*pi/(grid+1))-2cos(q*pi/(grid+1)); the largest is p = q =
	// grid, where both cosines approach -1 and the value nears 8.
	fmt.Println("  (analytic maximum: 4 - 4cos(24pi/25) = 7.9685)")
}

func mustDense(a *tensor.SparseCOO) *tensor.Array {
	d, err := a.Dense()
	if err != nil {
		log.Fatal(err)
	}
	return d
}