"""Independent event-scheduled author equation and preserved frozen component arms."""
import sys
from pathlib import Path
import numpy as np

sys.path.insert(0,'/data/digital-drosophila/models/male-cns-lif/lif-semantics-audit-v1')
from math_reference import params, state


def component(nodes,edges,weights,config,policy,seed_node):
    p=params(config['model']);n=len(nodes);steps=round(config['duration_ms']/config['dt_ms'])
    ix={int(x):i for i,x in enumerate(nodes)}
    src=np.array([ix[int(x)] for x in edges.source_neuron_idx],np.int32)
    dst=np.array([ix[int(x)] for x in edges.target_neuron_idx],np.int32)
    v=np.full(n,p['rest'],np.float32);g=np.zeros(n,np.float32);ref=np.zeros(n,np.int32)
    qlen=p['delay_steps']+1;q=np.zeros((qlen,n),np.float32)
    history=np.empty((steps,3,n),np.float32);counts=np.zeros(n,np.int32);ts=[];ns=[]
    for t in range(steps):
        v,g,ref,fire=state(v,g,ref,q[t%qlen],p,policy);q[t%qlen]=0
        if t==0 and seed_node is not None:
            k=ix[seed_node];assert not fire[k] and ref[k]==0
            fire[k]=True;v[k]=p['reset'];ref[k]=p['ref_steps']
            if policy in ('reset_g_only','reset_and_refractory'):g[k]=np.float32(0)
        counts+=fire;active=np.flatnonzero(fire);ts.extend([t]*len(active));ns.extend(nodes[active].tolist())
        selected=np.flatnonzero(fire[src]);selected=selected[np.argsort(dst[selected],kind='stable')]
        for target in np.unique(dst[selected]):
            acc=np.float32(0)
            for e in selected[dst[selected]==target]:acc=np.float32(acc+weights[e])
            slot=(t+p['delay_steps'])%qlen;q[slot,target]=np.float32(q[slot,target]+acc)
        history[t]=v,g,ref
    return dict(step=np.array(ts,np.int32),neuron_idx=np.array(ns,np.int32),counts=counts,v=v,g=g,ref=ref,history=history,queue=q)


def author_numpy(nodes,edges,config,precision,seed_node,external_arrivals=None):
    """Analytic integration in SI volts; sequential arrival writes follow spike detection."""
    f=np.dtype(precision).type;m=config['model'];n=len(nodes);dt=config['dt_ms']
    steps=round(config['duration_ms']/dt);rs=round(m['refractory_ms']/dt);ds=round(m['synaptic_delay_ms']/dt)
    ix={int(x):i for i,x in enumerate(nodes)}
    src=np.array([ix[int(x)] for x in edges.source_neuron_idx],np.int32)
    dst=np.array([ix[int(x)] for x in edges.target_neuron_idx],np.int32)
    weights=np.asarray(edges.raw_count.to_numpy(np.float64)*(m['w_syn_mv_per_anatomical_synapse']*.001),dtype=precision)
    rest=m['resting_potential_mv']*.001;reset=m['reset_potential_mv']*.001;threshold=m['threshold_mv']*.001
    v=np.full(n,rest,dtype=precision);g=np.zeros(n,dtype=precision)
    last=np.full(n,-1000000,np.int64);counts=np.zeros(n,np.int32)
    em=np.exp(-dt/m['membrane_tau_ms']);es=np.exp(-dt/m['synaptic_tau_ms'])
    drive=m['synaptic_tau_ms']/(m['membrane_tau_ms']-m['synaptic_tau_ms'])*(em-es)
    history=np.empty((steps,3,n),np.float64);queue=[[] for _ in range(ds+1)];ts=[];ns=[]
    external=np.zeros((steps,n),dtype=precision) if external_arrivals is None else np.asarray(np.asarray(external_arrivals,np.float64)*.001,dtype=precision)
    assert external.shape==(steps,n)
    for t in range(steps):
        available=(t-last)>=rs
        # Use previous g for the coupled exact differential update, then cast to variable dtype.
        oldg=g.copy()
        v[available]=np.asarray(rest+em*(v[available]-rest)+drive*oldg[available],dtype=precision)
        g[available]=np.asarray(es*oldg[available],dtype=precision)
        if t==0 and seed_node is not None:v[ix[seed_node]]=f(threshold+config['seed_margin_mv']*.001)
        fire=available&(v>threshold);last[fire]=t;counts+=fire
        active=np.flatnonzero(fire);ts.extend([t]*len(active));ns.extend(nodes[active].tolist())
        accept=available&~fire
        due=queue[t%(ds+1)];queue[t%(ds+1)]=[]
        for e in due:
            if accept[dst[e]]:g[dst[e]]=f(g[dst[e]]+weights[e])
        for j in np.flatnonzero(external[t]):
            if accept[j]:g[j]=f(g[j]+external[t,j])
        queue[(t+ds)%(ds+1)].extend(np.flatnonzero(fire[src]).tolist())
        v[fire]=f(reset);g[fire]=f(0)
        ref=np.maximum(rs-(t-last),0).astype(np.int32)
        history[t]=np.asarray(v,np.float64)*1000,np.asarray(g,np.float64)*1000,ref
    return dict(step=np.array(ts,np.int32),neuron_idx=np.array(ns,np.int32),counts=counts,
        v=np.asarray(v,np.float64)*1000,g=np.asarray(g,np.float64)*1000,ref=ref,history=history,weights=np.asarray(weights,np.float64)*1000)