"""Tests for moju.piratio.Laws (conservation-law residuals)."""

import pytest
import jax.numpy as jnp
from moju.piratio import Groups, Laws


class TestLawsMassIncompressible:
    """div(u) = 0 residual."""

    def test_zero_gradient_gives_zero_residual(self, rtol, atol):
        """Constant velocity field: u_grad = 0 => div u = 0."""
        u_grad = jnp.zeros((2, 2))
        residual = Laws.mass_incompressible(u_grad)
        assert jnp.allclose(residual, 0.0, rtol=rtol, atol=atol)

    def test_trace_zero_gives_zero_residual(self, rtol, atol):
        """Solenoidal field: trace(u_grad) = 0 => div u = 0."""
        # e.g. u = (y, -x) => u_grad = [[0,1],[-1,0]], trace = 0
        u_grad = jnp.array([[0.0, 1.0], [-1.0, 0.0]])
        residual = Laws.mass_incompressible(u_grad)
        assert jnp.allclose(residual, 0.0, rtol=rtol, atol=atol)

    def test_non_solenoidal_gives_nonzero_residual(self, rtol, atol):
        """trace != 0 => residual != 0."""
        u_grad = jnp.array([[1.0, 0.0], [0.0, 1.0]])  # trace = 2
        residual = Laws.mass_incompressible(u_grad)
        assert not jnp.allclose(residual, 0.0, atol=1e-6)
        assert jnp.allclose(residual, 2.0, rtol=rtol, atol=atol)

    def test_batch_mass_incompressible(self, rtol, atol):
        """Batch of velocity gradients returns 1D residual."""
        u_grad = jnp.zeros((3, 2, 2))  # 3 points, 2x2 gradient each
        residual = Laws.mass_incompressible(u_grad)
        assert residual.shape == (3,)
        assert jnp.allclose(residual, 0.0, rtol=rtol, atol=atol)


class TestLawsLaplaceAndWave:
    """Simple PDE residuals."""

    def test_laplace_equation_residual(self, rtol, atol):
        """Laplace equation: residual = Laplacian(phi). Zero when Laplacian is 0."""
        residual = Laws.laplace_equation(0.0)
        assert jnp.allclose(residual, 0.0, rtol=rtol, atol=atol)

    def test_wave_equation_residual_zero(self, rtol, atol):
        """Wave equation: phi_tt - c^2 Laplacian(phi) = 0 with c = omega*L/st_wave. Zero when both terms cancel."""
        phi_tt = 1.0
        phi_laplacian = 1.0
        omega = 1.0
        L = 1.0
        st_wave = 1.0  # c = omega*L/st_wave = 1.0
        residual = Laws.wave_equation(phi_tt, phi_laplacian, st_wave, omega, L)
        expected = phi_tt - 1.0 * phi_laplacian
        assert jnp.allclose(residual, expected, rtol=rtol, atol=atol)
        assert jnp.allclose(residual, 0.0, rtol=rtol, atol=atol)


class TestLawsFourierConduction:
    """Fourier conduction: T_t - alpha * T_laplacian with optional grid broadcasting."""

    def test_broadcasts_alpha_when_t_1d_and_fields_2d(self, rtol, atol):
        """Mesh time ``t(n_t,)`` with ``T_t``, ``T_laplacian`` (n_t, n_x) — alpha must broadcast."""
        nt, nx = 4, 5
        T_t = jnp.ones((nt, nx))
        T_laplacian = jnp.ones((nt, nx))
        t = jnp.linspace(0.5, 2.0, nt)
        fo = jnp.array(0.1)
        L = jnp.array(1.0)
        out = Laws.fourier_conduction(T_t, T_laplacian, fo, t, L)
        assert out.shape == (nt, nx)
        alpha = fo * (L**2) / t
        alpha2 = jnp.broadcast_to(
            alpha.reshape(nt, 1), (nt, nx)
        )
        assert jnp.allclose(out, T_t - alpha2 * T_laplacian, rtol=rtol, atol=atol)
        assert jnp.all(jnp.isfinite(out))


class TestLawsSchrodingerSteady:
    """Nondimensional steady Schrödinger matches scaled dimensional TISE."""

    def test_nd_residual_equals_k_times_dimensional_residual(self, rtol, atol):
        # Order-one m, L, h_bar so float32 JAX does not underflow h_bar**2.
        lap_dim = jnp.array(-1.5)
        V = jnp.array(2.0)
        E = jnp.array(1.0)
        psi = jnp.array(0.5 + 0.2j)
        m = jnp.array(1.0)
        h_bar = jnp.array(1.0)
        L = jnp.array(1.0)
        K = Groups.schrodinger_kinetic_length_squared(m, L, h_bar)
        lap_L2 = (L**2) * lap_dim
        r_new = Laws.schrodinger_steady(lap_L2, V, E, psi, K)
        r_dim = (-(h_bar**2) / (2 * m)) * lap_dim + (V - E) * psi
        assert jnp.allclose(r_new, K * r_dim, rtol=rtol, atol=atol)


class TestLawsStokesFlow:
    """Stokes flow residual: grad p - (1/Re) Laplacian(u) = 0."""

    def test_stokes_flow_residual_shape(self):
        """Stokes residual is vector of same dimension as p_grad."""
        p_grad = jnp.array([1.0, 0.0])
        u_laplacian = jnp.array([0.0, 0.0])
        re = 100.0
        residual = Laws.stokes_flow(p_grad, u_laplacian, re)
        assert residual.shape == (2,)


class TestLawsNewtonianLaplacianMomentum:
    def test_incompressible_balance_zero(self, rtol, atol):
        u_t = jnp.array([0.0, 0.0])
        u = jnp.array([1.0, 0.0])
        u_grad = jnp.zeros((2, 2))
        p_grad = jnp.array([0.0, 0.0])
        u_laplacian = jnp.array([0.0, 0.0])
        nu = jnp.array(0.1)
        r = Laws.momentum_incompressible_newtonian_laplacian(
            u_t, u, u_grad, p_grad, u_laplacian, nu
        )
        assert jnp.allclose(r, 0.0, rtol=rtol, atol=atol)

    def test_compressible_balance_zero(self, rtol, atol):
        rho = jnp.array(1.2)
        u_t = jnp.array([0.0, 0.0])
        u = jnp.array([0.0, 0.0])
        u_grad = jnp.zeros((2, 2))
        p_grad = jnp.array([0.0, 0.0])
        u_laplacian = jnp.array([0.0, 0.0])
        nu = jnp.array(0.0)
        r = Laws.momentum_compressible_newtonian_laplacian(
            rho, u_t, u, u_grad, p_grad, u_laplacian, nu
        )
        assert jnp.allclose(r, 0.0, rtol=rtol, atol=atol)
