"""Static slot targets and checkpoint compatibility, using the actual CUDA helper."""
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


@pytest.mark.parametrize('width', [0, 16, 48])
@pytest.mark.parametrize('amplitude', [0., .002, -.002])
def test_actual_slot_targets_and_requested_volume(tmp_path, width, amplitude):
    cfg = LBMConfig(nx=48, velocity=0., y_wall_bc='on_site_velocity',
                    wall_slot_width=width, wall_normal_velocity=amplitude)
    macros = cfg.to_macros()
    values = np.array(compile_run(tmp_path, 'slot_profile.cpp', {
        k: macros[k] for k in ('NX', 'WALL_SLOT_WIDTH', 'WALL_NORMAL_VELOCITY')}), dtype=float)
    # Independent node quadrature; includes both zero flanks and the full-width exception.
    x = np.arange(48)-23.5
    expected = amplitude*np.clip(1-(x/8)**2, 0, None)**2 if width == 16 else np.full(48, amplitude)
    np.testing.assert_allclose(values, expected, atol=5e-10, rtol=0)
    np.testing.assert_array_equal(values, values[::-1])
    assert abs(values.sum()-48*cfg.top_wall_normal_velocity) <= 1e-8
    assert abs(float(macros['TOP_WALL_NORMAL_VELOCITY'][:-1])-expected.mean()) < 1e-16


@pytest.mark.parametrize('change', [dict(wall_slot_width=w) for w in (-1, 1, 49, 2.5, True)]
                         + [dict(y_wall_bc='bounce_back'), dict(velocity=.01)])
def test_slot_invalid_configuration(change):
    values = dict(nx=48, velocity=0., y_wall_bc='on_site_velocity', wall_slot_width=16)
    values.update(change)
    with pytest.raises(ValueError, match='wall_slot_width'):
        LBMConfig(**values).validate_runtime()


def test_slot_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', wall_slot_width=16,
                          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.wall_slot_width == 16
    before = cfg.allocation_signature()
    cfg.wall_slot_width = 18
    assert before != cfg.allocation_signature()


@pytest.mark.parametrize('width, saved_width, accepted', [(6, 4, False), (6, None, False),
                                                        (0, None, True), (6, 6, True)])
def test_slot_checkpoint_width_before_device_writes(device, tmp_path, width, saved_width, accepted):
    field = make_field(device, velocity=0., y_wall_bc='on_site_velocity', wall_slot_width=width)
    bodies = manager(field)
    bodies.sync_to_gpu(field)
    stepper = SimpleNamespace(_step_count=0)
    path = save_checkpoint(field, SimpleNamespace(step_count=3), field.cfg, bodies,
                           str(tmp_path/'slot.h5'))
    with h5py.File(path, 'a') as hf:
        assert hf.attrs['ckpt_version'] == 3
        config = json.loads(hf.attrs['config_json'])
        config.pop('wall_slot_width')
        if saved_width is not None:
            config['wall_slot_width'] = saved_width
        hf.attrs['config_json'] = json.dumps(config)
    before, writes = field.ddf.copy(), device.writes
    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='wall_slot_width'):
            load_checkpoint(path, field, stepper, field.cfg, bodies)
        assert device.writes == writes and stepper._step_count == 0