"""Unchanged deterministic-v2 reference rerun, exclusively into this new stage."""
import os
import sys
import time
import numpy as np
import pandas as pd
import yaml
from context import ROOT, EXP, CFG, check, save, bit_equal


def main():
    check()
    assert not (EXP / 'regression/frozen-five.json').exists()
    from run_case import preflight
    preflight()
    runtime_path = ROOT / 'models/male-cns-lif/deterministic-runtime-v2'
    pc = yaml.safe_load((runtime_path / 'phase4p.yaml').read_text())
    sys.path.insert(0, str(ROOT / 'models/male-cns-lif'))
    sys.path.insert(0, str(runtime_path))
    import torch
    import phase4p
    assert torch.cuda.mem_get_info()[0] > 2 * 1024**3
    torch.cuda.reset_peak_memory_stats()
    started = time.perf_counter()
    runtime = phase4p.CanonicalRuntime(pc, runtime_path, EXP)
    rows = []
    for spec in pc['reference_trials']:
        preflight()
        reference_dir = ROOT / 'experiments/phase4p-deterministic-runtime/phase4c-v2/reference' / spec['name']
        events = pd.read_parquet(reference_dir / 'spike-events.parquet')
        reference = dict(step=events.step.to_numpy(), neuron_idx=events.neuron_idx.to_numpy(),
            counts=np.load(reference_dir / 'per-neuron-spike-counts.int32.npy'),
            v=np.load(reference_dir / 'final-voltage.float32.npy'), g=np.load(reference_dir / 'final-g.float32.npy'))
        result = runtime.run(np.flatnonzero(runtime.pops[spec['population']]), spec['rate_hz'], spec['seed'], spec['duration_ms'])
        row = dict(trial=spec['name'], exact={key: bool(bit_equal(value, result[key])) for key, value in reference.items()},
            dtype_shape={key: dict(reference_dtype=str(value.dtype), actual_dtype=str(result[key].dtype),
                reference_shape=list(value.shape), actual_shape=list(result[key].shape)) for key, value in reference.items()},
            wall_seconds=result['metrics']['wall_seconds'])
        rows.append(row)
        print(row, flush=True)
    passed = len(rows) == 5 and all(all(row['exact'].values()) for row in rows)
    save(EXP / 'regression/frozen-five.json', dict(pass_all=passed, trials=rows,
        wall_seconds=time.perf_counter() - started, peak_allocated_vram=int(torch.cuda.max_memory_allocated()),
        scope='Old frozen deterministic-v2 five references only; not the new D18 hypothesis'))
    assert passed


if __name__ == '__main__':
    main()