"""Exact float32 frozen state arithmetic plus explicitly labelled synthetic ablations."""
import numpy as np


def params(model):
    f = np.float32
    return {'gd': f(1-model['dt_ms']/model['synaptic_tau_ms']),
            'alpha': f(model['dt_ms']/model['membrane_tau_ms']),
            'rest': f(model['resting_potential_mv']), 'reset': f(model['reset_potential_mv']),
            'threshold': f(model['threshold_mv']),
            'ref_steps': round(model['refractory_ms']/model['dt_ms']),
            'delay_steps': round(model['synaptic_delay_ms']/model['dt_ms'])}


def state(v, g, ref, arrival, p, policy='frozen'):
    if policy not in ('frozen', 'reset_g_only', 'refractory_gate_only', 'reset_and_refractory'):
        raise ValueError('Unknown synthetic policy: '+policy)
    r = np.maximum(ref-1, 0).astype(np.int32)
    available = r == 0
    gi = np.add(np.multiply(g, p['gd'], dtype=np.float32), arrival, dtype=np.float32)
    if policy in ('refractory_gate_only', 'reset_and_refractory'):
        gi = np.where(available, gi, g).astype(np.float32)
    leak = np.subtract(v, p['rest'], dtype=np.float32)
    dv = np.multiply(p['alpha'], np.subtract(gi, leak, dtype=np.float32), dtype=np.float32)
    vi = np.where(available, np.add(v, dv, dtype=np.float32), p['reset']).astype(np.float32)
    fire = available & (vi >= p['threshold'])
    vi[fire] = p['reset']; r[fire] = p['ref_steps']
    if policy in ('reset_g_only', 'reset_and_refractory'): gi[fire] = np.float32(0)
    return vi, gi, r, fire


def single(arrivals, p, policy='frozen'):
    v = np.array([p['rest']], np.float32); g = np.zeros(1, np.float32); r = np.zeros(1, np.int32)
    history = np.empty((len(arrivals), 3), np.float32); events = []
    for t, a in enumerate(arrivals):
        v, g, r, fire = state(v, g, r, np.array([a], np.float32), p, policy)
        if fire[0]: events.append(t)
        history[t] = v[0], g[0], r[0]
    return np.asarray(events, np.int32), history


def cpu_motif(nodes, edges, weights, p, steps, seed):
    assert np.all(np.diff(nodes) > 0), 'Canonical nodes must be sorted and unique'
    assert np.all(np.diff(edges.edge_id.to_numpy()) > 0), 'Canonical edge order required'
    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)
    qlen = p['delay_steps']+1; queue = np.zeros((qlen, n), np.float32)
    v = np.full(n, p['rest'], np.float32); g = np.zeros(n, np.float32); ref = np.zeros(n, np.int32)
    counts = np.zeros(n, np.int32); history = np.empty((steps, 3, n), np.float32)
    ts, ns = [], []
    for t in range(steps):
        slot = t % qlen
        v, g, ref, fire = state(v, g, ref, queue[slot], p); queue[slot] = 0
        if t == 0 and seed is not None:
            i = lookup[seed]
            assert not fire[i] and ref[i] == 0
            fire[i] = True; v[i] = p['reset']; ref[i] = p['ref_steps']
        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]):
            ix = selected[dst[selected] == target]; acc = np.float32(0)
            for e in ix: acc = np.float32(acc+weights[e])
            outslot = (t+p['delay_steps']) % qlen
            queue[outslot, target] = np.float32(queue[outslot, target]+acc)
        history[t] = v, g, ref
    return {'step': np.asarray(ts, np.int32), 'neuron_idx': np.asarray(ns, np.int32), 'counts': counts,
            'v': v, 'g': g, 'ref': ref, 'history': history, 'queue': queue}