"""Fixed surface actuation: discrete operator/ledger checks, not flux accuracy."""
from dataclasses import asdict
import json

import h5py
import numpy as np
import pycuda.driver as cuda
import pytest

from CelerisLab.lbm.descriptors import FLUID, SOLID
from ._cases import make_sim
from .test_numerical_gpu import CX, CY, install_spatial_state, spatial_state


WEIGHTS = np.array([4/9] + [1/9]*4 + [1/36]*4)
OPPOSITE = np.array([0, 2, 1, 4, 3, 6, 5, 8, 7])
GEOMETRY = ("fluid_idx", "dir", "q", "body_id", "rx", "ry", "rz", "fallback_class")


def feq(rho, ux, uy):
    """Independent FP64 polynomial; round only the uploaded populations."""
    cu = CX[:, None, None]*ux + CY[:, None, None]*uy
    return WEIGHTS[:, None, None]*rho*(1 + 3*cu + 4.5*cu**2 - 1.5*(ux**2 + uy**2))


@pytest.mark.parametrize("speed", [-.01, .01])
def test_actual_surface_kernel_equilibrium_branches(speed):
    """Both q branches, q=.5 and donor fallback; non-unit rho and signed normal."""
    sim = make_sim(nx=24, ny=16, velocity=0)
    allocations = []
    try:
        sim.initialize()
        normal = np.array([-.6, .8])
        expected = feq(1.17, *(speed*normal))[:, 0, 0]
        initial = np.broadcast_to(expected[:, None, None], (9, 16, 24)).astype(np.float32).copy()
        # Distinct source/destination neighborhoods; no synthetic link race.
        cells = np.array([4*24+12, 6*24+12, 8*24+12, 10*24+12], np.uint32)
        directions = np.ones(4, np.uint8)
        radii = np.tile(5*normal, (4, 1)).astype(np.float32)
        columns = [cells, directions, np.array([.2, .5, .8, .2], np.float32),
                   radii[:, 0].copy(), radii[:, 1].copy(), np.zeros(4, np.float32),
                   np.array([0, 0, 0, 1], np.uint8), np.zeros(4, np.int32)]
        with sim.ctx.activate():
            for array in columns + [np.tile(speed*normal, (4, 1)).astype(np.float32),
                                    np.zeros((4, 3), np.float32)]:
                allocations.append(cuda.to_device(array))
            cuda.memcpy_htod(sim.field.ddf_gpu, initial)
            sim.stepper.curved_fn(sim.field.ddf_gpu, *allocations[:8],
                                  np.uintp(0), np.uintp(0), np.uint32(4),
                                  *allocations[8:], block=(32, 1, 1), grid=(1, 1))
            actual = np.empty_like(initial)
            exchange = np.empty((4, 3), np.float32)
            cuda.memcpy_dtoh(actual, sim.field.ddf_gpu)
            cuda.memcpy_dtoh(exchange, allocations[-1])
        reflected = actual.reshape(9, -1)[2, cells+1]
        np.testing.assert_allclose(reflected, expected[2], atol=7e-8, rtol=0)
        # Equilibrium identity is independent of interpolation implementation.
        fx = expected[1] + expected[2]
        wanted = np.column_stack((np.full(4, fx), np.zeros(4), -radii[:, 1]*fx))
        np.testing.assert_allclose(exchange, wanted, atol=4e-7, rtol=0)
    finally:
        with sim.ctx.activate():
            for allocation in allocations:
                allocation.free()
        sim.close()


def circle_sim(speed):
    sim = make_sim(nx=32, ny=24, velocity=0)
    spec = None if speed is None else dict(angle=.7, half_width=.85, normal_velocity=speed)
    try:
        bid = sim.add_body("circle", center=(16, 12), radius=4.3, surface_actuation=spec)
        sim.initialize()
        flags = sim.field.flag.copy()
        flags.reshape(24, 32)[1:-1, [0, -1]] = FLUID  # x periodic, y halfway walls
        y, x = np.mgrid[:24, :32]
        initial = feq(1.17+.001*np.cos(.2*x+.3*y),
                      .003*np.sin(.3*y), .002*np.cos(.4*x)).astype(np.float32)
        install_spatial_state(sim, initial)
        sim.field.flag[:] = flags  # installer intentionally resets all flags
        sim.field.upload_flags()
        return sim, bid, initial
    except BaseException:
        sim.close()
        raise


@pytest.mark.parametrize("speed", [-.01, .01])
def test_surface_population_ledger_and_public_diagnostics(speed):
    sim, bid, initial = circle_sim(speed)
    try:
        cl = sim.field.curved
        frozen = {name: getattr(cl, name).copy() for name in GEOMETRY}
        flags = sim.field.flag.copy()
        body_state = vars(sim.bodies.get(bid).state).copy()
        physical = initial.astype(np.float64)
        flat = physical.reshape(9, -1)
        cell, direction = cl.fluid_idx.astype(int), cl.dir.astype(int)
        ex, ey, q = CX[direction], CY[direction], cl.q.astype(float)
        # Independently project the prescribed arc at true intersections.
        angle = np.arctan2(cl.ry.astype(float), cl.rx.astype(float))
        delta = (angle-.7+np.pi) % (2*np.pi)-np.pi
        aperture = np.abs(delta) < .85
        radial = np.column_stack((cl.rx, cl.ry)).astype(float)
        normal = radial / np.linalg.norm(radial, axis=1)[:, None]
        velocity = speed*np.maximum(0, 1-(delta/.85)**2)[:, None]**2*normal
        np.testing.assert_array_equal(cl.aperture, aperture)
        np.testing.assert_allclose(cl.surface_velocity, velocity, atol=1e-8, rtol=0)
        assert aperture.any() and (~aperture).any()
        assert (q < .5).any() and (q >= .5).any()
        a = flat[direction, cell]
        b = np.empty(len(cell))
        passive = np.empty(len(cell))
        for i, (k, d) in enumerate(zip(cell, direction)):
            rho = flat[:, k].sum()
            dot = CX[d]*velocity[i, 0] + CY[d]*velocity[i, 1]
            if cl.fallback_class[i]:
                passive[i] = a[i]
                factor = 6
            elif q[i] < .5:
                passive[i] = 2*q[i]*a[i] + (1-2*q[i])*flat[d, k-CX[d]-32*CY[d]]
                factor = 6
            else:
                passive[i] = a[i]/(2*q[i]) + (1-1/(2*q[i]))*flat[OPPOSITE[d], k]
                factor = 3/q[i]
            b[i] = passive[i]-factor*WEIGHTS[d]*rho*dot
        force = np.column_stack((ex, ey))*(a+b)[:, None]
        torque = cl.rx*force[:, 1]-cl.ry*force[:, 0]
        wanted = np.column_stack((force, torque))
        # Fluid-node and intersection lever arms differ by q*c: same link torque.
        fluid_r = radial-q[:, None]*np.column_stack((ex, ey))
        np.testing.assert_allclose(torque, fluid_r[:, 0]*force[:, 1]-fluid_r[:, 1]*force[:, 0],
                                   atol=2e-15, rtol=0)
        with sim.ctx.activate():
            sim.stepper.curved_fn(sim.field.ddf_gpu,
                *(getattr(cl, name+"_gpu") for name in GEOMETRY[:3]),
                cl.rx_gpu, cl.ry_gpu, cl.rz_gpu, cl.fallback_class_gpu, cl.body_id_gpu,
                np.uintp(0), np.uintp(0), np.uint32(cl.count),
                cl.surface_velocity_gpu, cl.link_exchange_gpu,
                block=(64, 1, 1), grid=((cl.count+63)//64, 1))
            staged = np.empty_like(initial)
            cuda.memcpy_dtoh(staged, sim.field.ddf_gpu)
        incoming = staged.reshape(9, -1)[OPPOSITE[direction], cell+ex+32*ey].astype(float)
        np.testing.assert_allclose(incoming, b, atol=8e-8, rtol=0)
        # Incoming populations alone respond with the expected signed normal drive.
        assert np.all((-np.sum(normal[aperture]*np.column_stack((ex, ey))[aperture], axis=1)) > 0)
        core = aperture & (np.abs(delta) < .5*.85)
        assert core.any()
        assert np.all(speed*(incoming-passive)[core] > 0)

        # Independent pull ledger, including all planar-wall links. Use actual
        # rounded incoming values to isolate bookkeeping from CUDA arithmetic.
        arrived = np.stack([np.roll(physical[d], (CY[d], CX[d]), (0, 1)) for d in range(9)])
        arrived.reshape(9, -1)[OPPOSITE[direction], cell] = incoming
        mass_exchange = float(np.sum(incoming-a))
        momentum_exchange = -np.sum(np.column_stack((ex, ey))*(a+incoming)[:, None], axis=0)
        for y, outgoing in ((1, (4, 6, 7)), (22, (3, 5, 8))):
            for d in outgoing:
                values = physical[d, y, :]
                arrived[OPPOSITE[d], y, :] = values
                momentum_exchange -= 2*values.sum()*np.array([CX[d], CY[d]])
        fluid = ((flags & FLUID) != 0) & ((flags & SOLID) == 0)
        change = (arrived-physical).reshape(9, -1)[:, fluid].sum(axis=1)
        np.testing.assert_allclose([change.sum(), change@CX, change@CY],
                                   [mass_exchange, *momentum_exchange], atol=3e-13, rtol=0)
        # Start public run from the original state, not a previously staged field.
        with sim.ctx.activate():
            cuda.memcpy_htod(sim.field.ddf_gpu, initial)
        sim.run(1)
        actual = spatial_state(sim).astype(float).reshape(9, -1)[:, fluid]
        expected = arrived.reshape(9, -1)[:, fluid]
        moments = np.stack((np.ones(9), CX, CY))
        np.testing.assert_allclose(moments@actual, moments@expected, atol=5e-7, rtol=0)
        # FP32 collision conserves moments up to rounding, not bitwise globally.
        change = (actual-physical.reshape(9, -1)[:, fluid]).sum(axis=1)
        budget = 8*np.finfo(np.float32).eps*np.abs(expected).sum()
        np.testing.assert_allclose([change.sum(), change@CX, change@CY],
                                   [mass_exchange, *momentum_exchange], atol=budget, rtol=0)
        diagnostic = sim.read_surface_exchange(bid)
        for name, mask in (("wall", ~aperture), ("aperture", aperture)):
            np.testing.assert_allclose(diagnostic[name], wanted[mask].sum(axis=0),
                                       atol=4e-6, rtol=0)
        total = diagnostic["wall"]+diagnostic["aperture"]
        np.testing.assert_allclose(sim.read_force(bid, normalize=False), total[:2], atol=4e-6, rtol=0)
        np.testing.assert_allclose(sim.read_torque(bid, normalize=False), total[2:], atol=4e-6, rtol=0)
        for name, value in frozen.items():
            np.testing.assert_array_equal(getattr(cl, name), value)
        np.testing.assert_array_equal(sim.field.flag, flags)
        assert vars(sim.bodies.get(bid).state) == body_state
    finally:
        sim.close()


def test_zero_surface_drive_preserves_legacy_populations_and_geometry():
    records = []
    for speed in (None, 0.0):
        sim, bid, _ = circle_sim(speed)
        try:
            sim.run(2)
            records.append((spatial_state(sim), sim.field.flag.copy(),
                            {name: getattr(sim.field.curved, name).copy() for name in GEOMETRY}))
        finally:
            sim.close()
    np.testing.assert_array_equal(records[0][0], records[1][0])
    np.testing.assert_array_equal(records[0][1], records[1][1])
    for name in GEOMETRY:
        np.testing.assert_array_equal(records[0][2][name], records[1][2][name])


def test_surface_checkpoint_continues_and_resets_latest_exchange(tmp_path):
    sim, bid, _ = circle_sim(.01)
    try:
        sim.run(1)
        saved = spatial_state(sim)
        count = sim.stepper.step_count
        spec = asdict(sim.bodies.get(bid).surface_actuation)
        action = sim.bodies.action.copy()
        path = sim.save_checkpoint(str(tmp_path / "surface.h5"))
        with h5py.File(path, "r") as checkpoint:
            assert checkpoint.attrs["ckpt_version"] == 4
            geometry = json.loads(checkpoint.attrs["geometry_json"])
            assert len(geometry[0]["projection_sha256"]) == 64
            assert json.loads(checkpoint.attrs["surface_actuation_json"]) == [spec]
            applied = checkpoint["applied_action"][:]
        sim.run(2)
        expected = spatial_state(sim)
        expected_exchange = sim.read_surface_exchange(bid)
        assert np.linalg.norm(expected_exchange["aperture"]) > 0
        sim.load_checkpoint(path)
        assert sim.stepper.step_count == count
        np.testing.assert_array_equal(spatial_state(sim), saved)
        np.testing.assert_array_equal(sim.bodies.action, action)
        actual_action = np.empty_like(applied)
        with sim.ctx.activate():
            cuda.memcpy_dtoh(actual_action, sim.bodies.action_gpu)
        np.testing.assert_array_equal(actual_action, applied)
        assert asdict(sim.bodies.get(bid).surface_actuation) == spec
        for values in sim.read_surface_exchange(bid).values():
            np.testing.assert_array_equal(values, np.zeros(3))
        restored = sim.save_checkpoint(str(tmp_path / "restored.h5"))
        with h5py.File(restored, "r") as checkpoint:
            assert json.loads(checkpoint.attrs["geometry_json"]) == geometry
            assert json.loads(checkpoint.attrs["surface_actuation_json"]) == [spec]
        sim.run(2)
        np.testing.assert_array_equal(spatial_state(sim), expected)
        for name, values in sim.read_surface_exchange(bid).items():
            np.testing.assert_array_equal(values, expected_exchange[name])
    finally:
        sim.close()