#!/usr/bin/env python3
"""Reproduce the separately registered v2 exploratory social-factor study.

Dependencies: numpy, scipy. No downloads and no edits to the v1 panel/protocol.
Run prepare first to freeze row IDs, folds and input/code hashes, then fit.
"""
import argparse, csv, hashlib, json, math, sys
from datetime import datetime, timezone
from pathlib import Path
import numpy as np
import scipy
from scipy.optimize import minimize
from scipy.special import gammaln, digamma, xlogy

ROOT=Path(__file__).resolve().parents[1]
OUT=ROOT/'public/research/social-factors/exploratory-v2'
BASE=ROOT/'db/research/social-factors-panel.json'
EDU=ROOT/'db/research/social-exploratory-inputs/education.json'
CRIME=ROOT/'db/research/social-exploratory-inputs/crime.json'
PROTOCOL=OUT/'protocol.json'
SEED=20261010

def sha(p): return hashlib.sha256(p.read_bytes()).hexdigest()
def save(p,d):
    p.parent.mkdir(parents=True,exist_ok=True)
    p.write_text(json.dumps(d,indent=2,allow_nan=False)+'\n')
def read(p): return json.loads(p.read_text())
def stamp(): return datetime.now(timezone.utc).isoformat()
def deviance(y,mu):
    y=np.asarray(y,dtype=float); mu=np.asarray(mu,dtype=float)
    return 2*(xlogy(y,y/mu)-y+mu)

def nb_objective(theta,X,y,offset):
    eta=offset+X@theta[:-1]
    if np.max(eta)>700: return 1e100, np.zeros_like(theta)
    mu=np.exp(eta); r=np.exp(-theta[-1]); den=r+mu
    ll=gammaln(y+r)-gammaln(r)-gammaln(y+1)+r*(np.log(r)-np.log(den))+y*(eta-np.log(den))
    grad_beta=X.T@((mu-y)*r/den)
    grad_alpha=r*np.sum(digamma(y+r)-digamma(r)+np.log(r)+1-np.log(den)-(r+y)/den)
    # Tiny, fixed numerical ridge; no holdout tuning. Intercept is unpenalised.
    penalty=1e-6*np.sum(theta[1:-1]**2)/2
    grad_beta[1:]+=1e-6*theta[1:-1]
    return float(-ll.sum()+penalty),np.r_[grad_beta,grad_alpha]

def fit_nb(X,y,offset):
    start=np.zeros(X.shape[1]+1)
    start[0]=np.log(np.sum(y)/np.sum(np.exp(offset)))
    start[-1]=np.log(.1)
    fit=minimize(nb_objective,start,args=(X,y,offset),jac=True,method='L-BFGS-B',
                 bounds=[(None,None)]*X.shape[1]+[(np.log(1e-6),np.log(10))],
                 options={'maxiter':2000,'ftol':1e-12,'gtol':1e-5,'maxls':50})
    gradient=float(np.max(np.abs(fit.jac)))
    if not fit.success or gradient>.05 or not np.isfinite(fit.fun):
        raise ValueError(f'NB fit failed: {fit.message}; gradient={gradient}')
    return fit.x[:-1],{'alpha':float(np.exp(fit.x[-1])),'converged':bool(fit.success),'gradientMax':gradient,'iterations':int(fit.nit),'negativeLogLikelihood':float(fit.fun)}

BASE_FEATURES=['logDensity2011','ageUnder25Share2011','age65PlusShare2011','lagCarTaxiShare','lagLogMilesPerRoadKm','year']
MODELS={'B':BASE_FEATURES,'D':BASE_FEATURES+['incomeScore','employmentScore'],
        'DE':BASE_FEATURES+['incomeScore','employmentScore','noQualificationsShare2011'],
        'DE4':BASE_FEATURES+['incomeScore','employmentScore','level4PlusShare2011'],
        'BH':BASE_FEATURES+['logHistoryRate'],
        'DH':BASE_FEATURES+['logHistoryRate','incomeScore','employmentScore'],
        'DEH':BASE_FEATURES+['logHistoryRate','incomeScore','employmentScore','noQualificationsShare2011'],
        'DC':BASE_FEATURES+['incomeScore','employmentScore','logBurglaryIntensity'],
        'DCH':BASE_FEATURES+['logHistoryRate','incomeScore','employmentScore','logBurglaryIntensity'],
        'DCE':BASE_FEATURES+['incomeScore','employmentScore','logBurglaryIntensity','noQualificationsShare2011'],
        'DCEH':BASE_FEATURES+['logHistoryRate','incomeScore','employmentScore','logBurglaryIntensity','noQualificationsShare2011']}
COMPARISONS=[('deprivation','D','B'),('education','DE','D'),('education-level4','DE4','D'),
             ('deprivation-history','DH','BH'),('education-history','DEH','DH'),('combined-history','DEH','BH')]
CRIME_COMPARISONS=[('burglary','DC','D'),('burglary-history','DCH','DH'),('education-after-burglary','DCE','DC'),
                   ('education-after-burglary-history','DCEH','DCH'),('all-social-history','DCEH','BH')]

def design(train,test,features):
    # All transforms, including credibility prior and unseen-region intercept, use training only.
    prior_rate=sum(r['collisions'] for r in train)/sum(r['motorMiles']/1e5 for r in train)
    prior_exposure=float(np.median([r['lagMotorMiles']/1e5 for r in train]))
    def matrix(rows):
        return np.array([[np.log((r['lagCollisions']+prior_exposure*prior_rate)/(r['lagMotorMiles']/1e5+prior_exposure))
                          if k=='logHistoryRate' else r[k] for k in features] for r in rows],dtype=float)
    a,b=matrix(train),matrix(test); mean=a.mean(axis=0); scale=a.std(axis=0); scale[scale<1e-12]=1
    a=(a-mean)/scale; b=(b-mean)/scale
    regions=sorted({r['region'] for r in train}); ref=regions[0]
    shares={g:sum(r['region']==g for r in train)/len(train) for g in regions}
    def region_matrix(rows):
        return np.array([[float(r['region']==g) if r['region'] in regions else shares[g] for g in regions[1:]] for r in rows]).reshape(len(rows),len(regions)-1)
    X=np.column_stack([np.ones(len(train)),a,region_matrix(train)])
    Z=np.column_stack([np.ones(len(test)),b,region_matrix(test)])
    meta={'featureNames':['intercept']+features+['region:'+g for g in regions[1:]],'mean':mean.tolist(),'scale':scale.tolist(),
          'referenceRegion':ref,'trainingRegionShares':shares,'unseenRegion':'Training-row-weighted average regional effect',
          'historyPriorRatePer100k':prior_rate,'historyPriorExposureIn100kMiles':prior_exposure}
    return X,Z,meta

def assemble():
    panel=read(BASE); edu=read(EDU); educ={r['areaId']:r for r in edu['rows']}
    areas={r['id']:r for r in panel['areas']}; orig={(r['areaId'],r['year']):r for r in panel['rows']}
    complete={a for a in areas if all((a,y) in orig and orig[a,y]['countExposureEligible'] and orig[a,y]['incomeScore'] is not None and orig[a,y]['employmentScore'] is not None for y in range(2016,2026))}
    included=sorted(complete&set(educ)); rows=[]
    for a in included:
        e=educ[a]
        for year in range(2018,2026):
            r=orig[a,year]; lag=orig[a,year-2]
            row={'id':f'{a}:{year}','areaId':a,'name':areas[a]['name'],'region':areas[a]['region'],'year':year,
                 'collisions':r['collisions'],'motorMiles':r['motorMiles'],'incomeScore':r['incomeScore'],'employmentScore':r['employmentScore'],
                 'deprivationVintage':r['deprivationVintage'],'deprivationPublished':r['deprivationPublished'],
                 'lagYear':year-2,'lagCollisions':lag['collisions'],'lagMotorMiles':lag['motorMiles'],
                 'lagCarTaxiShare':lag['carTaxiExposureShare'],'lagLogMilesPerRoadKm':math.log(lag['motorMiles']/lag['roadLengthKm']),
                 **{k:e[k] for k in ['population2011','populationDensity2011','ageUnder25Share2011','age65PlusShare2011','noQualificationsShare2011','level4PlusShare2011']}}
            row['logDensity2011']=math.log(row['populationDensity2011'])
            assert row['deprivationPublished']<f'{year}-01-01'
            assert row['collisions']>=0 and row['motorMiles']>0
            rows.append(row)
    exclusions=[{'areaId':a,'name':areas[a]['name'],'reason':'Incomplete 2016–2025 eligible count/mileage/deprivation history' if a not in complete else 'No exact-code 2011 Census match'} for a in sorted(set(areas)-set(included))]
    coverage=[]
    for y in range(2016,2026):
        allrows=[r for r in panel['rows'] if r['year']==y]; kept=[r for r in allrows if r['areaId'] in included]
        coverage.append({'year':y,'includedCollisions':sum(r['collisions'] for r in kept),'targetCollisions':sum(r['collisions'] for r in allrows),
                         'includedMotorMiles':sum(r['motorMiles'] for r in kept),'knownTargetMotorMiles':sum(r['sourceMotorMiles'] or 0 for r in allrows)})
    studies=[{'id':'authority','label':'Authority study','rows':rows,'comparisons':COMPARISONS}]
    if CRIME.exists():
        crime=read(CRIME)
        if crime.get('rows'):
            flagged={r['areaId'] for r in crime['rows'] if r['recordingBreakSensitivityExcluded'] or r['crimeQualityExclusions']}
            studies.append({'id':'crime','label':'Police-force burglary extension','rows':[r for r in crime['rows'] if r['areaId'] not in flagged],'comparisons':CRIME_COMPARISONS})
            studies.append({'id':'crime-recording-sensitivity','label':'Crime sensitivity: include recording-break force','rows':crime['rows'],'comparisons':CRIME_COMPARISONS})
    return studies,exclusions,coverage

def make_folds(rows):
    folds=[]
    for year in [2023,2024,2025]:
        train=[r['id'] for r in rows if r['year']<year and r['year'] not in [2020,2021]]
        test=[r['id'] for r in rows if r['year']==year]
        folds.append({'id':f'temporal-{year}','kind':'temporal','year':year,'train':train,'test':test})
    for region in sorted({r['region'] for r in rows}):
        train=[r['id'] for r in rows if r['year']<2025 and r['year'] not in [2020,2021] and r['region']!=region]
        test=[r['id'] for r in rows if r['year']==2025 and r['region']==region]
        folds.append({'id':'spatial-'+region,'kind':'spatial','year':2025,'region':region,'train':train,'test':test})
    # Stress-only fits use pre-pandemic data, with no three-year minimum claim.
    for year in [2020,2021]:
        folds.append({'id':f'stress-{year}','kind':'stress','year':year,
                      'train':[r['id'] for r in rows if r['year']<2020], 'test':[r['id'] for r in rows if r['year']==year]})
    return folds

def prepare():
    if (OUT/'frozen-design.json').exists(): raise ValueError('Frozen design already exists; create a new version to amend it.')
    studies,exclusions,coverage=assemble()
    files=[BASE,EDU,PROTOCOL,Path(__file__)] + ([CRIME] if CRIME.exists() else [])
    design={'id':'social-exploratory-2026-10-10-v2','preparedAt':stamp(),'inputs':[{'path':str(p.relative_to(ROOT)),'sha256':sha(p)} for p in files],
            'studies':[{'id':s['id'],'label':s['label'],'rows':s['rows'],'comparisons':s['comparisons'],'folds':make_folds(s['rows'])} for s in studies],
            'exclusions':exclusions,'coverage':coverage,'software':{'numpy':np.__version__,'scipy':scipy.__version__}}
    for s in design['studies']:
        assert len({r['id'] for r in s['rows']})==len(s['rows'])
        for f in s['folds']:
            assert not set(f['train'])&set(f['test'])
    save(OUT/'frozen-design.json',design)
    print(json.dumps({'prepared':design['preparedAt'],'studies':[{ 'id':s['id'],'rows':len(s['rows']),'areas':len({r['areaId'] for r in s['rows']})} for s in studies]}))

def metric(rows,model):
    y=np.array([r['collisions'] for r in rows]); mu=np.array([r['predictions'][model] for r in rows]); miles=np.array([r['motorMiles'] for r in rows])
    return {'deviance':float(deviance(y,mu).sum()),'observedExpected':float(y.sum()/mu.sum()),
            'observedRatePer100k':float(1e5*y.sum()/miles.sum()),'predictedRatePer100k':float(1e5*mu.sum()/miles.sum()),
            'exposureWeightedAbsoluteRateError':float(1e5*np.abs(y-mu).sum()/miles.sum()),'rows':len(rows)}

def compare(rows,model,base):
    m,b=metric(rows,model),metric(rows,base)
    regions=sorted({r['region'] for r in rows})
    blocks=[]
    for g in regions:
        sub=[r for r in rows if r['region']==g]
        blocks.append([metric(sub,base)['deviance'],metric(sub,model)['deviance']])
    blocks=np.array(blocks); rng=np.random.default_rng(SEED)
    ix=rng.integers(0,len(blocks),size=(2000,len(blocks)))
    totals=blocks[ix].sum(axis=1); gains=1-totals[:,1]/totals[:,0]
    return {'model':model,'comparator':base,'gain':1-m['deviance']/b['deviance'],
            'gainInterval95':np.quantile(gains,[.025,.975]).tolist(),'independentRegionBlocks':len(regions),'modelMetrics':m,'baselineMetrics':b}

def run():
    frozen=read(OUT/'frozen-design.json')
    for item in frozen['inputs']:
        assert sha(ROOT/item['path'])==item['sha256'], 'Frozen input changed: '+item['path']
    predictions=[]; fits=[]; summaries=[]
    for study in frozen['studies']:
        lookup={r['id']:r for r in study['rows']}
        modelids=sorted({m for _,a,b in study['comparisons'] for m in [a,b]})
        for fold in study['folds']:
            train=[lookup[k] for k in fold['train']]; test=[lookup[k] for k in fold['test']]
            if not train or not test: raise ValueError('Empty fold '+fold['id'])
            if fold['kind']!='stress': assert len({r['year'] for r in train})>=3
            rows=[{k:r[k] for k in ['id','areaId','name','region','year','collisions','motorMiles']}|{'study':study['id'],'fold':fold['id'],'kind':fold['kind'],'predictions':{}} for r in test]
            for model in modelids:
                X,Z,transform=design(train,test,MODELS[model]); y=np.array([r['collisions'] for r in train],dtype=float)
                offset=np.log([r['motorMiles']/1e5 for r in train]); beta,diagnostic=fit_nb(X,y,offset)
                pred=np.exp(np.log([r['motorMiles']/1e5 for r in test])+Z@beta)
                assert np.isfinite(pred).all() and (pred>0).all()
                for r,p in zip(rows,pred): r['predictions'][model]=float(p)
                fits.append({'study':study['id'],'fold':fold['id'],'model':model,'trainingRows':len(train),'testRows':len(test),'transform':transform,'coefficients':beta.tolist(),**diagnostic})
            predictions+=rows
            print(study['id'],fold['id'],'fitted',flush=True)
        sr=[r for r in predictions if r['study']==study['id']]
        comparisons=[]
        for cid,model,base in study['comparisons']:
            temporal=[r for r in sr if r['kind']=='temporal']; spatial=[r for r in sr if r['kind']=='spatial']
            annual=[{'year':y,**compare([r for r in temporal if r['year']==y],model,base)} for y in [2023,2024,2025]]
            regional=[{'region':g,**compare([r for r in spatial if r['region']==g],model,base)} for g in sorted({r['region'] for r in spatial})]
            stress=[{'year':y,**compare([r for r in sr if r['kind']=='stress' and r['year']==y],model,base)} for y in [2020,2021]]
            total=compare(temporal,model,base); transfer=compare(spatial,model,base)
            passing=sum(a['gain']>=.01 and .9<=a['modelMetrics']['observedExpected']<=1.1 for a in annual)
            spatial_positive=sum(a['gain']>0 for a in regional)
            candidate=total['gain']>=.01 and total['gainInterval95'][0]>0 and passing==3 and spatial_positive>=5 and transfer['gain']>0 and .9<=total['modelMetrics']['observedExpected']<=1.1
            comparisons.append({'id':cid,**total,'annual':annual,'spatial':transfer,'regions':regional,'stress':stress,
                                'supportiveYears':passing,'positiveRegions':spatial_positive,
                                'status':'Exploratory candidate' if candidate else 'Not consistently supported'})
        summaries.append({'id':study['id'],'label':study['label'],'areas':len({r['areaId'] for r in study['rows']}),'rows':len(study['rows']),
                          'regions':sorted({r['region'] for r in study['rows']}),'comparisons':comparisons})
    save(OUT/'predictions.json',predictions); save(OUT/'fits.json',fits)
    with (OUT/'predictions.csv').open('w') as f:
        fields=['study','fold','kind','areaId','name','region','year','collisions','motorMiles','model','predictedCollisions','predictedRatePer100k']
        w=csv.DictWriter(f,fieldnames=fields); w.writeheader()
        for r in predictions:
            for m,p in r['predictions'].items():w.writerow({k:r[k] for k in fields[:9]}|{'model':m,'predictedCollisions':p,'predictedRatePer100k':p/r['motorMiles']*1e5})
    summary={'id':frozen['id'],'protocolVersion':'v2','baseEvidenceRelease':'lens-2026-10-09-cbec43fe76c4',
             'preparedAt':frozen['preparedAt'],'fittedAt':stamp(),'publicationStatus':'Exploratory research; no human actuarial review recorded',
             'studies':summaries,'coverage':frozen['coverage'],'exclusions':frozen['exclusions'],
             'fitCount':len(fits),'allFitsConverged':all(f['converged'] for f in fits),
             'hashes':{'frozenDesign':sha(OUT/'frozen-design.json'),'protocol':sha(PROTOCOL),'code':sha(Path(__file__))},
             'uncertainty':'2,000 paired bootstrap resamples of English-region blocks using frozen out-of-sample predictions. Nine authority-study blocks; intervals are conditional on fitted models, do not include source measurement error or refitting uncertainty, and are not multiplicity-adjusted.'}
    save(OUT/'results.json',summary)
    save(ROOT/'db/research/social-exploratory-results.json',summary)
    manifest={'id':summary['id'],'createdAt':stamp(),'files':[{'path':str(p.relative_to(ROOT/'public')),'sha256':sha(p),'bytes':p.stat().st_size} for p in sorted(OUT.iterdir()) if p.is_file() and p.name!='manifest.json']}
    save(OUT/'manifest.json',manifest)
    print(json.dumps({'summary':[{ 'id':s['id'],'comparisons':[{ 'id':c['id'],'gain':c['gain'],'interval':c['gainInterval95'],'status':c['status']} for c in s['comparisons']]} for s in summaries]}))

if __name__=='__main__':
    parser=argparse.ArgumentParser(); parser.add_argument('action',choices=['prepare','fit']);args=parser.parse_args()
    prepare() if args.action=='prepare' else run()