"""Actual reservoir helper moments and static boundary restart contracts."""
import json
from types import SimpleNamespace

import h5py
import numpy as np
import pytest

from CelerisLab.config import LBMConfig, load_lbm_config
from CelerisLab.common.checkpoint import save_checkpoint, load_checkpoint
from .test_numerical_operators import compile_run, ROOT
from .test_state_regressions import device, make_field, manager


def test_reservoir_moments_and_known_populations(tmp_path):
    lines = compile_run(tmp_path, 'reservoir.cpp', {})
    cx = np.array([0,1,-1,0,0,1,-1,1,-1])
    cy = np.array([0,0,0,1,-1,1,-1,-1,1])
    weights = np.array([4/9]+[1/9]*4+[1/36]*4)
    for start in range(0, len(lines), 3):
        rho, v, neq = np.fromstring(lines[start], sep=' ')
        before, after = [np.fromstring(s, sep=' ') for s in lines[start+1:start+3]]
        known = cy >= 0
        np.testing.assert_array_equal(after[known], before[known])
        # The normal momentum reference uses ONLY prescribed rho and original known values.
        jy = before[cy == 0].sum()+2*before[cy == 1].sum()-rho
        np.testing.assert_allclose([after.sum(), after@cx, after@cy], [rho, 0., jy], atol=3e-7, rtol=0)
        if not neq:
            cu = cy*v
            expected = weights*rho*(1+3*cu+4.5*cu**2-1.5*v*v)
            np.testing.assert_allclose(after, expected, atol=3e-7, rtol=0)
            assert abs(after@cy/rho-v) < 3e-7


@pytest.mark.parametrize('override', [dict(top_wall_bc='unknown'), dict(y_wall_bc='bounce_back'),
                                    dict(velocity=.01), dict(rho=0.)])
def test_reservoir_invalid_configuration(override):
    cfg = dict(y_wall_bc='on_site_velocity', top_wall_bc='reservoir', velocity=0.)
    cfg.update(override)
    with pytest.raises(ValueError):
        LBMConfig(**cfg).validate_runtime()


def test_reservoir_json_and_static_signature(tmp_path):
    data = json.loads((ROOT/'src/CelerisLab/configs/config_lbm.json').read_text())
    data['physics']['velocity'] = 0.
    data['method'].update(y_wall_bc='on_site_velocity', top_wall_bc='reservoir',
                          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(path)
    assert cfg.to_macros()['TOP_WALL_BC'] == 1
    before = cfg.allocation_signature()
    cfg.rho *= 1.01
    assert before != cfg.allocation_signature()
    cfg.rho /= 1.01
    cfg.top_wall_bc = 'velocity'
    assert before != cfg.allocation_signature()


@pytest.mark.parametrize('mode, damage, accepted', [
    ('reservoir', 'mode', False), ('reservoir', 'missing', False),
    ('reservoir', 'rho', False), ('reservoir', 'u', False),
    ('reservoir', None, True), ('velocity', 'missing', True), ('velocity', 'mode', False),
])
def test_reservoir_checkpoint_contract(device, tmp_path, mode, damage, accepted):
    field = make_field(device, velocity=0., y_wall_bc='on_site_velocity', top_wall_bc=mode)
    bodies = manager(field)
    bodies.sync_to_gpu(field)
    path = save_checkpoint(field, SimpleNamespace(step_count=3), field.cfg, bodies,
                           str(tmp_path/'reservoir.h5'))
    with h5py.File(path, 'a') as hf:
        assert hf.attrs['ckpt_version'] == 3
        cfg = json.loads(hf.attrs['config_json'])
        if damage == 'missing': cfg.pop('top_wall_bc')
        elif damage == 'mode': cfg['top_wall_bc'] = 'velocity' if mode == 'reservoir' else 'reservoir'
        elif damage == 'rho': cfg['rho'] = 1.01
        elif damage == 'u': cfg['velocity'] = .01
        hf.attrs['config_json'] = json.dumps(cfg)
    writes, before = device.writes, field.ddf.copy()
    stepper = SimpleNamespace(_step_count=0)
    if accepted:
        load_checkpoint(path, field, stepper, field.cfg, bodies)
        np.testing.assert_array_equal(field.ddf, before)
        assert stepper._step_count == 3
    else:
        with pytest.raises(ValueError, match='boundary mismatch'):
            load_checkpoint(path, field, stepper, field.cfg, bodies)
        assert device.writes == writes and stepper._step_count == 0