"""Rerun old deterministic-v2 references without changing their semantics."""
import json
import os
import sys
import time
import urllib.request
import numpy as np
import pandas as pd
import yaml
from context import ROOT, EXP, check, save, bit_equal


def preflight():
    os.environ['CUDA_VISIBLE_DEVICES'] = 'GPU-79f3a85d-adce-7833-7003-6488140ca596'
    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'
    for port in (8188, 8189, 8190):
        with urllib.request.urlopen(f'http://192.168.10.51:{port}/queue', timeout=5) as response:
            queue = json.load(response)
        assert not queue['queue_running'] and not queue['queue_pending'], 'Comfy busy: defer, do not evict'


def main():
    check(); preflight()
    assert not (EXP / 'regression/frozen-five.json').exists()
    runtime_path = ROOT / 'models/male-cns-lif/deterministic-runtime-v2'
    config = 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(config, runtime_path, EXP)
    rows = []
    for spec in config['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={k: bit_equal(v, result[k]) for k, v in reference.items()},
            dtype_shape={k: dict(reference_dtype=str(v.dtype), actual_dtype=str(result[k].dtype),
                reference_shape=list(v.shape), actual_shape=list(result[k].shape)) for k, v 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 reference only, not D18 equality to old dynamics'))
    assert passed


if __name__ == '__main__':
    main()