"""Fixed Brian/GPU stress suite, with source-checked immutable inputs."""
import argparse
import json
import os
from pathlib import Path
import sys
import time
import numpy as np
from context import ROOT,MODEL,EXP,CFG,save,check

def load(path):
    with np.load(path,allow_pickle=False) as f:return {k:f[k].copy() for k in f.files}

def compare(a,b,case,precision):
    checks={k:bool(np.array_equal(a[k],b[k])) for k in ('step','neuron_idx','counts','ref')}
    ah=np.asarray(a['history'],np.float64);bh=np.asarray(b['history'],np.float64)
    assert ah.shape==bh.shape
    delta=np.abs(ah[:,:2]-bh[:,:2]);limit=CFG['max_abs_volt'][precision]
    checks['history_ref_exact']=bool(np.array_equal(ah[:,2],bh[:,2]))
    checks['finite_states']=bool(np.isfinite(ah).all() and np.isfinite(bh).all())
    checks['state_tolerance']=bool(delta.max()<=limit)
    checks['weights_exact']=bool(np.array_equal(a['weights'],b['weights']))
    # Distinguish logical event requests from realised requests after refractory.
    for key in ('request_attempted','request_accepted','request_rejected'):
        checks[key+'_exact']=bool(key in a and key in b and np.array_equal(a[key],b[key]))
    data=load(case/'arrays.npz');targets=data['dst']
    def arrival_sequence(result):
        tick=result['arrival_step'];eid=result['arrival_edge_index'];guard=result['arrival_accepted']
        assert tick.shape==eid.shape==guard.shape
        # Different targets can interleave; order of writes to one target cannot.
        order=np.lexsort((np.arange(len(eid)),targets[eid],tick))
        return np.stack((tick[order],eid[order],guard[order]),axis=1)
    aa=arrival_sequence(a);bb=arrival_sequence(b)
    checks['arrival_order_and_guard_exact']=bool(np.array_equal(aa,bb))
    first=None
    for t in range(len(ah)):
        if not np.array_equal(a['neuron_idx'][a['step']==t],b['neuron_idx'][b['step']==t]):first=t;break
    bits_equal=a['history'].dtype==b['history'].dtype and np.array_equal(np.ascontiguousarray(a['history']).view(np.uint8),np.ascontiguousarray(b['history']).view(np.uint8))
    return dict(pass_all=all(checks.values()),checks=checks,bitwise_history=bool(bits_equal),
        max_abs_volt=float(delta.max()),mean_abs_volt=float(delta.mean()),fixed_limit_volt=limit,
        first_event_divergence_step=first,oracle_events=int(len(a['step'])),gpu_events=int(len(b['step'])),
        due_edges=int(len(aa)),accepted_edges=int(a['arrival_accepted'].sum()),
        attempted_requests=int(a['request_attempted'].sum()),accepted_requests=int(a['request_accepted'].sum()))

def oracle():
    from oracle_stress import run_case
    t0=time.perf_counter();rows=[]
    for case in sorted((EXP/'cases').iterdir()):
        for precision in CFG['precisions']:
            output=EXP/'oracle'/(case.name+'-'+precision)
            run_case(case,output,precision)
            r=load(output/'state.npz');assert np.isfinite(r['history']).all()
            rows.append(dict(case=case.name,precision=precision,events=len(r['step'])))
            print(json.dumps(rows[-1]),flush=True)
    save(EXP/'regression/oracle-complete.json',dict(cases=rows,wall_seconds=time.perf_counter()-t0))

def gpu():
    assert (EXP/'regression/oracle-complete.json').exists(),'No complete oracle suite'
    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()
    from gpu_backend import run_gpu
    t0=time.perf_counter();rows=[]
    for case in sorted((EXP/'cases').iterdir()):
        for precision in CFG['precisions']:
            queues();refdir=EXP/'oracle'/(case.name+'-'+precision)
            reference=load(refdir/'state.npz');coeff=json.loads((refdir/'coefficients.json').read_text());first=None
            for repeat in range(CFG['gpu_repeats']):
                out=EXP/'gpu'/(case.name+'-'+precision+f'-r{repeat}')
                run_gpu(case,out,precision,coeff,EXP/'extension-build',weights=reference['weights'])
                result=load(out/'state.npz');cmp=compare(reference,result,case,precision)
                if first is None:first=result
                rep={k:bool(np.array_equal(v,result[k])) for k,v in first.items()}
                row=dict(case=case.name,precision=precision,repeat=repeat,comparison=cmp,repeat_arrays=rep,
                    pass_all=cmp['pass_all'] and all(rep.values()))
                save(out/'comparison.json',row);rows.append(row)
                print(json.dumps(dict(case=case.name,precision=precision,repeat=repeat,pass_all=row['pass_all'],
                    error_volt=cmp['max_abs_volt'],events=cmp['oracle_events'],bitwise_history=cmp['bitwise_history'])),flush=True)
    save(EXP/'regression/gpu-complete.json',dict(pass_all=all(x['pass_all'] for x in rows),trials=rows,
        wall_seconds=time.perf_counter()-t0,whole_brain_tested=False,learning_achieved=False))

if __name__=='__main__':
    p=argparse.ArgumentParser();p.add_argument('mode',choices=['oracle','gpu']);args=p.parse_args()
    check();globals()[args.mode]()