"""Evidence-only audit and unchanged-dynamics real motif tests; no learning."""
import argparse
import ast
from datetime import datetime, timezone
import json
import os
from pathlib import Path
import shutil
import subprocess
import sys
import time
import unittest
import urllib.request
import numpy as np
import pandas as pd
import yaml
from math_reference import params, single, state, cpu_motif

CFG = yaml.safe_load(Path(__file__).with_name('config.yaml').read_text())
ROOT = Path(CFG['root']); MODEL = ROOT/CFG['model']; EXP = ROOT/CFG['experiment']
sys.path.insert(0, str(ROOT/'models/male-cns-lif/apl-pn-scope-audit-v1'))
from audit import save, digest, immutable_check, seal
sys.path.insert(0, str(ROOT/'models/male-cns-lif/terminal-scope-diagnostic-v1'))
from diagnostic import manifest_check, queues


def integrity():
    result = {'original': immutable_check(ROOT, ROOT/CFG['old_snapshot'], 4)}
    for name in CFG['additional_frozen']: result[name] = manifest_check(ROOT/name, digest)
    assert all(x['pass'] for x in result.values()), result
    return result


def prepare():
    if EXP.exists() and any(EXP.iterdir()): raise RuntimeError('Refusing nonempty output')
    EXP.mkdir(exist_ok=True)
    for d in ['protocol', 'external-resources', 'audit', 'synthetic', 'motifs', 'regression', 'performance']:
        (EXP/d).mkdir()
    save(EXP/'protocol/preregistered.json', {'config': CFG, 'before_analysis': True,
        'timestamp_utc': datetime.now(timezone.utc).isoformat(),
        'parent_commit': subprocess.check_output(['git', '-C', str(ROOT), 'rev-parse', 'HEAD'], text=True).strip(),
        'sources': {p.name: digest(p) for p in sorted(MODEL.iterdir()) if p.is_file()}})
    save(EXP/'regression/preflight.json', integrity())
    external = EXP/'external-resources'; commit = CFG['author']['commit']; records = []
    for filename in ['model.py', 'LICENSE', 'environment_full.yml']:
        url = f'https://raw.githubusercontent.com/philshiu/Drosophila_brain_model/{commit}/{filename}'
        with urllib.request.urlopen(url, timeout=60) as response: data = response.read()
        target = external/filename
        with target.open('xb') as out: out.write(data)
        records.append({'path': filename, 'url': url, 'bytes': len(data), 'sha256': digest(target), 'license': 'MIT'})
    assert digest(external/'model.py') == CFG['author']['model_sha256']
    assert 'MIT License' in (external/'LICENSE').read_text()
    text = (external/'model.py').read_text(); ast.parse(text)
    clauses = ['dg/dt = -g / tau', '(unless refractory)', 'v = v_rst; w = 0; g = 0 * mV',
               "method='linear'", "'v > v_th'", 'neu[i].rfc = 0 * ms']
    assert all(clause in text for clause in clauses)
    lines = text.splitlines()
    save(EXP/'audit/author-source-clauses.json', [{'clause': clause, 'line_numbers_1based': [i+1 for i, line in enumerate(lines) if clause in line]} for clause in clauses])
    save(external/'provenance.json', {'files': records, 'commit': commit, 'paper': CFG['author']['paper'],
        'downloaded_code_executed': False, 'new_dependencies_installed': False})
    can = ROOT/CFG['canonical']; metadata = pd.read_parquet(can/'neuron-index/unified-metadata.parquet').sort_values('neuron_idx').reset_index(drop=True)
    ptr = np.load(can/'gpu/csr-indptr.int64.npy', mmap_mode='r'); dst = np.load(can/'gpu/csr-indices.int32.npy', mmap_mode='r'); raw = np.load(can/'gpu/csr-weights.int64.npy', mmap_mode='r')
    assert len(metadata) == 166700 and len(dst) == 25582938 and int(raw.sum()) == 124177617
    nodes = np.array(CFG['selected_nodes'], np.int32)
    meta = metadata.iloc[nodes].copy(); assert np.array_equal(meta.neuron_idx, nodes)
    assert meta.neurotransmitter.eq('acetylcholine').all(), 'No sign inference or reassignment'
    alln = meta[meta.annotation_class.eq('ALLN')]; assert alln.annotation_type.eq('lLN1_bc').all() and alln.nt_ground_truth.eq('acetylcholine').all()
    ids = np.concatenate([np.arange(ptr[n], ptr[n+1], dtype=np.int64) for n in nodes])
    src = np.concatenate([np.full(ptr[n+1]-ptr[n], n, np.int32) for n in nodes])
    keep = np.isin(dst[ids], nodes); ids = ids[keep]; src = src[keep]
    edges = pd.DataFrame({'edge_id': ids, 'source_neuron_idx': src, 'target_neuron_idx': dst[ids], 'raw_count': raw[ids]})
    edges['source_body_id'] = metadata.body_id.to_numpy()[src]; edges['target_body_id'] = metadata.body_id.to_numpy()[dst[ids]]
    assert len(edges) == 13
    meta.to_parquet(EXP/'audit/selected-metadata.parquet', index=False); edges.to_csv(EXP/'audit/selected-real-edges.csv', index=False)
    # Explicit metadata mapping only, not physiological parameter assignment.
    graded = metadata[metadata.annotation_type.isin(['lLN2P_a', 'lLN2P_b', 'lLN2P_c', 'lLN2F_b'])].copy()
    cols = ['neuron_idx','body_id','annotation_type','annotation_hemibrainType','annotation_flywireType','annotation_matchingNotes','neurotransmitter','nt_ground_truth']
    graded[cols].to_csv(EXP/'audit/local-inhibition-crosswalk.csv', index=False)
    frozen_yaml = ROOT/CFG['runtime']/'phase4p.yaml'; pc = yaml.safe_load(frozen_yaml.read_text())
    save(EXP/'protocol/frozen-input-hashes.json', {'runtime_config': digest(frozen_yaml),
        'runtime_cpp': digest(ROOT/CFG['runtime']/'phase4p_ext.cpp'), 'runtime_cuda': digest(ROOT/CFG['runtime']/'phase4p_ext.cu'),
        'baseline_phase4b': digest(ROOT/'models/male-cns-lif/phase4b.py'),
        'metadata': digest(can/'neuron-index/unified-metadata.parquet')})
    save(EXP/'protocol/dynamics.json', pc['model'])
    evidence = [
        ['spike reset g', 'reset to zero', 'retain g'],
        ['g during refractory', 'ODE paused; synaptic writes blocked', 'decay and accept arrivals'],
        ['integration', 'linear exact state updater', 'forward Euler with decayed g'],
        ['threshold', 'strict >', '>='],
        ['within-step arrivals', 'after integration, before reset', 'before integration'],
        ['stimulated neurons', 'voltage Poisson input; refractory=0', 'forced spikes subject to 22-step refractory'],
        ['W_syn', 'free parameter calibrated on FlyWire feeding response', 'same value transferred to MaleCNS; not independently measured']]
    pd.DataFrame(evidence, columns=['item', 'author', 'frozen_adapted']).to_csv(EXP/'audit/semantics-differences.csv', index=False)
    prepared = ['protocol/dynamics.json','audit/selected-real-edges.csv','audit/selected-metadata.parquet',
                'protocol/frozen-input-hashes.json','audit/local-inhibition-crosswalk.csv']
    save(EXP/'protocol/prepared-input-hashes.json', {name:digest(EXP/name) for name in prepared})
    print('PREPARED: pinned author source, 13 actual motif edges, frozen integrity PASS', flush=True)


def synthetic():
    m = json.loads((EXP/'protocol/dynamics.json').read_text()); p = params(m)
    edges = pd.read_csv(EXP/'audit/selected-real-edges.csv'); counts = np.unique(np.r_[1, edges.raw_count.to_numpy(np.int64)])
    results, traces = [], {}
    cases = {}
    for count in counts:
        a = np.zeros(round(100/m['dt_ms']), np.float32); a[0] = np.float32(count)*np.float32(m['w_syn_mv_per_anatomical_synapse'])
        cases[f'one_impulse_contacts_{count}'] = a
    # The separate refractory boundary fixture is a synthetic unit test, not a biological stimulation.
    a = np.zeros(round(20/m['dt_ms']), np.float32)
    a[0] = np.float32(2000)  # enough for a spike at t=0 in every component policy
    a[[1,10,20,22,23]] += np.float32(10)
    cases['refractory_boundary'] = a
    a = np.zeros(round(200/m['dt_ms']), np.float32); a[:round(100/m['dt_ms'])] = np.float32(10)
    cases['sustained_then_off'] = a
    cases['zero_input'] = np.zeros(round(100/m['dt_ms']), np.float32)
    for name, a in cases.items():
        for policy in CFG['synthetic_policies']:
            events, history = single(a, p, policy)
            repeat = single(a, p, policy)
            assert np.array_equal(events, repeat[0]) and np.array_equal(history, repeat[1])
            if policy in ('reset_g_only', 'reset_and_refractory'): assert np.all(history[events,1] == 0)
            if name.startswith('one_impulse') and policy in ('reset_g_only', 'reset_and_refractory'): assert len(events) <= 1
            if name == 'zero_input': assert not len(events)
            if name == 'refractory_boundary' and policy == 'reset_and_refractory':
                assert np.all(history[:22,1] == 0) and history[22,1] > 0
            results.append({'case': name, 'policy': policy, 'spikes': len(events),
                'spike_times_ms': json.dumps((events.astype(float)*m['dt_ms']).tolist()),
                'post_input_spikes': int((events >= 1000).sum()) if name == 'sustained_then_off' else None,
                'final_g': float(history[-1,1]), 'finite': bool(np.isfinite(history).all()), 'exact_repeat': True,
                'not_full_Brian2_replication': True})
            traces[name+'__'+policy] = history
        traces[name+'__input'] = a
    pd.DataFrame(results).to_csv(EXP/'synthetic/results.csv', index=False)
    np.savez_compressed(EXP/'synthetic/traces.npz', **traces)
    save(EXP/'regression/synthetic.json', {'pass': True, 'cases': len(cases), 'policies': 4,
        'proves_only_operator_difference': True, 'brain_model_modified': False, 'Brian2_executed': False})


def gpu_motif(rt, nodes, edges, weights, p, steps, seed):
    import torch
    n = len(nodes); lookup = {int(node): i for i, node in enumerate(nodes)}
    src = np.array([lookup[int(x)] for x in edges.source_neuron_idx], np.int32)
    dst = np.array([lookup[int(x)] for x in edges.target_neuron_idx], np.int32)
    d = rt.device; qlen = p['delay_steps']+1
    v = torch.full((n,), float(p['rest']), dtype=torch.float32, device=d); g = torch.zeros_like(v)
    ref = torch.zeros(n, dtype=torch.int32, device=d); spike = torch.zeros(n, dtype=torch.uint8, device=d); counts = torch.zeros_like(ref)
    queue = torch.zeros((qlen,n), dtype=torch.float32, device=d); wg = torch.as_tensor(weights, device=d)
    history = np.empty((steps,3,n), np.float32); ts, ns = [], []
    for t in range(steps):
        rt.ext.state_step(v,g,ref,queue[t%qlen],spike,counts,float(p['gd']),float(p['alpha']),float(p['rest']),float(p['reset']),float(p['threshold']),p['ref_steps'])
        if t == 0 and seed is not None:
            targets = torch.tensor([lookup[seed]], dtype=torch.int64, device=d); force = torch.ones(1,dtype=torch.uint8,device=d)
            rt.ext.force_step(v,ref,spike,counts,targets,force,float(p['reset']),p['ref_steps'])
        active = spike.cpu().numpy().astype(bool)
        indices = np.flatnonzero(active); ts.extend([t]*len(indices)); ns.extend(nodes[indices].tolist())
        selected = np.flatnonzero(active[src]); selected = selected[np.argsort(dst[selected],kind='stable')]
        if len(selected):
            st = dst[selected]; starts = np.flatnonzero(np.r_[True,st[1:] != st[:-1]])
            rt.ext.segmented_sum(torch.as_tensor(selected,dtype=torch.int64,device=d),
                torch.as_tensor(starts,dtype=torch.int64,device=d), torch.as_tensor(st[starts],dtype=torch.int32,device=d),
                wg,queue[(t+p['delay_steps'])%qlen])
        history[t] = torch.stack([v,g,ref.float()]).cpu().numpy()
    return {'step': np.asarray(ts,np.int32),'neuron_idx':np.asarray(ns,np.int32),'counts':counts.cpu().numpy(),
        'v':v.cpu().numpy(),'g':g.cpu().numpy(),'ref':ref.cpu().numpy(),'history':history,'queue':queue.cpu().numpy()}


def run():
    pre = json.loads((EXP/'protocol/preregistered.json').read_text())
    assert all(digest(MODEL/name) == h for name,h in pre['sources'].items())
    assert not (EXP/'regression/run-start.json').exists(), 'Refusing overwrite'
    prepared = json.loads((EXP/'protocol/prepared-input-hashes.json').read_text())
    assert all(digest(EXP/name) == h for name,h in prepared.items()), 'Prepared-input drift'
    runtime = ROOT/CFG['runtime']; pc = yaml.safe_load((runtime/'phase4p.yaml').read_text())
    assert pc['model'] == json.loads((EXP/'protocol/dynamics.json').read_text())
    suite = unittest.defaultTestLoader.discover(str(MODEL), pattern='test_math.py')
    test_result = unittest.TextTestRunner(verbosity=2).run(suite)
    assert test_result.wasSuccessful() and test_result.testsRun == 7
    save(EXP/'regression/unit-tests.json', {'pass': True, 'tests': test_result.testsRun})
    queues(); os.environ['CUDA_VISIBLE_DEVICES'] = pc['device']['visible_gpu_uuid']
    os.environ['CUDA_HOME'] = '/usr/local/cuda'; os.environ['MAX_JOBS'] = '2'; os.environ['TORCH_CUDA_ARCH_LIST'] = '8.9'
    os.environ['PATH'] = str(ROOT/'models/male-cns-lif/.venv/bin')+':/usr/local/cuda/bin:'+os.environ['PATH']
    sys.path.insert(0,str(ROOT/'models/male-cns-lif')); sys.path.insert(0,str(runtime))
    import torch, phase4p
    free,total = torch.cuda.mem_get_info(); assert free > 2*1024**3
    save(EXP/'regression/run-start.json', {'free_VRAM': free,'device':torch.cuda.get_device_name(),
        'whole_brain_uses_only_frozen_reference':True,'alternative_policies_only_synthetic_single_cell':True})
    start = time.perf_counter(); torch.cuda.reset_peak_memory_stats()
    rt = phase4p.CanonicalRuntime(pc,runtime,EXP); p = params(pc['model'])
    synthetic()
    edges = pd.read_csv(EXP/'audit/selected-real-edges.csv').sort_values('edge_id')
    eid=edges.edge_id.to_numpy(np.int64); src=edges.source_neuron_idx.to_numpy(np.int32)
    assert np.all(eid >= rt.indptr_np[src]) and np.all(eid < rt.indptr_np[src+1])
    assert np.array_equal(rt.indices_np[eid],edges.target_neuron_idx)
    assert np.array_equal(rt.raw_np[eid],edges.raw_count)
    save(EXP/'regression/anatomical-edge-validation.json', {'pass':True,'edges':len(edges),
        'checks':['canonical edge_id in source CSR row','exact target','exact raw int64 count','exact effective float32 weight']})
    effective = rt.weights.cpu().numpy(); rows, regress = [], []
    assert np.array_equal(effective[edges.edge_id], edges.raw_count.to_numpy(np.float32)*np.float32(pc['model']['w_syn_mv_per_anatomical_synapse']))
    for spec in CFG['motifs']:
        queues(); nodes = np.sort(np.array(spec['nodes'],np.int32)); subset = edges[edges.source_neuron_idx.isin(nodes)&edges.target_neuron_idx.isin(nodes)].copy()
        if 'pairs' in spec:
            pairs = {tuple(x) for x in spec['pairs']}
            subset = subset[[tuple(x) in pairs for x in subset[['source_neuron_idx','target_neuron_idx']].to_numpy()]]
            assert len(subset) == len(pairs)
        w = effective[subset.edge_id]; seed = CFG['seed_node'] if spec.get('force',True) else None
        steps = round(CFG['duration_ms']/pc['model']['dt_ms']); t0 = time.perf_counter()
        cpu = cpu_motif(nodes,subset,w,p,steps,seed); gpu = gpu_motif(rt,nodes,subset,w,p,steps,seed)
        eq = {key:bool(np.array_equal(value,gpu[key])) for key,value in cpu.items()}
        save(EXP/'regression'/(spec['name']+'.json'), {'exact':eq,'pass':all(eq.values())}); assert all(eq.values()),eq
        if seed is None: assert len(cpu['step']) == 0
        subset.to_csv(EXP/'motifs'/(spec['name']+'-edges.csv'),index=False)
        pd.DataFrame({'step':cpu['step'],'neuron_idx':cpu['neuron_idx']}).to_parquet(EXP/'motifs'/(spec['name']+'-events.parquet'),index=False)
        np.savez_compressed(EXP/'motifs'/(spec['name']+'-state.npz'), nodes=nodes, **{k:v for k,v in cpu.items() if k not in ('step','neuron_idx')})
        late = cpu['step'] >= round(CFG['late_start_ms']/pc['model']['dt_ms'])
        for node in nodes:
            chosen = cpu['neuron_idx'] == node
            rows.append({'motif':spec['name'],'neuron_idx':int(node),'spikes':int(chosen.sum()),
                'late_spikes':int((chosen&late).sum()),'first_ms':float(cpu['step'][chosen][0]*pc['model']['dt_ms']) if chosen.any() else None,
                'last_ms':float(cpu['step'][chosen][-1]*pc['model']['dt_ms']) if chosen.any() else None,
                'external_force_requests':int(seed == node),'frozen_math_unchanged':True,'whole_brain_necessity_proven':False})
        regress.append({'motif':spec['name'],'CPU_GPU_exact':True,'wall_seconds':time.perf_counter()-t0})
        print(json.dumps({'motif':spec['name'],'events':len(cpu['step']),'late_events':int(late.sum()),'exact':True}),flush=True)
    pd.DataFrame(rows).to_csv(EXP/'motifs/summary.csv',index=False)
    save(EXP/'regression/motifs.json',regress)
    refs=[]
    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')
        expected={'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')}
        r=rt.run(np.flatnonzero(rt.pops[spec['population']]),spec['rate_hz'],spec['seed'],spec['duration_ms'])
        eq={k:bool(np.array_equal(r[k],v)) for k,v in expected.items()};assert all(eq.values()),eq
        refs.append({'trial':spec['name'],'exact':eq,'wall_seconds':r['metrics']['wall_seconds']})
        print(json.dumps({'frozen_trial':spec['name'],'exact':True}),flush=True)
    save(EXP/'regression/frozen-five.json',{'pass':True,'trials':refs,'oracle':'frozen deterministic-v2, not Brian2'})
    save(EXP/'performance/summary.json',{'wall_seconds':time.perf_counter()-start,'peak_allocated_VRAM':int(torch.cuda.max_memory_allocated()),
        'not_runtime_benchmark':True,'ComfyUI_unloaded':False})
    save(EXP/'audit/run-complete.json',{'synthetic_policies':4,'real_motifs_frozen_only':5,'all_CPU_GPU_exact':True,'learning_achieved':False,
        'alternative_MaleCNS_dynamics_run':False,'original_author_code_executed':False})


def finalize():
    assert (EXP/'audit/run-complete.json').exists()
    post=integrity();assert post==json.loads((EXP/'regression/preflight.json').read_text())
    save(EXP/'regression/final-integrity.json',post)
    pre=json.loads((EXP/'protocol/preregistered.json').read_text());assert all(digest(MODEL/n)==h for n,h in pre['sources'].items())
    from report import write_report
    write_report(EXP)
    save(EXP/'audit/final-status.json',{'status':'COMPLETE-SEMANTICS-AUDIT / LEARNING-NOT-ACHIEVED','frozen_integrity_pass':True,
        'documented_adaptation_not_GPU_regression':True,'new_model_authorization_separate':True,'goal_remains_active':True})
    for folder in (MODEL,EXP):
        print('SEALED',folder,seal(folder),flush=True);assert manifest_check(folder,digest)['pass']


if __name__=='__main__':
    parser=argparse.ArgumentParser();parser.add_argument('task',choices=['prepare','run','finalize']);a=parser.parse_args();globals()[a.task]()