#!/usr/bin/env python3
"""Mechanosensory runtime mode matrix and optional state-window forensics."""
from __future__ import annotations

import argparse
import hashlib
import json
import os
from pathlib import Path

import numpy as np
import pyarrow.parquet as pq
import torch
import yaml

from phase4o_a import EXP,FROZEN,ROOT,sha,first_event_diff,save

NAME='mechanosensory-50hz'
SEED=42003


def run(mode: str, focus_step: int, focus_neuron: int):
    from brain_runtime import BrainRuntime
    from phase4i import poisson_batch
    cfgpath=ROOT/'models/male-cns-lif/mb-mechanistic-v1/phase4i.yaml'
    if mode=='deterministic-flags':
        torch.use_deterministic_algorithms(True)
        torch.backends.cudnn.deterministic=True
        torch.backends.cudnn.benchmark=False
        torch.backends.cudnn.allow_tf32=False
        torch.backends.cuda.matmul.allow_tf32=False
        torch.set_float32_matmul_precision('highest')
    rt=BrainRuntime(cfgpath,EXP/'compiler-reports/extension-build')
    target=np.flatnonzero(rt.population_masks['mechanosensory']).astype(np.int64)
    refroot=FROZEN/NAME
    ref=pq.read_table(refroot/'spike-events.parquet')
    rstep=ref['step'].to_numpy().astype(np.int32);rnode=ref['neuron_idx'].to_numpy().astype(np.int32)
    rcount=np.load(refroot/'per-neuron-spike-counts.int32.npy')
    rv=np.load(refroot/'final-voltage.float32.npy');rg=np.load(refroot/'final-g.float32.npy')
    rt.reset(SEED)
    batch=poisson_batch(target,50.,1000.,NAME)
    backend={'python_module':str(Path(rt.runtime_module.__file__)), 'python_sha256':sha(Path(rt.runtime_module.__file__)),
             'extension':str(Path(rt.extension.__file__)), 'extension_sha256':sha(Path(rt.extension.__file__))}
    window=[]
    if mode in ('logging-off','logging-on'):
        p4cfg=yaml.safe_load((ROOT/'models/male-cns-lif/runtime-v1/phase4c.yaml').read_text())
        direct=rt.runtime_module.Runtime(p4cfg,rt.table,rt.indptr_np,rt.indices_np,rt.raw_np,rt.device,rt.extension)
        out=direct.run('hybrid-exact',target,50.,SEED,1000.,capture_events=(mode=='logging-on'))
        step=out['event_steps'];node=out['event_nodes'];counts=out['counts'];v=out['v'];g=out['g']
    elif mode in ('instrumentation-on','single-stream-sync'):
        for t in range(10000):
            oldlen=len(rt.spike_steps)
            if mode=='instrumentation-on' and focus_step-6<=t<=focus_step+5 and focus_neuron>=0:
                slot=t%rt.queue_len
                window.append({'step':t,'when':'before','neuron_idx':focus_neuron,
                               'voltage':float(rt.v[focus_neuron].item()),'g':float(rt.g[focus_neuron].item()),
                               'refractory':int(rt.refractory[focus_neuron].item()),
                               'arrival_buffer':float(rt.delay_queue[slot,focus_neuron].item()),
                               'ring_buffer_sha256':hashlib.sha256(rt.delay_queue.cpu().numpy().tobytes()).hexdigest()})
            rt.step(batch,0.1)
            torch.cuda.synchronize()
            if mode=='instrumentation-on' and focus_step-6<=t<=focus_step+5 and focus_neuron>=0:
                new_nodes=np.concatenate(rt.spike_nodes[oldlen:]) if len(rt.spike_nodes)>oldlen else np.empty(0,np.int32)
                window.append({'step':t,'when':'after','neuron_idx':focus_neuron,
                               'voltage':float(rt.v[focus_neuron].item()),'g':float(rt.g[focus_neuron].item()),
                               'refractory':int(rt.refractory[focus_neuron].item()),'focus_spiked':int(focus_neuron in new_nodes),
                               'active_presynaptic_count':len(new_nodes),
                               'fanout_edge_count':int(np.sum(rt.indptr_np[new_nodes+1]-rt.indptr_np[new_nodes])) if len(new_nodes) else 0,
                               'active_presynaptic_neuron_idx':new_nodes.tolist()[:500]})
        events=rt.read_spikes();step=events['step'].astype(np.int32);node=events['neuron_idx'].astype(np.int32)
        counts=rt.spike_counts.cpu().numpy().astype(np.int32);v=rt.v.cpu().numpy();g=rt.g.cpu().numpy()
    else:
        rt.step(batch,1000.)
        events=rt.read_spikes();step=events['step'].astype(np.int32);node=events['neuron_idx'].astype(np.int32)
        counts=rt.spike_counts.cpu().numpy().astype(np.int32);v=rt.v.cpu().numpy();g=rt.g.cpu().numpy()
    event_known=mode!='logging-off'
    result={'mode':mode,'backend':backend,'logging_events_captured':event_known,
            'spike_events_exact':bool(np.array_equal(step,rstep) and np.array_equal(node,rnode)) if event_known else None,
            'per_neuron_counts_exact':bool(np.array_equal(counts,rcount)),
            'final_voltage_max_abs_error':float(np.max(np.abs(v-rv))),
            'final_g_max_abs_error':float(np.max(np.abs(g-rg))),
            'first_event_difference':first_event_diff(rstep,rnode,step,node) if event_known else None,
            'event_sha256':hashlib.sha256(step.tobytes()+node.tobytes()).hexdigest() if event_known else None,
            'count_sha256':hashlib.sha256(counts.tobytes()).hexdigest(),
            'torch_deterministic_algorithms':torch.are_deterministic_algorithms_enabled(),
            'CUBLAS_WORKSPACE_CONFIG':os.getenv('CUBLAS_WORKSPACE_CONFIG'),
            'state_window':window}
    return result


def main():
    ap=argparse.ArgumentParser();ap.add_argument('mode',choices=['default','logging-off','logging-on','instrumentation-off','instrumentation-on','single-stream-sync','deterministic-flags'])
    ap.add_argument('--focus-step',type=int,default=-1000);ap.add_argument('--focus-neuron',type=int,default=-1)
    ap.add_argument('--repeat',type=int,default=0);a=ap.parse_args()
    try: result=run(a.mode,a.focus_step,a.focus_neuron)
    except Exception as exc: result={'mode':a.mode,'error':repr(exc),'unsupported_or_failed':True}
    result['repeat']=a.repeat
    save(EXP/'divergence'/f'mode-{a.mode}-repeat-{a.repeat}.json',result)
    print(json.dumps({k:result.get(k) for k in ('mode','repeat','spike_events_exact','per_neuron_counts_exact','first_event_difference','error')}),flush=True)


if __name__=='__main__':main()