"""Independent reconstruction from bundled cell coordinates and expression tables.
Does not import or execute the author's analysis code. Run inside a container.
"""
import json, pathlib
import numpy as np
import pandas as pd
from scipy.spatial import KDTree
from scipy.stats import spearmanr
r=pathlib.Path('/bundle'); d=r/'data'
p=json.loads((d/'parameters.json').read_text())
sub=pd.read_csv(d/'subclasses.csv',index_col=0)
ids=sub.index[sub['primary']]
x=pd.read_csv(d/'subclass_expression.csv',index_col=0).loc[ids]
genes=pd.read_csv(d/'energy_genes.csv'); genes=genes[genes['in_10x']]
standard=(x-x.mean())/x.std(ddof=0)
components=[]
for category in ['oxphos','pump']:
 selected=[row.symbol for row in genes.itertuples() if row.component==category and x[row.symbol].max()>=p['gene_min_max_subclass_mean_log2']]
 c=standard[selected].mean(axis=1)
 components.append((c-c.mean())/c.std(ddof=0))
e=(components[0]+components[1])/2
z=np.load(d/'cells_allen.npz',allow_pickle=False)
frames=[]
for sec in np.unique(z['sec']):
 within=z['sec']==sec
 neurons=within & (z['kind']==1)
 vessels=within & (z['kind']==2)
 if not vessels.any():continue
 xy=np.column_stack([z['x_um'][neurons],z['y_um'][neurons]]).astype(float)
 vxy=np.column_stack([z['x_um'][vessels],z['y_um'][vessels]]).astype(float)
 distances=KDTree(vxy).query(xy)[0]
 keep=distances<=p['distance_cap_um']
 f=pd.DataFrame({'subclass':z['subclass'][neurons][keep],'tx':np.floor(xy[keep,0]/p['tile_um']),'ty':np.floor(xy[keep,1]/p['tile_um']),'ld':np.log(np.maximum(distances[keep],p['distance_floor_um']))})
 f['residual']=f['ld']-f.groupby(['tx','ty'])['ld'].transform('mean')
 frames.append(f[['subclass','residual']])
allcells=pd.concat(frames,ignore_index=True)
g=allcells.groupby('subclass')['residual'].agg(['mean','count']).reindex(ids)
rho=float(spearmanr(e,g['mean']).statistic)
original=pd.read_csv(r/'results/per_subclass.csv',index_col=0).reindex(ids)
result={'rho':rho,'n_units':len(ids),'neurons_primary_subclasses':int(g['count'].sum()),'neurons_all_eligible':len(allcells),'max_abs_energy_delta_vs_declared_table':float(np.abs(e-original['E']).max()),'max_abs_excess_delta_vs_declared_table':float(np.abs(g['mean']-original['excess']).max()),'hazard_data_inventory':{}}
for file in sorted(d.glob('cells_*.npz')):
 data=np.load(file,allow_pickle=False)
 result['hazard_data_inventory'][file.name]={'arrays':{name:{'dtype':str(data[name].dtype),'shape':list(data[name].shape)} for name in data.files},'kind_codes':np.unique(data['kind']).tolist(),'finite_xy':bool(np.isfinite(data['x_um']).all() and np.isfinite(data['y_um']).all())}
assert abs(rho-(-0.018960304))<1e-7
assert len(ids)==239 and int(g['count'].sum())==2066795
assert result['max_abs_excess_delta_vs_declared_table']<1e-8
print(json.dumps(result,indent=2))