#!/usr/bin/env python3
"""Real Phase 4C event/CSR atomic-collision micro-reproduction (read-only inputs)."""
from __future__ import annotations

import hashlib
import json
from pathlib import Path

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

from phase4o_a import EXP,FROZEN,ROOT,save


def main():
    trial='mechanosensory-50hz'
    step=6625
    delay=18
    focus=[3942,3954]
    ref=pq.read_table(FROZEN/trial/'spike-events.parquet',columns=['step','neuron_idx']).to_pandas()
    src=ref.loc[ref.step==step-delay,'neuron_idx'].to_numpy(np.int32)
    graph=ROOT/'datasets/male-cns-v1.0/processed/neuron-only/gpu'
    indptr=np.load(graph/'csr-indptr.int64.npy',mmap_mode='r')
    indices=np.load(graph/'csr-indices.int32.npy',mmap_mode='r')
    raw=np.load(graph/'csr-weights.int64.npy',mmap_mode='r')
    cfg=yaml.safe_load((ROOT/'models/male-cns-lif/runtime-v1/phase4c.yaml').read_text())
    meta=pq.read_table(ROOT/'datasets/male-cns-v1.0/processed/neuron-only/neuron-index/unified-metadata.parquet',columns=['neurotransmitter']).column('neurotransmitter').to_pylist()
    signs=[]
    policy=cfg['sign_policy'];mp=policy['neurotransmitter_sign']
    for n in src:
        val=meta[int(n)]
        key='' if val is None else str(val).strip().lower()
        signs.append(policy['unknown_sign'] if not key else policy['unclear_sign'] if key=='unclear' else mp[key])
    edge_ids=np.concatenate([np.arange(int(indptr[n]),int(indptr[n+1]),dtype=np.int64) for n in src]) if len(src) else np.empty(0,np.int64)
    edge_src=np.repeat(src,(indptr[src+1]-indptr[src]).astype(np.int64))
    sign_edge=np.repeat(np.asarray(signs,np.int8),(indptr[src+1]-indptr[src]).astype(np.int64))
    tgt=np.asarray(indices[edge_ids],np.int32)
    w=(np.asarray(raw[edge_ids],np.float32)*sign_edge.astype(np.float32)*np.float32(cfg['model']['w_syn_mv_per_anatomical_synapse'])).astype(np.float32)
    counts=np.bincount(tgt,minlength=166700)
    # Run the unchanged scatter_add_ ordering against the real active-event edge subset.
    target_gpu=torch.as_tensor(tgt,dtype=torch.int64,device='cuda:0')
    weight_gpu=torch.as_tensor(w,dtype=torch.float32,device='cuda:0')
    samples=[];all_samples=[]
    for _ in range(100):
        out=torch.zeros(166700,dtype=torch.float32,device='cuda:0')
        out.scatter_add_(0,target_gpu,weight_gpu)
        all_samples.append(out.cpu().numpy().copy())
        samples.append(all_samples[-1][focus].copy())
    sampled=np.asarray(samples,np.float32)
    full=np.asarray(all_samples,np.float32)
    spread=np.ptp(full,axis=0)
    variable=np.flatnonzero(spread>0)
    detail=[]
    for j,n in enumerate(focus):
        loc=np.flatnonzero(tgt==n)
        contrib=w[loc]
        # Deterministic reference keyed by original anatomical edge_id.
        order=np.argsort(edge_ids[loc],kind='stable')
        acc=np.float32(0)
        for value in contrib[order]: acc=np.float32(acc+value)
        detail.append({'target_neuron_idx':n,'same_target_writes':len(loc),'incoming_source_neuron_idx':edge_src[loc].tolist(),
                       'incoming_edge_ids':edge_ids[loc].tolist(),'incoming_effective_weights_mv':contrib.tolist(),
                       'cpu_edge_order_float32_sum_mv':float(acc),'cpu_float64_sum_mv':float(np.sum(contrib.astype(np.float64))),
                       'gpu_atomic_min_mv':float(sampled[:,j].min()),'gpu_atomic_max_mv':float(sampled[:,j].max()),
                       'gpu_atomic_unique_values':np.unique(sampled[:,j]).tolist()})
    report={'reference_divergence_event_step':step,'arrival_source_step':step-delay,
            'arrival_presynaptic_count':len(src),'fanout_edges':len(edge_ids),
            'targets_with_multiple_concurrent_writes':int(np.count_nonzero(counts>1)),
            'maximum_same_target_writes':int(counts.max()),
            'targets_with_nonidentical_atomic_sum_across_100_repeats':len(variable),
            'max_atomic_sum_range_mv':float(spread.max()),
            'variable_target_examples':variable[:30].tolist(),
            'focus':detail,'gpu_atomic_repetitions':len(samples),
            'interpretation':'Collisions and float-order sensitivity are necessary evidence, not by themselves proof of the first state divergence.'}
    save(EXP/'atomic/real-edge-subset.json',report)
    print(json.dumps({'source_step':step-delay,'fanout_edges':len(edge_ids),'collision_targets':report['targets_with_multiple_concurrent_writes'],
                      'max_fanin':report['maximum_same_target_writes'],'focus':[{'neuron_idx':x['target_neuron_idx'],
                      'writes':x['same_target_writes'],'gpu_unique':len(x['gpu_atomic_unique_values'])} for x in detail]}),flush=True)


if __name__=='__main__':main()