import math

import pytest
import torch
import kernels

surf = kernels.get_kernel("phanerozoic/surf", version=1, trust_remote_code=True)
requires_cuda = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
G = 9.81


def _bumpy_bed(N, L):
    xs = torch.linspace(0, L, N, device="cuda")
    X, Y = torch.meshgrid(xs, xs, indexing="xy")
    b = (-3.0 + 0.8 * torch.exp(-((X - L / 2) ** 2 + (Y - L / 2) ** 2) / 200)
         + 0.5 * torch.sin(X * 0.3) * torch.cos(Y * 0.2))
    return b.float().contiguous(), X, Y


@requires_cuda
def test_nsw_lake_at_rest():
    """Still water over an uneven bed must stay still: the well-balancing test.

    A scheme that is not well-balanced spins up spurious currents from the bed
    slope and never settles.
    """
    N, L = 192, 100.0
    dx = dy = L / N
    b, _, _ = _bumpy_bed(N, L)
    h = (0.0 - b).clamp(min=0)
    hu = torch.zeros_like(h)
    hv = torch.zeros_like(h)
    dt = 0.3 * dx / math.sqrt(G * float(h.max()))
    for _ in range(200):
        h, hu, hv = surf.nsw_step(h, hu, hv, b, dx, dy, dt)
    u = hu / h.clamp(min=1e-3)
    v = hv / h.clamp(min=1e-3)
    assert float(torch.sqrt(u * u + v * v).max()) < 1e-4
    assert float((h + b).abs().max()) < 1e-4


@requires_cuda
def test_nsw_dam_break_bounded_and_finite():
    """A radial dam break stays finite and bounded (shock capturing works)."""
    N, L = 192, 100.0
    dx = dy = L / N
    xs = torch.linspace(0, L, N, device="cuda")
    X, Y = torch.meshgrid(xs, xs, indexing="xy")
    b = torch.full((N, N), -5.0, device="cuda")
    r = torch.sqrt((X - L / 2) ** 2 + (Y - L / 2) ** 2)
    eta = torch.where(r < 15, torch.ones_like(r), torch.zeros_like(r))
    h = (eta - b).clamp(min=0).contiguous()
    hu = torch.zeros_like(h)
    hv = torch.zeros_like(h)
    dt = 0.3 * dx / math.sqrt(G * float(h.max()))
    for _ in range(150):
        h, hu, hv = surf.nsw_step(h, hu, hv, b, dx, dy, dt)
    eta = h + b
    assert torch.isfinite(eta).all()
    assert -1.0 < float(eta.min()) and float(eta.max()) < 1.0


@requires_cuda
def test_nsw_mass_conservation_closed_bump():
    """With no in/outflow the total water volume is conserved."""
    N, L = 128, 60.0
    dx = dy = L / N
    b = torch.full((N, N), -4.0, device="cuda")
    xs = torch.linspace(0, L, N, device="cuda")
    X, Y = torch.meshgrid(xs, xs, indexing="xy")
    eta = 0.3 * torch.exp(-((X - L / 2) ** 2 + (Y - L / 2) ** 2) / 40)
    h = (eta - b).clamp(min=0).contiguous()
    hu = torch.zeros_like(h)
    hv = torch.zeros_like(h)
    m0 = float(h.sum())
    dt = 0.3 * dx / math.sqrt(G * float(h.max()))
    for _ in range(120):
        h, hu, hv = surf.nsw_step(h, hu, hv, b, dx, dy, dt)
    assert abs(float(h.sum()) - m0) / m0 < 1e-4


@requires_cuda
def test_sph_isolated_particles_freefall():
    """SPH particles with no neighbours inside 2h feel gravity and nothing else."""
    pos = torch.tensor([[0.0, 0.0, 0.0], [0.5, 0.0, 0.0]], device="cuda")
    vel = torch.zeros_like(pos)
    rho = torch.full((2,), 1000.0, device="cuda")
    ptype = torch.zeros(2, dtype=torch.int32, device="cuda")
    dp = 0.01
    h = 1.5 * dp
    csz = 2 * h
    origin = (-0.1, -0.1, -0.1)
    grid = (int(0.8 / csz) + 2,) * 3
    perm, cs, ce = surf.build_cells(pos, origin, csz, grid)
    pos, vel, rho, ptype = pos[perm], vel[perm], rho[perm], ptype[perm]
    drho, acc, xsph = surf.sph_forces(pos, vel, rho, ptype, cs, ce, grid, origin, csz,
                                      h, 1000.0 * dp ** 3, c0=20.0)
    assert torch.allclose(acc[:, 2], torch.full((2,), -G, device="cuda"), atol=1e-4)
    assert float(acc[:, :2].abs().max()) < 1e-5
    assert float(drho.abs().max()) < 1e-6


@requires_cuda
def test_sph_block_is_finite_and_repels():
    """A compressed block pushes apart: pressure forces are finite and outward."""
    dp = 0.02
    n = 12
    g1 = torch.arange(n, device="cuda", dtype=torch.float32) * dp
    X, Y, Z = torch.meshgrid(g1, g1, g1, indexing="ij")
    pos = torch.stack([X.reshape(-1), Y.reshape(-1), Z.reshape(-1)], 1).contiguous()
    vel = torch.zeros_like(pos)
    rho = torch.full((pos.shape[0],), 1050.0, device="cuda")   # compressed
    ptype = torch.zeros(pos.shape[0], dtype=torch.int32, device="cuda")
    h = 1.5 * dp
    csz = 2 * h
    origin = (-0.1, -0.1, -0.1)
    grid = (int((n * dp + 0.3) / csz) + 2,) * 3
    perm, cs, ce = surf.build_cells(pos, origin, csz, grid)
    pos, vel, rho, ptype = pos[perm], vel[perm], rho[perm], ptype[perm]
    drho, acc, xsph = surf.sph_forces(pos, vel, rho, ptype, cs, ce, grid, origin, csz,
                                      h, 1000.0 * dp ** 3, c0=30.0)
    assert torch.isfinite(acc).all() and torch.isfinite(drho).all()
    ctr = pos.mean(0)
    outward = ((pos - ctr) * (acc - torch.tensor([0.0, 0.0, -G], device="cuda"))).sum(1)
    assert float(outward.mean()) > 0        # compressed block pushes outward