"""Static circle apertures: real geometry and host/device-memory protocol tests."""
from dataclasses import FrozenInstanceError, asdict
import json
import math
from types import SimpleNamespace

import h5py
import numpy as np
import pytest

from CelerisLab.body.geometry.circle import CircleGeometry
from CelerisLab.body.manager import ObjectManager
from CelerisLab.body.objects import SimObject, SurfaceActuation
from CelerisLab.common.checkpoint import load_checkpoint, save_checkpoint
from CelerisLab.config import LBMConfig
from CelerisLab.lbm.curved_links import CurvedLinkSoA
from CelerisLab.lbm.descriptors import FLUID, SOLID
from CelerisLab.simulation import Simulation
from .test_state_regressions import device, make_field, manager


def circle(spec=None, *, center=(6., 5.), radius=2.999, **kwargs):
    return SimObject(0, geometry=CircleGeometry(*center, radius), center=center,
                     radius=radius, surface_actuation=spec, **kwargs)


def projected(obj):
    cfg = LBMConfig(nx=12, ny=10, viscosity=.1)
    cfg.validate()
    flags = ObjectManager.build_flags_for(
        [obj], np.full(cfg.nx*cfg.ny, FLUID, np.uint16), nx=cfg.nx, ny=cfg.ny, nz=1)
    cl = CurvedLinkSoA()
    cl.assign_host(*obj.get_curved_list(cfg.nx, cfg.ny, cfg.nq, domain_flags=flags))
    cl.assign_surface([obj], cfg, flags)
    return cl, cfg, flags


@pytest.mark.parametrize('speed', [.002, -.002])
def test_aperture_uses_true_hit_not_fluid_node_or_clamped_q(speed):
    # Fluid (9,5), direction (-1,+1), circle radius 2.999. Actual t < .01,
    # whereas the existing Bouzidi q is clipped to .01. The narrow opening
    # contains the true hit but neither the fluid-node angle nor clipped hit.
    t = (6-math.sqrt(36-8*(9-2.999**2)))/4
    angle = math.atan2(t, 3-t)
    spec = SurfaceActuation(angle, .0001, speed)
    cl, _, _ = projected(circle(spec))
    index, = np.flatnonzero((cl.fluid_idx == 9+5*12) & (cl.dir == 8))
    assert cl.q[index] == np.float32(.01)
    assert angle > spec.half_width
    assert math.atan2(.01, 2.99)-angle > spec.half_width
    assert cl.aperture[index]
    expected = speed*np.array([3-t, t])/2.999
    np.testing.assert_allclose(cl.surface_velocity[index], expected, rtol=2e-6, atol=1e-10)
    np.testing.assert_allclose(np.hypot(cl.rx, cl.ry), 2.999, rtol=0, atol=3e-7)
    np.testing.assert_array_equal(cl.surface_velocity[~cl.aperture], 0.)


def test_surface_zero_preserves_eight_column_geometry_and_mask():
    plain, zero = circle(), circle(SurfaceActuation(0., .7, 0.))
    old, new = plain.get_curved_list(12, 10, 9), zero.get_curved_list(12, 10, 9)
    assert len(old) == len(new) == 8
    for before, after in zip(old, new):
        np.testing.assert_array_equal(before, after)
    np.testing.assert_array_equal(plain.get_flag_mask(12, 10), zero.get_flag_mask(12, 10))
    cl, _, _ = projected(zero)
    assert cl.aperture.any()
    np.testing.assert_array_equal(cl.surface_velocity, 0.)
    assert projected(plain)[0].surface_velocity is None


def test_surface_spec_is_immutable_and_wraps_angles():
    spec = SurfaceActuation(math.pi, .2, .002)
    with pytest.raises(FrozenInstanceError):
        spec.normal_velocity = .003
    angles = np.array([math.pi-.1, -math.pi+.1, 0.])
    velocity, mask = spec.project(np.cos(angles), np.sin(angles))
    np.testing.assert_array_equal(mask, [True, True, False])
    np.testing.assert_allclose(np.linalg.norm(velocity[:2], axis=1), .002*(1-.5**2)**2,
                               rtol=1e-6)
    np.testing.assert_array_equal(velocity[2], 0.)


@pytest.mark.parametrize('values', [(0., 0., .002), (0., math.pi+.1, .002),
                                   (0., .2, 1.), (math.nan, .2, .002),
                                   (0., math.inf, .002), (0., .2, math.nan)])
def test_invalid_surface_spec(values):
    with pytest.raises(ValueError):
        SurfaceActuation(*values)


@pytest.mark.parametrize('override', [dict(streaming='esopull'), dict(store_precision='FP16S'),
                                     dict(ddf_shifting=True), dict(lattice_model='D3Q19')])
def test_surface_rejects_unsupported_mode_before_registration(override):
    cfg = LBMConfig(nx=12, ny=10, **override)
    cfg.validate()
    bodies = ObjectManager(cfg.nx, cfg.ny, cfg.nz, cfg.nq, cfg)
    with pytest.raises(ValueError, match='Surface actuation requires'):
        bodies.add(circle(SurfaceActuation(0., .7, .002)))
    assert bodies.count == 0


@pytest.mark.parametrize('role', ['is_sensor', 'is_force_region'])
def test_surface_rejects_nonwall_roles(role):
    with pytest.raises(ValueError):
        circle(SurfaceActuation(0., .7, .002), **{role: True}).validate_surface_actuation(LBMConfig())


@pytest.mark.parametrize('key', ['omega', 'vx', 'vy'])
def test_surface_set_body_rejects_motion_before_mutation(key):
    cfg = LBMConfig(nx=12, ny=10)
    cfg.validate()
    bodies = ObjectManager(cfg.nx, cfg.ny, cfg.nz, cfg.nq, cfg)
    item = circle(SurfaceActuation(0., .7, .002))
    bodies.add(item)
    sim = Simulation.__new__(Simulation)
    sim.lbm_cfg, sim.bodies = cfg, bodies
    state, action = asdict(item.state), bodies.action.copy()
    with pytest.raises(ValueError, match='stationary'):
        sim.set_body(0, **{key: .01})
    assert asdict(item.state) == state
    np.testing.assert_array_equal(bodies.action, action)


def test_surface_projection_rejects_unresolved_aperture_and_bad_writers():
    with pytest.raises(ValueError, match='no resolved'):
        projected(circle(SurfaceActuation(.123456, 1e-9, .002)))
    item = circle(SurfaceActuation(0., .7, .002))
    cl, cfg, flags = projected(item)
    flags[cl.fluid_idx[0]] = SOLID
    with pytest.raises(ValueError, match='active fluid'):
        cl.assign_surface([item], cfg, flags)
    flags[cl.fluid_idx[0]] = FLUID
    cl.fluid_idx[1], cl.dir[1] = cl.fluid_idx[0], cl.dir[0]
    with pytest.raises(ValueError, match='unique'):
        cl.assign_surface([item], cfg, flags)


def test_surface_public_dict_entry_and_sync_copy():
    cfg = LBMConfig(nx=12, ny=10)
    cfg.validate()
    sim = Simulation.__new__(Simulation)
    sim.lbm_cfg, sim._initialized = cfg, False
    sim.bodies = ObjectManager(cfg.nx, cfg.ny, cfg.nz, cfg.nq, cfg)
    spec = dict(angle=0., half_width=.7, normal_velocity=.002)
    oid = sim.add_body('circle', center=(6., 5.), radius=2., surface_actuation=spec)
    copied, = sim.bodies.build_next_objects()
    assert asdict(copied.surface_actuation) == spec
    assert copied.surface_actuation is sim.bodies.get(oid).surface_actuation


def case(device, spec):
    field = make_field(device)
    bodies = manager(field, [circle(spec, radius=2.)])
    bodies.sync_to_gpu(field)
    return field, bodies


def test_surface_gpu_buffer_lifecycle_with_memory_double(device):
    field, bodies = case(device, SurfaceActuation(0., .7, .002))
    cl = field.curved
    velocity, exchange = cl.surface_velocity_gpu, cl.link_exchange_gpu
    np.testing.assert_array_equal(np.frombuffer(velocity.data, np.float32).reshape(-1, 2),
                                  cl.surface_velocity)
    np.testing.assert_array_equal(np.frombuffer(exchange.data, np.float32), 0.)
    cl.assign_surface([], field.cfg, field.flag)
    cl.upload()
    assert velocity.freed and exchange.freed
    assert cl.surface_velocity_gpu is None and cl.link_exchange_gpu is None
    cl.assign_surface(bodies.objects, field.cfg, field.flag)
    cl.upload()
    velocity, exchange = cl.surface_velocity_gpu, cl.link_exchange_gpu
    cl.free()
    assert velocity.freed and exchange.freed and cl.count == 0


@pytest.mark.parametrize('damage', ['missing_spec', 'legacy_missing_spec', 'different_spec',
                                   'omega', 'vx', 'vy', 'applied_omega'])
def test_surface_checkpoint_rejects_mismatch_before_any_write(device, tmp_path, damage):
    field, bodies = case(device, SurfaceActuation(0., .7, .002))
    path = save_checkpoint(field, SimpleNamespace(step_count=3), field.cfg, bodies,
                           str(tmp_path/'surface.h5'))
    with h5py.File(path, 'a') as hf:
        assert hf.attrs['ckpt_version'] == 4
        if damage in ('missing_spec', 'legacy_missing_spec'):
            del hf.attrs['surface_actuation_json']
            if damage == 'legacy_missing_spec':
                hf.attrs['ckpt_version'] = 3
        elif damage == 'different_spec':
            records = json.loads(hf.attrs['surface_actuation_json'])
            records[0]['normal_velocity'] *= -1
            hf.attrs['surface_actuation_json'] = json.dumps(records)
        elif damage == 'applied_omega':
            action = hf['applied_action'][:]
            action[-1] = .01
            hf['applied_action'][:] = action
        else:
            records = json.loads(hf.attrs['objects_json'])
            records[0][damage] = .01
            hf.attrs['objects_json'] = json.dumps(records)
    writes, state = device.writes, asdict(bodies.get(0).state)
    raw = [bytes(ptr.data) for ptr in (field.ddf_gpu, field.temp_gpu, field.flag_gpu)]
    stepper = SimpleNamespace(_step_count=99)
    with pytest.raises(ValueError):
        load_checkpoint(path, field, stepper, field.cfg, bodies)
    assert device.writes == writes and stepper._step_count == 99
    assert asdict(bodies.get(0).state) == state
    assert raw == [bytes(ptr.data) for ptr in (field.ddf_gpu, field.temp_gpu, field.flag_gpu)]


@pytest.mark.parametrize('enabled', [False, True])
def test_surface_checkpoint_rebuild_and_legacy_absence(device, tmp_path, enabled):
    spec = SurfaceActuation(0., .7, .002) if enabled else None
    field, bodies = case(device, spec)
    path = save_checkpoint(field, SimpleNamespace(step_count=3), field.cfg, bodies,
                           str(tmp_path/'surface.h5'))
    with h5py.File(path, 'a') as hf:
        assert hf.attrs['ckpt_version'] == (4 if enabled else 3)
        if not enabled:
            del hf.attrs['surface_actuation_json']
    restored, restored_bodies = case(device, spec)
    stepper = SimpleNamespace(_step_count=0)
    load_checkpoint(path, restored, stepper, restored.cfg, restored_bodies)
    assert stepper._step_count == 3
    for source, target in ((field.ddf_gpu, restored.ddf_gpu), (field.temp_gpu, restored.temp_gpu),
                           (field.flag_gpu, restored.flag_gpu)):
        assert bytes(source.data) == bytes(target.data)
    if enabled:
        np.testing.assert_array_equal(field.curved.surface_velocity, restored.curved.surface_velocity)
    else:
        assert restored.curved.surface_velocity is None