"""Independent complete-pair-set and weight check (including entirely missing pairs)."""
from pathlib import Path
import json
import numpy as np
import pandas as pd
from audit import save

ROOT=Path('/data/digital-drosophila');EXP=ROOT/'experiments/phase4aa-apl-pn-scope'
c=ROOT/'datasets/male-cns-v1.0/processed/neuron-only'
m=pd.read_parquet(c/'neuron-index/unified-metadata.parquet').sort_values('neuron_idx').reset_index(drop=True)
ptr=np.load(c/'gpu/csr-indptr.int64.npy',mmap_mode='r');dst=np.load(c/'gpu/csr-indices.int32.npy',mmap_mode='r');w=np.load(c/'gpu/csr-weights.int64.npy',mmap_mode='r')
src=np.repeat(np.arange(len(m),dtype=np.int32),np.diff(ptr))
apl=np.flatnonzero(m.annotation_type.eq('APL'));kc=np.flatnonzero(m.annotation_class.eq('Kenyon_Cell'));alln=np.flatnonzero(m.annotation_class.eq('ALLN'))
pn=pd.read_parquet(EXP/'connectivity/APL-PN-reciprocity.parquet').target_neuron_idx.unique()
watch=json.loads((ROOT/'experiments/phase4z-pn-feedback-audit/protocol/watched.json').read_text())['all_incoming_ALPNs']
selectors={
    'APL-incident':np.isin(src,apl)|np.isin(dst,apl),
    'APL-target-PN-to-KC':np.isin(src,pn)&np.isin(dst,kc),
    'watched-PN-to-ALLN':np.isin(src,watch)&np.isin(dst,alln),
    'ALLN-to-watched-PN':np.isin(src,alln)&np.isin(dst,watch)}
results=[]
for name,mask in selectors.items():
    frame=pd.read_parquet(EXP/'localization'/(name+'.parquet'))
    frame=frame[frame.source_neuron_idx.notna()&frame.target_neuron_idx.notna()]
    counts=frame.groupby(['source_neuron_idx','target_neuron_idx']).size().reset_index(name='weight')
    actual=counts.to_numpy(np.int64)
    ids=np.flatnonzero(mask)
    expected=np.column_stack((src[ids],dst[ids],w[ids])).astype(np.int64)
    actual=actual[np.lexsort((actual[:,1],actual[:,0]))];expected=expected[np.lexsort((expected[:,1],expected[:,0]))]
    passed=np.array_equal(actual,expected)
    results.append({'group':name,'expected_CSR_pairs':len(expected),'observed_pairs':len(actual),'complete_pair_set_and_weight_exact':passed})
    assert passed,name
save(EXP/'localization/complete-pair-set-validation.json',{'pass':all(r['complete_pair_set_and_weight_exact'] for r in results),'groups':results})
print(json.dumps(results,indent=2))