"""Full v630 reference-model experiment. Offline CPU runner; no graphical output rules.

Equations, thresholds, resets, weights, delays and sensory inputs come from the
vendored Shiu/Spiller reference. Only connectivity loading and trial orchestration
are changed to fit local memory and to record reproducible controlled trials.
"""
from pathlib import Path
import argparse, ast, csv, gc, hashlib, json, resource, sys, time
from datetime import datetime, timezone
import numpy as np
import pyarrow.parquet as pq
import brian2 as b

HERE = Path(__file__).resolve().parent
sys.path.insert(0, str(HERE / 'reference'))
import model as reference


def notebook_ids():
    found = {}
    for cell in json.loads((HERE/'reference/figures.ipynb').read_text())['cells']:
        try:
            tree = ast.parse(''.join(cell['source']))
        except SyntaxError:
            continue
        for node in tree.body:
            if isinstance(node, ast.Assign) and isinstance(node.targets[0], ast.Name):
                name = node.targets[0].id
                if name in ('neu_sugar', 'neu_bitter', 'ids_mn9'):
                    found[name] = ast.literal_eval(node.value)
    return found


def peak_mib():
    return round(resource.getrusage(resource.RUSAGE_SELF).ru_maxrss / 1024, 1)


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument('--seeds', nargs='+', type=int, default=[1729, 1730, 1731])
    ap.add_argument('--conditions', nargs='+', choices=['rest','sugar','bitter','mixed','disconnected'],
                    default=['rest','sugar','bitter','mixed','disconnected'])
    ap.add_argument('--duration-ms', type=int, default=1000)
    ap.add_argument('--backend', choices=['numpy','cython'], default='cython')
    args = ap.parse_args()
    (HERE/'results').mkdir(exist_ok=True)
    b.prefs.codegen.target = args.backend
    b.defaultclock.dt = .1*b.ms
    params = dict(reference.default_params)
    params['r_poi'] = 100*b.Hz
    params['r_poi2'] = 100*b.Hz
    started = time.perf_counter()
    print('Loading full v630 network; one process, no graph pruning.', flush=True)
    with (HERE/'data/2023_03_23_completeness_630_final.csv').open() as f:
        rows = csv.reader(f);next(rows)
        roots = np.array([int(row[0]) for row in rows], dtype=np.int64)
    id2i = {int(root): i for i,root in enumerate(roots)}
    sets = notebook_ids()
    sugar = [id2i[x] for x in sets['neu_sugar']]
    bitter = [id2i[x] for x in sets['neu_bitter']]
    motor = [id2i[x] for x in sets['ids_mn9']]
    n = b.NeuronGroup(len(roots), model=params['eqs'], method='linear',
                     threshold=params['eq_th'], reset=params['eq_rst'],
                     refractory='rfc', namespace=params, name='default_neurons')
    n.v=params['v_0'];n.g=0;n.rfc=params['t_rfc']
    syn = b.Synapses(n,n,'w : volt',on_pre='g += w',delay=params['t_dly'],name='default_synapses')
    path=HERE/'data/2023_03_23_connectivity_630_final.parquet'
    parquet=pq.ParquetFile(path)
    offset=0;physical_synapses=0;positive_edges=0;negative_edges=0
    cols=['Presynaptic_Index','Postsynaptic_Index','Presynaptic_ID','Postsynaptic_ID',
          'Connectivity','Excitatory x Connectivity']
    for batch in parquet.iter_batches(batch_size=250000,columns=cols,use_threads=False):
        col={k:batch.column(k).to_numpy() for k in cols}
        pre=col['Presynaptic_Index'].astype(np.int32)
        post=col['Postsynaptic_Index'].astype(np.int32)
        assert np.all((pre>=0)&(pre<len(roots))) and np.all((post>=0)&(post<len(roots)))
        assert np.array_equal(roots[pre],col['Presynaptic_ID'])
        assert np.array_equal(roots[post],col['Postsynaptic_ID'])
        signed=col['Excitatory x Connectivity']
        assert np.array_equal(np.abs(signed),col['Connectivity'])
        syn.connect(i=pre,j=post)
        syn.w[offset:offset+len(pre)]=signed*params['w_syn']
        offset+=len(pre)
        physical_synapses+=int(col['Connectivity'].sum())
        positive_edges+=int((signed>0).sum());negative_edges+=int((signed<0).sum())
    del batch,col,pre,post,signed,parquet
    gc.collect()
    assert offset==14687178 and len(roots)==127400
    print(f'Graph ready: {len(roots):,} neurons, {offset:,} edges; peak {peak_mib()} MiB',flush=True)
    # Exact author input helper: one independent PoissonInput per selected cell.
    inputs,n=reference.poi(n,sugar,bitter,params)
    for p in inputs:p.active=False
    n.rfc=params['t_rfc']
    monitor=b.SpikeMonitor(n)
    net=b.Network(n,syn,monitor,*inputs)
    net.store('resting')
    loaded=time.perf_counter()-started
    report={
        'schema':1,'generatedAt':datetime.now(timezone.utc).isoformat(),
        'referenceCommit':'91bdd1e7dcf193f3e7ca5a8933497fcef63b7960',
        'dataset':'FlyWire v630 / author-supplied complete model',
        'neuronCount':len(roots),'edgeCount':offset,'physicalSynapseCount':physical_synapses,
        'positiveEdges':positive_edges,'negativeEdges':negative_edges,
        'inputIds':{k:[str(x) for x in v] for k,v in sets.items()},
        'parameters':{'dtMs':.1,'durationMs':args.duration_ms,'sugarHz':100,'bitterHz':100,
                      'weightMv':.275,'baselineHz':0,'seedPolicy':'reset before every trial'},
        'versions':{'python':sys.version.split()[0],'brian2':b.__version__,'numpy':np.__version__},
        'backend':args.backend,'loadSeconds':round(loaded,3),'trials':[],
        'manifest':json.loads((HERE/'data/manifest.json').read_text()),
        'deviations':['Python 3.13 / Brian2 2.10.1 environment; original equations retained',
                      'Streamed parquet loading with all rows retained; IDs checked on every edge',
                      'Three seeded one-second trials per condition instead of 30 unseeded repeats',
                      'One reusable network restored to identical resting state between trials']}
    for seed in args.seeds:
        for condition in args.conditions:
            net.restore('resting');b.seed(seed)
            excite_sugar=condition in ('sugar','mixed','disconnected')
            excite_bitter=condition in ('bitter','mixed')
            for i,p in enumerate(inputs):p.active=excite_sugar if i<len(sugar) else excite_bitter
            if excite_sugar:n.rfc[sugar]=0*b.ms
            if excite_bitter:n.rfc[bitter]=0*b.ms
            if condition=='disconnected':syn.w=0*b.mV
            print(f'Running {condition} seed {seed}; peak {peak_mib()} MiB',flush=True)
            t0=time.perf_counter();net.run(args.duration_ms*b.ms)
            elapsed=time.perf_counter()-t0
            idx=np.asarray(monitor.i,dtype=np.int32)
            times=np.asarray(monitor.t/b.ms,dtype=np.float64)
            motor_times=[np.round(times[idx==i],4).tolist() for i in motor]
            sugar_mask=np.isin(idx,sugar);bitter_mask=np.isin(idx,bitter)
            motor_mask=np.isin(idx,motor)
            bins=np.arange(0,args.duration_ms+10,10)
            populations={
                'sugar':np.histogram(times[sugar_mask],bins=bins)[0].tolist(),
                'bitter':np.histogram(times[bitter_mask],bins=bins)[0].tolist(),
                'network':np.histogram(times[~(sugar_mask|bitter_mask|motor_mask)],bins=bins)[0].tolist(),
                'motor':np.histogram(times[motor_mask],bins=bins)[0].tolist()}
            name=f'{condition}-{seed}'
            np.savez_compressed(HERE/'results'/f'{name}.npz',indices=idx,times_ms=times,
                                root_ids=roots,seed=np.array(seed))
            trial={'id':name,'condition':condition,'seed':seed,'durationMs':args.duration_ms,
                   'wallSeconds':round(elapsed,3),'motorSpikes':motor_times,
                   'motorCounts':[len(x) for x in motor_times],
                   'activeNeurons':int(len(np.unique(idx))),'totalSpikes':len(idx),
                   'inputSpikes':{'sugar':int(sugar_mask.sum()),'bitter':int(bitter_mask.sum())},
                   'binMs':10,'populations':populations,
                   'rawFile':f'{name}.npz'}
            report['trials'].append(trial);report['peakRssMiB']=peak_mib()
            report['totalWallSeconds']=round(time.perf_counter()-started,3)
            (HERE/'results/results.json').write_text(json.dumps(report,indent=2)+'\n')
            print(f'RESULT {name}: MN9 {trial["motorCounts"]}, active {trial["activeNeurons"]}, '
                  f'{elapsed:.2f}s, peak {peak_mib()} MiB',flush=True)
    print('COMPLETE',flush=True)

if __name__=='__main__':main()