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

// Command pde solves the two canonical one-dimensional partial
// differential equations: the heat equation by Crank-Nicolson and the
// wave equation by velocity Verlet. Both start from the same Gaussian
// pulse on a rod, and the diagnostics show diffusion flattening the
// pulse while the wave keeps its shape and travels.
//
// Usage: go run ./examples/pde
package main

import (
	"fmt"
	"log"
	"math"

	"sourcedock.dev/petrbalvin/tensor"
)

func main() {
	const (
		n  = 201
		dx = 0.01 // metres, a 2 m rod
		k  = 1e-3 // thermal diffusivity, m^2/s
		c  = 0.5  // wave speed, m/s
	)
	pulse := make([]float64, n)
	for i := range n {
		x := float64(i)*dx - 1.0
		pulse[i] = math.Exp(-(x * x) / 0.01)
	}
	u0, err := tensor.FromFloats(pulse, n)
	if err != nil {
		log.Fatal(err)
	}

	fmt.Println("heat equation (Crank-Nicolson, ends held at zero):")
	fmt.Println("  time     peak      mean (interior)")
	for _, tf := range []float64{0, 0.5, 2, 5} {
		var row *tensor.Array
		if tf == 0 {
			row = u0 // the initial pulse itself
		} else {
			u, err := tensor.IntegrateHeat1D(u0, k, dx, tf, tf/400, 5, 0, 0)
			if err != nil {
				log.Fatal(err)
			}
			// The last sample row holds the final state. The ends are
			// Dirichlet zeros, so heat drains out of the rod once the
			// pulse reaches them; by the last time printed it has not,
			// which is why the interior mean barely moves.
			row, err = tensor.Slice(u, 0, 4, 5)
			if err != nil {
				log.Fatal(err)
			}
		}
		mx, err := tensor.Max(row)
		if err != nil {
			log.Fatal(err)
		}
		mean, err := tensor.Mean(row)
		if err != nil {
			log.Fatal(err)
		}
		fmt.Printf("  %.2f s   %7.4f   %7.4f\n", tf, mx.Float(), mean)
	}

	fmt.Println()
	fmt.Println("wave equation (velocity Verlet, fixed ends):")
	fmt.Println("  time     peak     position of the peak")
	rest, err := tensor.Zeros(tensor.Float, n)
	if err != nil {
		log.Fatal(err)
	}
	for _, tf := range []float64{0, 0.5, 1.0, 1.5} {
		var row *tensor.Array
		if tf == 0 {
			row = u0
		} else {
			u, err := tensor.IntegrateWave1D(u0, rest, c, dx, tf, tf/600, 5)
			if err != nil {
				log.Fatal(err)
			}
			row, err = tensor.Slice(u, 0, 4, 5)
			if err != nil {
				log.Fatal(err)
			}
		}
		mx, err := tensor.Max(row)
		if err != nil {
			log.Fatal(err)
		}
		peak, peakVal := 0, -1.0
		for i := range n {
			if v := row.FloatAt(i); v > peakVal {
				peak, peakVal = i, v
			}
		}
		fmt.Printf("  %.2f s   %6.4f   x = %.2f m\n",
			tf, mx.Float(), float64(peak)*dx-1.0)
	}
}