"""Run unchanged frozen reference into new output; never rewrite old evidence."""
import os
import sys
import time
import numpy as np
import pandas as pd
import yaml
from context import ROOT, EXP, CFG, save, check, bit_equal


def main():
    check()
    assert not (EXP / 'regression/frozen-five.json').exists()
    rtpath = ROOT / 'models/male-cns-lif/deterministic-runtime-v2'
    pc = yaml.safe_load((rtpath / 'phase4p.yaml').read_text())
    os.environ['CUDA_VISIBLE_DEVICES'] = CFG['gpu_uuid']
    os.environ['CUDA_HOME'] = '/usr/local/cuda'
    os.environ['PATH'] = str(ROOT / 'models/male-cns-lif/.venv/bin') + ':/usr/local/cuda/bin:' + os.environ['PATH']
    os.environ['MAX_JOBS'] = '2'
    os.environ['TORCH_CUDA_ARCH_LIST'] = '8.9'
    sys.path.insert(0, str(ROOT / 'models/male-cns-lif/terminal-scope-diagnostic-v1'))
    from diagnostic import queues
    queues()
    sys.path.insert(0, str(ROOT / 'models/male-cns-lif'))
    sys.path.insert(0, str(rtpath))
    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, rtpath, EXP)
    rows = []
    for spec in pc['reference_trials']:
        queues()
        base = ROOT / 'experiments/phase4p-deterministic-runtime/phase4c-v2/reference' / spec['name']
        events = pd.read_parquet(base / 'spike-events.parquet')
        reference = dict(step=events.step.to_numpy(), neuron_idx=events.neuron_idx.to_numpy(),
            counts=np.load(base / 'per-neuron-spike-counts.int32.npy'),
            v=np.load(base / 'final-voltage.float32.npy'), g=np.load(base / '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 = 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()),
        oracle='Old frozen deterministic-v2 only; not this independent dynamics branch'))
    assert passed


if __name__ == '__main__':
    main()