"""Continuing a run from a .cymsim file reproduces the uninterrupted run exactly."""
import copy

import h5py
import numpy as np
import pytest

from cymsim.cli import main
from cymsim.cymsim_file import CymsimWriter, load_cymsim
from cymsim.simulate import Simulator


def _config(base, duration):
    cfg = copy.deepcopy(base)
    cfg.numerics.duration = duration
    cfg.output.record_acceleration = True
    cfg.output.displacement_rate = 200.0
    cfg.output.probes = [(0.1, 0.05), (0.05, -0.12)]
    return cfg


def _datasets(path) -> dict:
    """Every dataset of a .cymsim file except the energy log, whose report times depend on how
    the run was split."""
    out = {}
    with h5py.File(path) as f:
        f.visititems(lambda name, obj: out.__setitem__(name, obj[()]) if isinstance(obj, h5py.Dataset) else None)
    out.pop("energy/table")
    return out


def _assert_same_files(a, b):
    da, db = _datasets(a), _datasets(b)
    assert sorted(da) == sorted(db)
    for name in da:
        assert np.array_equal(da[name], db[name]), name
    with h5py.File(a) as fa, h5py.File(b) as fb:
        for key in ("n_steps_done", "complete", "duration", "n_samples"):
            assert fa["run"].attrs[key] == fb["run"].attrs[key], key


@pytest.fixture(scope="module")
def straight(coarse_config, tmp_path_factory):
    path = tmp_path_factory.mktemp("straight") / "straight.cymsim"
    Simulator(_config(coarse_config, 0.05)).run(save_path=str(path))
    return path


def test_continued_run_equals_the_uninterrupted_one(coarse_config, straight, tmp_path):
    prev = None
    for k, duration in enumerate((0.02, 0.035, 0.05)):
        path = tmp_path / f"seg{k}.cymsim"
        Simulator(_config(coarse_config, duration)).run(save_path=str(path), resume=prev)
        prev = load_cymsim(path)
        assert prev.complete and len(prev.pressure_rayleigh) == round(44100 * duration)
        assert len(prev.state.overhang_front) > 0                    # the overshoot past the duration is kept
    _assert_same_files(path, straight)
    assert len(prev.energy_log) > len(load_cymsim(straight).energy_log)


def test_continue_from_a_killed_run(coarse_config, straight, tmp_path, monkeypatch):
    """A run killed right after a partial save continues from the partial file."""
    path = tmp_path / "killed.cymsim"
    update, calls = CymsimWriter.update, []

    def update_then_die(self, result):
        update(self, result)
        calls.append(1)
        if len(calls) == 3:
            self.close()
            raise KeyboardInterrupt

    monkeypatch.setattr(CymsimWriter, "update", update_then_die)
    with pytest.raises(KeyboardInterrupt):
        Simulator(_config(coarse_config, 0.05)).run(save_path=str(path), save_every=0)
    monkeypatch.undo()
    part = load_cymsim(path)
    assert not part.complete and 0 < part.n_steps_done
    out = tmp_path / "continued.cymsim"
    Simulator(_config(coarse_config, 0.05)).run(save_path=str(out), resume=part)
    _assert_same_files(out, straight)


def test_shorter_duration_cuts_the_records(coarse_config, straight, tmp_path):
    short = Simulator(_config(coarse_config, 0.03)).run()
    cut = Simulator(_config(coarse_config, 0.03)).run(resume=load_cymsim(straight))
    n = round(44100 * 0.03)
    assert len(cut.pressure_rayleigh) == n and cut.complete
    assert np.array_equal(cut.pressure_rayleigh.astype(np.float32), short.pressure_rayleigh.astype(np.float32))
    assert np.array_equal(cut.acceleration.data, short.acceleration.data)
    assert np.array_equal(cut.displacement.w, short.displacement.w)
    assert np.array_equal(cut.probe_acceleration.astype(np.float32), short.probe_acceleration.astype(np.float32))


def test_a_different_configuration_is_refused(coarse_config, straight):
    cfg = _config(coarse_config, 0.08)
    cfg.hits[0].velocity *= 1.001
    with pytest.raises(ValueError, match=r"hits"):
        Simulator(cfg).run(resume=load_cymsim(straight))
    cfg = _config(coarse_config, 0.08)
    cfg.output.prefix = "elsewhere"                                  # output naming may change
    Simulator(cfg).check_resumable(load_cymsim(straight))


def test_cli_resume(coarse_config, straight, tmp_path, capsys):
    cfg = _config(coarse_config, 0.05)
    cfg.output.wav = cfg.output.npz = cfg.output.cascade_plot = False
    cfg.output.prefix = str(tmp_path / "x")
    ypath = tmp_path / "c.yaml"
    cfg.to_yaml(ypath)
    out = tmp_path / "longer.cymsim"
    assert main(["run", str(ypath), "--resume", str(straight), "--duration", "0.06", "-o", str(out), "--quiet"]) == 0
    assert load_cymsim(out).complete and len(load_cymsim(out).pressure_rayleigh) == round(44100 * 0.06)
    cfg.hits[0].r += 0.01
    cfg.to_yaml(ypath)
    assert main(["run", str(ypath), "--resume", str(straight), "-o", str(out), "--quiet"]) == 2
    assert "hits" in capsys.readouterr().err