#!/usr/bin/env python3
"""Inclusive actuarial review; no outcome-driven dropping or national rebands."""
import argparse,copy,gzip,hashlib,importlib.util,json,os,time
from pathlib import Path
import numpy as np
import pandas as pd
from statsmodels.stats.multitest import multipletests
s=importlib.util.spec_from_file_location('glm',Path(__file__).with_name('ew-glm-engine.py'));g=importlib.util.module_from_spec(s);s.loader.exec_module(g)
ROOT=g.b.ROOT;OUT=ROOT/'public/research/ew-deciles/glm-review-v1';p=g.f7.p
INPUTS=g.f7.INPUTS+['public/research/ew-deciles/f7-v1/register.json','public/research/ew-deciles/f6-v1/quality-audit.json']
def setup(out=None):
 global OUT
 if out:OUT=Path(out)
 OUT.mkdir(parents=True,exist_ok=True);p.OUT=OUT

def dictionary():return json.loads((ROOT/'public/research/ew-deciles/f6-v1/feature-dictionary.json').read_text())['features']
def freeze():
 metas=dictionary();primary=[m['id'] for m in metas if m.get('bandedCount',0)>0];diagnostic=[m['id'] for m in metas if m.get('validCount',0)>0 and not m.get('bandedCount',0)];spec={'version':'ew-glm-review-v1','frozenDate':'2026-10-13','sourceCommit':'0ee1c697a43a4c599ed9adab35752790a6ffa0b9','userInstruction':'Include all features for which a GLM can be fitted; user decides drops and simplifications.','primaryFeatures':primary,'diagnosticFeatures':diagnostic,'primary':'Unpenalised Poisson log-link frequency GLM, all40 frozen nationally banded features simultaneously, offset log(vehicleMiles/100000). All eligible zero counts retained. No screening on prior p-values or gains.','core':['intercept','year factor','Wales x year','sum-zero region','urban/rural','four allocation-quality shares','estimated-major-mileage share'],'form':'Primary features categorical, replacing their raw copies in the old baseline. Reference D5 or structural zero where available; nearest observed band otherwise. No automatic merging, dropping or shrinkage.','benchmark':'Frozen Phase4 rich_poisson refit on identical observations. Different functional form, so non-nested held-out comparison; not a likelihood-ratio test.','development':[2022,2023],'validation':'2022 fit ->2023 test;10 leave-one-region-out2022–23 folds. Refit representations and core within each training fold. Unsupported held-out categorical levels withheld from scoring for both comparator and candidate, counted explicitly.','laterYears':'2024/2025 previously inspected diagnostics only. Latest trained year effect carried forward. No untouched validation.','relativities':'Pooled2022–23 conditional coefficient contrasts and cluster-sandwich95% intervals. Report exact reference, raw rates, summed exposure/collisions, fitted rates, weak support, rank and estimability. No causal interpretation of counterfactual changes in compositional features.','dropOne':'Refit without each of40 features in2022; score2023 on the full-model supported cohort. Gain=1-D(full)/D(without). Conditional on every other feature; may be negative. Also report development deviance change; not out-of-sample evidence.','diagnostics':'Each of10 unbanded raw numeric features added separately to same40-feature model, refitting both on its complete-case cohort. Restricted cubic spline on source values, training quantiles5/27.5/50/72.5/95. Four crime rates and mean terrain gradient log1p; other diagnostics identity. Also show linear per-unit sensitivity. Raw grid values are not national deciles. Speed and terrain are unvalidated context proxies; current speed is not historic speed.','diagnosticFailure':'If sparse subset prevents a finite/identifiable joint fit, retain explicit failure; do not invent neutral effects or silently replace with reduced controls.','inference':'162 permanent source-area/authority unions inherited from Phase5 where available. Finite-sample clustered sandwich. No full-model omnibus Wald because coefficient count may exceed cluster covariance rank. Per-feature conditional Wald p and BH q over40 primary or10 exploratory diagnostics separately, unestimable p=1. These are review aids, never drop rules.','adjacent':'Pairwise adjacent-band contrasts use full covariance; intervals including1 do not establish equivalence or approve a merge. Structural zero remains labelled.','overlap':'Unweighted unique-development-LSOA Spearman raw-value correlations and adjusted GVIF^(1/(2df)) where model full-rank. These are overlap diagnostics, not drop criteria.','uncertainty':'Paired fixed-prediction block intervals,1000 draws seed20261012; no refit/selection/exposure uncertainty. Exact null-space contrast checks; zero-event levels and unidentifiable contrasts withheld.','sensitivities':'Refit primary model under benchmark and road-length offsets, same cohort; matched-counted-route outcome/physical-mileage fit with allocation quality set zero. These are sensitivity models, not new selection.','decisionWorkflow':'All feature decisions initially Undecided. User can mark Keep, Drop or Simplify and describe grouping. Browser-local/exported proposals; changes require explicit refit, updating all other coefficients and validation. No multiplying away dropped factors or averaging relativities.','scope':'User-led inclusive GLM review, not completion of the earlier automated F8 penalised/nested-selection protocol. No commercial index/insurer-claims validation. National values and bands and previous results immutable.'}
 path=OUT/'specification.json'
 if path.exists():assert json.loads(path.read_text())==spec,'Frozen review protocol differs'
 else:p.dump('specification.json',spec)
 checks=[{'path':x,'bytes':(ROOT/x).stat().st_size,'sha256':p.sha(ROOT/x)} for x in INPUTS];p.dump('input-checksums.json',checks);return spec

def cohort(d,features,extra=None):
 mask=d.commonEligible.copy()
 for f in features:mask &= np.isfinite(d[f+'_value']) & d[f+'_band'].notna()
 if extra:mask &= np.isfinite(d[extra+'_value'])
 return d[mask].copy()
def label(meta,band):return 'Z0' if band==0 else ('P' if meta.get('zeroBand') else 'D')+str(int(band))
def score(d,base,candidate,sample,scenario='fitted',feature=None):
 bm=base.predict(d);bm=bm[0] if isinstance(bm,tuple) else bm;mu,unknown=candidate.predict(d);valid=~unknown
 if isinstance(base,g.Fit):valid &= ~base.design.unknown(d)
 r,blocks=g.e.paired(d.loc[valid],bm[valid],mu[valid],np.zeros(valid.sum(),bool));r.update(sample=sample,scenario=scenario,feature=feature,eligibleRows=len(d),unsupportedRows=int((~valid).sum()));return r,blocks,mu,valid

def band_effects(fit,d,meta,master,sample='development',scenario='fitted'):
 fid=meta['id'];rows=[];ref=fit.design.refs[fid];mu=fit.predict(d)[0];allbands=range(0 if meta.get('zeroBand') else 1,11)
 for k in allbands:
  mask=(d[fid+'_band']==k).to_numpy();z=d.loc[mask];a=fit.band_contrast(fid,k);effect=fit.contrast(a) if a is not None else {'estimable':False,'rr':None,'lower95':None,'upper95':None,'logEffect':None,'se':None};coll=int(z.motorCollisions.sum());miles=float(z[fit.mileage].sum());expected=float(mu[mask].sum())
  if not len(z) or coll==0:effect.update(estimable=False,rr=None,lower95=None,upper95=None,logEffect=None,se=None)
  raw=master.loc[master[fid+'_band']==k,fid+'_value'];rows.append({'feature':fid,'sample':sample,'scenario':scenario,'band':label(meta,k),'bandValue':k,'referenceBand':label(meta,ref),'nationalLSOAs':len(raw),'minimumValue':float(raw.min()) if len(raw) else None,'maximumValue':float(raw.max()) if len(raw) else None,'rows':len(z),'lsoas':int(z.lsoa.nunique()),'collisions':coll,'miles':miles,'expected':expected,'rawFrequency':coll/miles*100000 if miles else None,'fittedFrequency':expected/miles*100000 if miles else None,'clusters':int(z.cluster.nunique()),'weakSupport':len(z)<100 or coll<100 or z.cluster.nunique()<10,**effect})
 adjacent=[];observed=[r for r in rows if r['rows']]
 for left,right in zip(observed[:-1],observed[1:]):
  a=fit.band_contrast(fid,right['bandValue'],left['bandValue']);v=fit.contrast(a) if a is not None else {'estimable':False,'rr':None,'lower95':None,'upper95':None,'logEffect':None,'se':None}
  if not left['collisions'] or not right['collisions']:v.update(estimable=False,rr=None,lower95=None,upper95=None)
  adjacent.append({'fromBand':left['band'],'toBand':right['band'],'includesOne':v['lower95']<=1<=v['upper95'] if v['estimable'] else None,**v})
 return rows,adjacent

def extra_effects(fit,d,fid):
 vals=d[fid+'_value'];grid=np.unique(np.quantile(vals,np.linspace(0,1,21)));ref=fit.design.extraRef;t=pd.concat([d.iloc[[0]]]*(len(grid)+1),ignore_index=True);t[fid+'_value']=np.r_[grid,ref];x=fit.design.matrix(t,full=True).toarray();rows=[]
 for i,value in enumerate(grid):rows.append({'value':float(value),'referenceValue':ref,**fit.contrast(x[i]-x[-1])})
 return rows

def save_model(name,fit):p.dump('models/'+name+'.json.gz',fit.record())
def run(args):
 setup(args.output);spec=freeze()
 if args.freeze_only:return
 started=time.time();d,master=g.f7.load();metas=dictionary();md={m['id']:m for m in metas};features=spec['primaryFeatures'];eligible=cohort(d,features);dev=eligible[eligible.year<=2023];tr=dev[dev.year==2022];te=dev[dev.year==2023]
 p.csv('model-cohort.csv.gz',eligible[['lsoa','year','cluster','region','country','urbanRural','motorCollisions','preferredMiles']]);p.csv('cluster-map.csv.gz',d[['lsoa','authority','cluster','source_authority']].drop_duplicates());coverage=[]
 for dim in ['year','region','urbanRural','country']:
  for name,z in d.groupby(dim,sort=True):
   keys=eligible.index.intersection(z.index);x=eligible.loc[keys];coverage.append({'dimension':dim,'group':str(name),'nationalRows':len(z),'baselineEligibleRows':int(z.commonEligible.sum()),'jointRows':len(x),'excludedRows':len(z)-len(x),'collisions':int(x.motorCollisions.sum()),'miles':float(x.preferredMiles.sum())})
 p.csv('coverage.csv',coverage)
 print('Main cohort',len(dev),'rows',dev.lsoa.nunique(),'LSOAs; fitting all40',flush=True)
 final=g.Fit(dev,features,covariance=True);save_model('joint-development',final);model22=g.Fit(tr,features,covariance=True);save_model('joint-train2022',model22);base22=g.e.Baseline(tr,'preferredMiles');basefinal=g.e.Baseline(dev,'preferredMiles');metrics=[];prediction=[];blocks=[];subgroups=[]
 def evaluation(z,baseline,model,sample,scenario='fitted'):
  r,bb,mu,valid=score(z,baseline,model,sample,scenario);metrics.append(r);blocks.extend([{'sample':sample,'scenario':scenario,**v} for v in bb.to_dict('records')]);zz=z.loc[valid].copy();zz['expected']=mu[valid];zz['miles']=zz[model.mileage]
  for row in g.b.group_rows(zz):subgroups.append({'sample':sample,'scenario':scenario,**row})
  if scenario=='fitted':
   pp=z[['lsoa','year','region','motorCollisions']].copy();pp['sample']=sample;pp['miles']=z[model.mileage];pp['predictedCount']=np.where(valid,mu,np.nan);pp['predictedFrequencyPer100k']=pp.predictedCount/pp.miles*100000;pp['supported']=valid;prediction.extend(pp.to_dict('records'))
  return r,bb,mu,valid
 evaluation(te,base22,model22,'temporal_2023');pieces=[]
 for region in sorted(dev.region.unique()):
  train=dev[dev.region!=region];test=dev[dev.region==region];model=g.Fit(train,features);base=g.e.Baseline(train,'preferredMiles');r,bb,mu,valid=evaluation(test,base,model,'region_'+region);bm=base.predict(test);z=test.loc[valid,['lsoa','cluster','motorCollisions']].copy();z['base']=bm[valid];z['mu']=mu[valid];pieces.append(z);save_model('holdout-'+region.replace(' ','_'),model);print('Geographic fold',region,round(r['gain']*100,3),'%',flush=True)
 zz=pd.concat(pieces).sort_index();r,bb=g.e.paired(zz,zz.base.to_numpy(),zz.mu.to_numpy(),np.zeros(len(zz),bool));r.update(sample='spatial_pooled',scenario='fitted',eligibleRows=len(dev),unsupportedRows=len(dev)-len(zz),feature=None);metrics.append(r);blocks.extend([{'sample':'spatial_pooled','scenario':'fitted',**v} for v in bb.to_dict('records')])
 for year in [2024,2025]:evaluation(eligible[eligible.year==year],basefinal,final,'diagnostic_'+str(year))
 # In-sample fitted frequencies are provided separately from held-out scores.
 mu=final.predict(dev)[0];fitted=dev[['lsoa','year','region','motorCollisions','preferredMiles']].copy();fitted['predictedCount']=mu;fitted['predictedFrequencyPer100k']=mu/dev.preferredMiles*100000;p.csv('fitted-development.csv.gz',fitted)
 effects={};cards=[];uni=dev.drop_duplicates('lsoa');corr=uni[[f+'_value' for f in features]].corr(method='spearman');p.csv('raw-feature-correlations.csv',corr.rename_axis('feature').reset_index())
 for fid in features:
  meta=md[fid];rows,adj=band_effects(final,dev,meta,master);annual22,_=band_effects(model22,tr,meta,master,'year_2022');effects[fid]=rows+annual22
  others=[f for f in features if f!=fid];reduced=g.Fit(tr,others);r,bb,_,_=score(te,reduced,model22,'drop_one_temporal',feature=fid);save_model('without-'+fid+'-2022',reduced)
  # Development comparison is explicitly in-sample; no likelihood-ratio p.
  reddev=g.Fit(dev,others);dr=float(g.b.deviances(dev.motorCollisions.to_numpy(),reddev.fit['mu']).sum());df=float(g.b.deviances(dev.motorCollisions.to_numpy(),final.fit['mu']).sum())
  overlaps=sorted([{'feature':f,'label':md[f]['label'],'spearman':float(corr.loc[fid+'_value',f+'_value'])} for f in others],key=lambda x:-abs(x['spearman']))[:5]
  card={'feature':fid,'label':meta['label'],'metadata':meta,'scope':'joint40','tested':True,'decision':'undecided','representation':'Frozen national categorical bands','reference':label(meta,final.design.refs[fid]),'developmentRows':len(dev),'developmentLSOAs':dev.lsoa.nunique(),'clusters':final.clusters,'effects':effects[fid],'adjacent':adj,'overlap':overlaps,'adjustedGVIF':final.gvif(fid),'dropOne':r,'inSampleDevianceIncreaseIfDropped':dr-df,'inSampleDfRemoved':final.design.rank-reddev.design.rank,'metrics':[],'curves':[],**final.wald(fid)}
  cards.append(card);p.dump('features/'+fid+'.json',card);p.csv('drop-blocks/'+fid+'.csv.gz',bb);print('Feature',fid,'drop-one gain',round(r['gain']*100,3),'%',flush=True)
 # Annual2023 coefficients and alternative-mileage sensitivities.
 annual23=g.Fit(te,features,covariance=True);save_model('joint-year2023',annual23)
 for card in cards:
  rows,_=band_effects(annual23,te,md[card['feature']],master,'year_2023');card['effects'].extend(rows)
 for scenario,mileage in [('benchmark','benchmarkMiles'),('length','roadLengthMiles')]:
  mt=g.Fit(tr,features,mileage);mf=g.Fit(dev,features,mileage,covariance=True);bf=g.e.Baseline(tr,mileage);evaluation(te,bf,mt,'temporal_2023',scenario);save_model('joint-'+scenario+'-2022',mt);save_model('joint-'+scenario+'-development',mf)
  for card in cards:
   rows,_=band_effects(mf,dev,md[card['feature']],master,'development',scenario);card['effects'].extend(rows)
 counted=g.e.counted_scope(eligible);ct=counted[counted.year==2022];cv=counted[counted.year==2023];cd=counted[counted.year<=2023];cm=g.Fit(ct,features,'physicalCountedMiles');cf=g.Fit(cd,features,'physicalCountedMiles',covariance=True);cb=g.e.Baseline(ct,'physicalCountedMiles');evaluation(cv,cb,cm,'temporal_2023','counted_routes');save_model('joint-counted-2022',cm);save_model('joint-counted-development',cf)
 for card in cards:
  rows,_=band_effects(cf,cd,md[card['feature']],master,'development','counted_routes');card['effects'].extend(rows)
 # Broad inclusion of every raw-valued diagnostic, without inventing national bands.
 for fid in spec['diagnosticFeatures']:
  meta=md[fid];z=cohort(d,features,fid);zd=z[z.year<=2023];zt=zd[zd.year==2022];zv=zd[zd.year==2023];card={'feature':fid,'label':meta['label'],'metadata':meta,'scope':'restricted_extension','tested':False,'decision':'undecided','representation':'Restricted cubic spline of raw values; no national bands','developmentRows':len(zd),'developmentLSOAs':zd.lsoa.nunique(),'clusters':zd.cluster.nunique(),'reference':None,'effects':[],'adjacent':[],'overlap':[],'adjustedGVIF':None,'dropOne':None,'metrics':[],'curves':[],'pvalue':1.,'df':0,'waldEstimable':False}
  try:
   ef=g.Fit(zd,features,extra=fid,covariance=True);et=g.Fit(zt,features,extra=fid);eb=g.Fit(zt,features);ref=g.Fit(zd,features);r,bb,_,_=score(zv,eb,et,'temporal_2023',feature=fid);card.update(tested=True,reference=ef.design.extraRef,curves=extra_effects(ef,zd,fid),metrics=[r],adjustedGVIF=ef.gvif(fid),**ef.wald(fid));save_model('extra-'+fid+'-development',ef);save_model('extra-'+fid+'-2022',et);save_model('extra-'+fid+'-baseline2022',eb);p.csv('drop-blocks/'+fid+'.csv.gz',bb)
   linear=g.Fit(zd,features,extra=fid,extra_form='linear',covariance=True);save_model('extra-'+fid+'-linear',linear);lr=extra_effects(linear,zd,fid);card['linearCurves']=lr
   v=zd[fid+'_value'];card['numericSupport']={'minimum':float(v.min()),'p05':float(v.quantile(.05)),'median':float(v.median()),'p95':float(v.quantile(.95)),'maximum':float(v.max()),'collisions':int(zd.motorCollisions.sum()),'miles':float(zd.preferredMiles.sum()),'zeroRows':int((zd.motorCollisions==0).sum()),'rawFrequency':float(zd.motorCollisions.sum()/zd.preferredMiles.sum()*100000),'fittedFrequency':float(ef.fit['mu'].sum()/zd.preferredMiles.sum()*100000),'regions':zd.region.nunique(),'conditionNumber':ef.design.condition,'aliasedColumns':ef.design.dropped}
   for dim in ['region','urbanRural']:
    card.setdefault('coverage',[]).extend([{'dimension':dim,'group':str(name),'rows':len(gg),'lsoas':gg.lsoa.nunique(),'collisions':int(gg.motorCollisions.sum())} for name,gg in zd.groupby(dim,sort=True)])
  except (ArithmeticError,np.linalg.LinAlgError,ValueError) as ex:card['failureReason']=str(ex);card['tested']=False;card['curves']=[];card['pvalue']=1.;card['waldEstimable']=False
  cards.append(card);p.dump('features/'+fid+'.json',card);print('Diagnostic',fid,card['tested'],card.get('failureReason',''),flush=True)
 for meta in metas:
  if meta['id'] not in features+spec['diagnosticFeatures']:cards.append({'feature':meta['id'],'label':meta['label'],'metadata':meta,'scope':'no_data','tested':False,'decision':'unavailable','representation':'No numeric observations','effects':[],'curves':[],'adjacent':[],'overlap':[],'metrics':[],'failureReason':meta.get('note','No numerical source values'),'pvalue':1.,'df':0,'waldEstimable':False})
 for scope in ['joint40','restricted_extension']:
  group=[c for c in cards if c['scope']==scope];qs=multipletests([c['pvalue'] for c in group],method='fdr_bh')[1]
  for c,q in zip(group,qs):c['qvalue']=float(q)
 register=[]
 for c in cards:
  p.dump('features/'+c['feature']+'.json',c);row={k:c.get(k) for k in ['feature','label','scope','tested','decision','representation','reference','developmentRows','developmentLSOAs','clusters','adjustedGVIF','pvalue','qvalue','df','waldEstimable','inSampleDevianceIncreaseIfDropped','inSampleDfRemoved','failureReason']};row['family']=c['metadata'].get('family');row['temporalConditionalGain']=c['dropOne']['gain'] if c.get('dropOne') else c['metrics'][0]['gain'] if c['metrics'] else None;row['card']='features/'+c['feature']+'.json';row['nationalNumericLSOAs']=c['metadata'].get('validCount',0);row['nationalBandedLSOAs']=c['metadata'].get('bandedCount',0);register.append(row)
 p.dump('register.json',register);p.csv('feature-review.csv',register);p.csv('all-relativities.csv.gz',[r for c in cards for r in c['effects']]);p.csv('adjacent-band-contrasts.csv',[{'feature':c['feature'],**r} for c in cards for r in c['adjacent']]);p.csv('diagnostic-curves.csv',[{'feature':c['feature'],**r} for c in cards for r in c['curves']]);p.csv('metrics.csv',metrics);p.csv('validation-blocks.csv.gz',blocks);p.csv('calibration.csv',subgroups);p.csv('validation-predictions.csv.gz',prediction)
 p.dump('decision-template.json',{'release':'ew-glm-review-v1','status':'proposed decisions only; refit required','features':[{'feature':r['feature'],'decision':'undecided' if r['tested'] else 'unavailable','simplification':'','notes':''} for r in register]})
 summary={'version':'ew-glm-review-v1','nationalLSOAs':35672,'primaryFeatures':len(features),'numericDiagnostics':len(spec['diagnosticFeatures']),'diagnosticsFitted':sum(c['tested'] and c['scope']=='restricted_extension' for c in cards),'noNumericData':sum(c['scope']=='no_data' for c in cards),'jointDevelopmentRows':len(dev),'jointDevelopmentLSOAs':dev.lsoa.nunique(),'jointAllYearRows':len(eligible),'developmentCollisions':int(dev.motorCollisions.sum()),'developmentMiles':float(dev.preferredMiles.sum()),'observedDevelopmentFrequency':float(dev.motorCollisions.sum()/dev.preferredMiles.sum()*100000),'fittedDevelopmentFrequency':float(final.fit['mu'].sum()/dev.preferredMiles.sum()*100000),'zeroDevelopmentRows':int((dev.motorCollisions==0).sum()),'coefficients':final.design.rank,'designColumns':len(final.design.names),'designCondition':final.design.condition,'aliasedColumns':final.design.dropped,'clusters':final.clusters,'pearsonDispersion':final.dispersion,'metrics':metrics,'automaticDrops':0,'automaticMerges':0,'priorResultsChanged':False,'nationalBandsChanged':False,'claimsValidation':False,'independentValidation':False,'fullF8Completed':False}
 p.dump('summary.json',summary);print(json.dumps({k:v for k,v in summary.items() if k!='metrics'}),flush=True);print('Seconds',round(time.time()-started,1),flush=True)
if __name__=='__main__':
 ap=argparse.ArgumentParser();ap.add_argument('--output');ap.add_argument('--freeze-only',action='store_true');run(ap.parse_args())