"""J1 actual-header moments and bounded host capability contracts."""
import json
from pathlib import Path
from types import SimpleNamespace

import h5py
import numpy as np
import pytest

import CelerisLab
from CelerisLab.config import LBMConfig, load_lbm_config
from CelerisLab.cuda import compiler_v2 as compiler
from CelerisLab.common.checkpoint import save_checkpoint, load_checkpoint
from .test_numerical_operators import compile_run
from .test_state_regressions import device, make_field, manager


def test_planar_helper_moments_and_known_populations(tmp_path):
    root = Path(__file__).resolve().parents[2]
    assert Path(CelerisLab.__file__).resolve() == root/'src/CelerisLab/__init__.py'
    assert Path(compiler.kernel_path('kernel_v2.cu')).resolve() == root/'src/CelerisLab/lbm/kernels/kernel_v2.cu'
    lines = compile_run(tmp_path, 'planar_velocity.cpp', dict(WALL_NORMAL_VELOCITY='0.0f'))
    cx = np.array([0,1,-1,0,0,1,-1,1,-1])
    cy = np.array([0,0,0,1,-1,1,-1,-1,1])
    w = np.array([4/9]+[1/9]*4+[1/36]*4)
    for start in range(0, len(lines), 3):
        top, rho, u, v, neq = np.fromstring(lines[start], sep=' ')
        before, after = [np.fromstring(s, sep=' ') for s in lines[start+1:start+3]]
        known = cy != (-1 if top else 1)
        np.testing.assert_array_equal(after[known], before[known])
        density = after.sum()
        np.testing.assert_allclose([after@cx, after@cy], density*np.array([u,v]), atol=3e-7, rtol=0)
        if not neq:
            cu = cx*u+cy*v
            expected = w*rho*(1+3*cu+4.5*cu**2-1.5*(u*u+v*v))
            np.testing.assert_allclose(after, expected, atol=3e-7, rtol=0)
            assert abs(density-rho) <= 3e-7


@pytest.mark.parametrize('override', [
    dict(collision='TRT'), dict(collision='MRT'), dict(streaming='esopull'),
    dict(store_precision='FP16S'), dict(ddf_shifting=True), dict(les_enabled=True),
    dict(lattice_model='D3Q19'), dict(wall_normal_velocity=1.),
    dict(wall_normal_velocity=float('nan')),
])
def test_planar_unsupported_configuration(override):
    with pytest.raises(ValueError):
        LBMConfig(y_wall_bc='on_site_velocity', **override).validate_runtime()


def test_planar_config_roundtrip_and_fixed_targets(tmp_path):
    path = Path(__file__).resolve().parents[2]/'src/CelerisLab/configs/config_lbm.json'
    data = json.loads(path.read_text())
    data['method'].update(y_wall_bc='on_site_velocity', wall_normal_velocity=-.01,
                          collision='SRT', streaming='double_buffer',
                          store_precision='FP32', ddf_shifting=False)
    data['method']['les']['enabled'] = False
    path = tmp_path/'config.json'
    path.write_text(json.dumps(data))
    cfg = load_lbm_config(str(path))
    assert float(cfg.to_macros()['WALL_NORMAL_VELOCITY'][:-1]) == -.01
    with pytest.raises(ValueError, match='object-free'):
        compiler.generate_config(cfg, n_objects=1, build_root=tmp_path)
    before = cfg.allocation_signature()
    cfg.wall_normal_velocity = .01
    assert cfg.allocation_signature() != before


def test_planar_force_rejected_before_state_change(device):
    field = make_field(device, y_wall_bc='on_site_velocity')
    before, writes = field._runtime_params.copy(), device.writes
    with pytest.raises(ValueError, match='body forces'):
        field.update_params(fx=.001)
    assert field._runtime_params == before and device.writes == writes


@pytest.mark.parametrize('damage', ['target', 'mode', 'force'])
def test_planar_checkpoint_contract_before_writes(device, tmp_path, damage):
    field = make_field(device, y_wall_bc='on_site_velocity')
    bodies = manager(field)
    bodies.sync_to_gpu(field)
    path = save_checkpoint(field, SimpleNamespace(step_count=3), field.cfg, bodies,
                           str(tmp_path/'plane.h5'))
    with h5py.File(path, 'a') as hf:
        if damage == 'force':
            hf.attrs['runtime_params_json'] = '{"fx":0.001}'
        else:
            cfg = json.loads(hf.attrs['config_json'])
            cfg['wall_normal_velocity' if damage == 'target' else 'y_wall_bc'] = (
                .01 if damage == 'target' else 'bounce_back')
            hf.attrs['config_json'] = json.dumps(cfg)
    writes = device.writes
    with pytest.raises(ValueError):
        load_checkpoint(path, field, SimpleNamespace(_step_count=0), field.cfg, bodies)
    assert device.writes == writes