"""Run old reference unchanged, writing only new experiment output."""
import os
import sys
import time
import numpy as np
import pandas as pd
import yaml
from context import ROOT,EXP,CFG,save,check


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,phase4p
    assert torch.cuda.mem_get_info()[0]>2*1024**3
    torch.cuda.reset_peak_memory_stats();t0=time.perf_counter();rt=phase4p.CanonicalRuntime(pc,rtpath,EXP);rows=[]
    for spec in pc['reference_trials']:
        queues();base=ROOT/'experiments/phase4p-deterministic-runtime/phase4c-v2/reference'/spec['name']
        ev=pd.read_parquet(base/'spike-events.parquet')
        ref=dict(step=ev.step.to_numpy(),neuron_idx=ev.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=rt.run(np.flatnonzero(rt.pops[spec['population']]),spec['rate_hz'],spec['seed'],spec['duration_ms'])
        row=dict(trial=spec['name'],exact={k:bool(np.array_equal(v,result[k])) for k,v in ref.items()},wall_seconds=result['metrics']['wall_seconds'])
        rows.append(row);print(row,flush=True)
    save(EXP/'regression/frozen-five.json',dict(pass_all=all(all(r['exact'].values()) for r in rows),trials=rows,
        wall_seconds=time.perf_counter()-t0,peak_allocated_vram=int(torch.cuda.max_memory_allocated()),oracle='old frozen deterministic-v2 only'))
    assert all(all(r['exact'].values()) for r in rows)


if __name__=='__main__':main()