#!/usr/bin/env python3
"""Frozen London / South East / South West LSOA mileage pilot.
Run with Python >=3.11, numpy pandas scipy scikit-learn shapely>=2 pyogrio
geopandas pyarrow pyproj networkx. Sources are pinned in source-register.json.
python scripts/build-frost-mileage.py --sources /path/to/mileage-sources --cache /path/to/cache
Stages are cached; --stage prepares one stage (prepare, major, model, outcomes, package).
"""
from pathlib import Path
import argparse, calendar, collections, gzip, hashlib, importlib.util, io, json, math, os, platform, struct, sys, time, zipfile
import numpy as np
import pandas as pd
import geopandas as gpd
import pyogrio, shapely, networkx as nx, sklearn, scipy, pyproj, pyarrow
from shapely.strtree import STRtree
from scipy.spatial import cKDTree
from sklearn.compose import ColumnTransformer
from sklearn.preprocessing import OneHotEncoder, StandardScaler
from sklearn.linear_model import PoissonRegressor
from sklearn.pipeline import make_pipeline
from sklearn.metrics import mean_poisson_deviance
ROOT=Path(__file__).resolve().parents[1]; OLD=ROOT/'public/research/frost/v1'; OUT=ROOT/'public/research/frost/mileage-v1'
REGIONS=['London','South East','South West']; YEARS=[2021,2022,2023,2024]
spec=importlib.util.spec_from_file_location('frostgeometry',ROOT/'scripts/build-frost-features.py'); geo=importlib.util.module_from_spec(spec); spec.loader.exec_module(geo)
def log(s): print(time.strftime('%H:%M:%S ')+s,flush=True)
def sha(p): return geo.sha256(p)
def js(p,d): geo.json_file(p,d)
def csv(p,d):
    d.to_csv(p,index=False,compression={'method':'gzip','mtime':0} if str(p).endswith('.gz') else None)
def truth(v): return v.astype(str).str.lower().eq('true')
def load_inputs():
    f=pd.read_csv(OLD/'lsoa-frost-features.csv'); f=f[f.region.isin(REGIONS)].sort_values('lsoa').reset_index(drop=True)
    a=pd.read_csv(OLD/'authority-year-panel.csv'); a=a[a.region.isin(REGIONS)&a.year.isin(YEARS)]
    return f,a

def prepare(args):
    f,a=load_inputs(); C=args.cache; C.mkdir(parents=True,exist_ok=True)
    pinned={'os-open-roads-april2026.zip':'bb2bbff8be3d637b38ed5d06277856c15f41ea262614edf974be8f4b96fdc791','ons-lsoa21-bfc-v10.zip':'3f9e1af25956e4e0bf28e3a3823af44ce7c315cad19c50660ab99352cdf23969'}
    for name,h in pinned.items(): assert sha(args.sources/name)==h,name
    bzip=args.sources/'ons-lsoa21-bfc-v10.zip'
    with zipfile.ZipFile(bzip) as z: member=next(n for n in z.namelist() if n.lower().endswith('.shp'))
    bounds=pyogrio.read_dataframe(f'/vsizip/{bzip}/{member}'); code=next(c for c in bounds if c.upper()=='LSOA21CD')
    bounds=bounds.set_index(code).loc[f.lsoa].reset_index(); bounds.geometry=shapely.make_valid(bounds.geometry.values)
    assert bounds.crs.to_epsg()==27700 and len(bounds)==13972
    polygons=np.array(bounds.geometry); tree=STRtree(polygons); bb=bounds.total_bounds
    bounds.to_parquet(C/'boundaries.parquet',index=False)
    pop=pd.read_csv(ROOT/'public/research/parking/v4/inputs/oa-population-rurality.csv.gz')
    pop['urbanPop']=pop.population*pop.settlement.eq('Urban'); pop=pop.groupby('lsoa')[['population','urbanPop']].sum()
    f=f.join(pop,on='lsoa'); assert f.population.notna().all();f['urbanFraction']=f.urbanPop/f.population.clip(lower=1);f['density']=f.population/f.areaKm2;f['logDensity']=np.log1p(f.density);f['settlement']=np.where(f.urbanFraction>=.5,'urban_resident','rural_resident');f['li']=np.arange(len(f))
    f.to_parquet(C/'features.parquet',index=False)
    rzip=args.sources/'os-open-roads-april2026.zip'; allroads=[]; excluded=collections.Counter()
    with zipfile.ZipFile(rzip) as z:
        for n in sorted(z.namelist()):
            if not n.endswith('_RoadLink.shp'): continue
            with z.open(n) as h: header=h.read(100)
            tile=struct.unpack('<4d',header[36:68])
            if tile[2]<bb[0] or tile[0]>bb[2] or tile[3]<bb[1] or tile[1]>bb[3]: continue
            d=pyogrio.read_dataframe(f'/vsizip/{rzip}/{n}',bbox=tuple(bb),columns=['identifier','fictitious','class','roadNumber','function','startNode','endNode','formOfWay'])
            fict=truth(d.fictitious); restricted=d['function'].eq('Restricted Local Access Road');excluded['fictitious']+=int(fict.sum());excluded['restricted']+=int((~fict&restricted).sum())
            d=d[~fict&~restricted].copy(); allroads.append(d);log(f'Road tile {n}: {len(d):,} access-screened links')
    roads=gpd.GeoDataFrame(pd.concat(allroads,ignore_index=True),crs=27700)
    duplicates=int(roads.identifier.duplicated().sum())
    # Duplicate tile identifiers must have the same geometry and attributes.
    for _,g in roads[roads.identifier.duplicated(False)].groupby('identifier'):
        assert len(set(shapely.to_wkb(shapely.normalize(g.geometry.values))))==1
        assert g.drop(columns='geometry').fillna('').drop_duplicates().shape[0]==1
    roads=roads.drop_duplicates('identifier').reset_index(drop=True)
    roads['ri']=np.arange(len(roads));roads['roadNumber']=roads.roadNumber.fillna('').str.upper().str.replace(' ','',regex=False)
    roads['roadClass']=np.where(roads['class'].eq('Motorway'),'M',np.where(roads['class'].eq('A Road'),'A',np.where(roads['class'].eq('B Road'),'B','minor')))
    degrees=pd.concat([roads.startNode,roads.endNode]).value_counts(); roads['degree']=(roads.startNode.map(degrees)+roads.endNode.map(degrees))/2
    log(f'Clipping {len(roads):,} links to {len(polygons):,} LSOAs')
    pieces=[]; duplicate_length=0
    for start in range(0,len(roads),20000):
        ids,ls,geoms,lengths,dup,n=geo.allocate_lines(np.array(roads.geometry.iloc[start:start+20000]),polygons,tree)
        if len(ids): pieces.append(pd.DataFrame({'ri':ids+start,'li':ls,'metres':lengths}))
        duplicate_length+=dup
    pieces=pd.concat(pieces,ignore_index=True).groupby(['ri','li'],as_index=False).metres.sum(); pieces=pieces[pieces.metres>1e-5]
    roads.to_parquet(C/'roads.parquet',index=False);pieces.to_parquet(C/'pieces.parquet',index=False)
    sources=json.load(open(OLD/'source-register.json')); aadf=sources['majorRoadSources'][0]
    zp=ROOT/aadf['archivedPath'];assert sha(zp)==aadf['archiveSha256'];raw=gzip.decompress(zp.read_bytes());assert hashlib.sha256(raw).hexdigest()==aadf['sha256']
    z=zipfile.ZipFile(io.BytesIO(raw)); x=pd.read_csv(z.open(next(n for n in z.namelist() if n.endswith('.csv'))),low_memory=False)
    x=x[x.year.isin(YEARS)&x.local_authority_code.isin(a.areaId)].copy();assert not x.duplicated(['count_point_id','year']).any()
    x['road_name']=x.road_name.fillna('').str.upper().str.replace(' ','',regex=False);x.to_parquet(C/'aadf.parquet',index=False)
    js(C/'geometry-audit.json',{'lsoas':len(f),'authorities':f.authority.nunique(),'linksInBoundingBox':len(roads),'linksIntersectingPilot':pieces.ri.nunique(),'pieces':len(pieces),'networkKm':pieces.metres.sum()/1000,'excludedBoundingBoxLinks':dict(excluded),'duplicateTileIdsRemoved':duplicates,'coincidentBoundaryKmAssignedOnce':duplicate_length/1000})
    js(OUT/'source-register.json',{'date':'2026-10-11','designSha256':sha(OUT/'design.json'),'sourceBaseline':'05be2ff301b326047d82f074b6fd8daa5411e5d8','newArchives':json.load(open(args.sources/'acquisition.json')),'majorRoadSources':sources['majorRoadSources'][:-1],'frozenInputs':{str(p.relative_to(ROOT)):sha(p) for p in [OLD/'lsoa-frost-features.csv',OLD/'lsoa-year-inventory.csv.gz',OLD/'authority-year-panel.csv',ROOT/'public/research/parking/v4/inputs/oa-population-rurality.csv.gz']},'licence':'Contains Ordnance Survey, ONS, Met Office and DfT data under the Open Government Licence; Crown copyright and database right 2026','limitations':['April 2026 network reused retrospectively for 2021-2024','Resident settlement proxy is not a road urban/rural classification','Public access not independently established by OS function']})
    log('Geometry preparation complete')

def major(args):
    C=args.cache;f=pd.read_parquet(C/'features.parquet');r=gpd.read_parquet(C/'roads.parquet');p=pd.read_parquet(C/'pieces.parquet');x=pd.read_parquet(C/'aadf.parquet'); a=load_inputs()[1]
    p=p.merge(f[['li','authority']],on='li').merge(r[['ri','roadNumber','roadClass']],on='ri'); majorp=p[p.roadClass.isin(['M','A'])]
    fallback_number={(auth,num):g for (auth,num),g in majorp.groupby(['authority','roadNumber'])}; fallback_class={(auth,cl):g for (auth,cl),g in majorp.groupby(['authority','roadClass'])}
    perroad={};cache={}; alloc=[];audit=[]; accepted=[]
    source=json.load(open(OLD/'source-register.json'))['majorRoadSources'];roadshares={k:g[['li','authority','metres']] for k,g in majorp.groupby('ri')}
    for year in YEARS:
        s=next(s for s in source if s['id']==f'geography-mrdb-{year}');b=gzip.decompress((ROOT/s['archivedPath']).read_bytes());assert hashlib.sha256(b).hexdigest()==s['sha256'];path=C/f'mrdb-{year}.zip';path.write_bytes(b)
        with zipfile.ZipFile(path) as z: member=next(n for n in z.namelist() if n.endswith('.shp'))
        endpoints=pyogrio.read_dataframe(f'/vsizip/{path}/{member}').set_index('CP_Number')
        rows=x[(x.year==year)&x.road_type.eq('Major')]
        for row in rows.itertuples():
            cp=int(row.count_point_id);num=row.road_name;auth=row.local_authority_code;cl='M' if num.startswith('M') or '(M)' in num else 'A'
            miles=float(row.all_motor_vehicles)*float(row.link_length_km)/1.609344*(366 if calendar.isleap(year) else 365)
            status='unmatched';reason='no_endpoint';routeids=[];ratio=np.nan;snap=np.nan;cpdist=np.nan;own=np.nan
            if cp in endpoints.index:
                chord=endpoints.loc[cp].geometry;coords=np.asarray(chord.coords); key=(num,tuple(coords[0]),tuple(coords[-1]),row.easting,row.northing,float(row.link_length_km),auth)
                if key in cache: status,reason,routeids,ratio,snap,cpdist,own=cache[key]
                else:
                    if num not in perroad:
                        rr=r[(r.roadNumber==num)&r.roadClass.isin(['M','A'])];g=nx.Graph();nodecoords={}
                        for q in rr.itertuples():
                            c=np.asarray(q.geometry.coords);u=q.startNode;v=q.endNode;length=q.geometry.length
                            nodecoords[u]=c[0,:2];nodecoords[v]=c[-1,:2]
                            if not g.has_edge(u,v) or length<g[u][v]['weight']:g.add_edge(u,v,weight=length,ri=q.ri)
                        nodes=list(nodecoords);perroad[num]=(g,nodes,cKDTree(np.array([nodecoords[n] for n in nodes])) if nodes else None)
                    g,nodes,nt=perroad[num]
                    reason='no_same_number_network'
                    if nt is not None:
                        distances,idx=nt.query(coords[[0,-1]]);snap=float(max(distances));reason='endpoint_snap'
                        if snap<=100:
                            try:
                                pathnodes=nx.shortest_path(g,nodes[idx[0]],nodes[idx[1]],weight='weight');routeids=[g[u][v]['ri'] for u,v in zip(pathnodes[:-1],pathnodes[1:])]
                                geom=shapely.union_all(r.geometry.iloc[routeids].values);ratio=geom.length/(row.link_length_km*1000);cpdist=shapely.distance(geom,shapely.Point(row.easting,row.northing));reason='route_length'
                                if .8<=ratio<=1.2:
                                    reason='count_point_off_route'
                                    if cpdist<=75:
                                        shares=pd.concat([roadshares[i] for i in routeids if i in roadshares]);own=shares.loc[shares.authority==auth,'metres'].sum()/geom.length;reason='authority_coverage'
                                        if own>=.99: status='routed';reason='accepted'
                            except (nx.NetworkXNoPath,nx.NodeNotFound,ValueError): reason='disconnected_route'
                    cache[key]=(status,reason,routeids,ratio,snap,cpdist,own)
            if status=='routed':
                shares=pd.concat([roadshares[i] for i in routeids if i in roadshares]);shares=shares[shares.authority==auth];accepted.extend({'year':year,'countPoint':cp,'ri':i,'roadNumber':num} for i in routeids)
            else:
                shares=fallback_number.get((auth,num));status='road_number_fallback'
                if shares is None:shares=fallback_class.get((auth,cl));status='road_class_fallback'
                if shares is None:status='unallocated'
            audit.append({'year':year,'authority':auth,'countPoint':cp,'roadNumber':num,'vehicleMiles':miles,'estimationMethod':row.estimation_method,'allocation':status,'routeFailure':reason,'routeLengthRatio':ratio,'endpointSnapM':snap,'countPointRouteDistanceM':cpdist,'ownAuthorityRouteShare':own})
            if shares is not None and len(shares):
                weights=shares.groupby('li').metres.sum();weights=weights/weights.sum()
                alloc.extend({'year':year,'authority':auth,'countPoint':cp,'li':int(li),'majorMiles':miles*w,'routedMajorMiles':miles*w if status=='routed' else 0} for li,w in weights.items())
        log(f'Major routing {year}: {len(rows):,} sections')
    audit=pd.DataFrame(audit);section_alloc=pd.DataFrame(alloc);section_alloc.to_parquet(C/'section-allocation.parquet',index=False);alloc=section_alloc.groupby(['year','authority','li'],as_index=False)[['majorMiles','routedMajorMiles']].sum()
    csv(OUT/'major-section-lsoa-allocation.csv.gz',section_alloc);csv(OUT/'accepted-major-route-links.csv.gz',pd.DataFrame(accepted).drop_duplicates());audit.to_parquet(C/'major-audit.parquet',index=False);alloc.to_parquet(C/'major-allocation.parquet',index=False);pd.DataFrame(accepted).drop_duplicates().to_parquet(C/'accepted-routes.parquet',index=False)
    csv(OUT/'major-section-audit.csv.gz',audit);log('Major routing complete')

CATS=['region','roadClass','function']; NUMS=['logDensity','urbanFraction','degree']
def model_fit(d):
    prep=ColumnTransformer([('cat',OneHotEncoder(handle_unknown='ignore',sparse_output=False),CATS),('num',StandardScaler(),NUMS)])
    model=make_pipeline(prep,PoissonRegressor(alpha=1,max_iter=1000,tol=1e-7))
    w=1/d.groupby('count_point_id').count_point_id.transform('size');model.fit(d[CATS+NUMS],d.all_motor_vehicles,poissonregressor__sample_weight=w)
    return model

def bench_fit(d):
    d=d.copy();d['w']=1/d.groupby('count_point_id').count_point_id.transform('size');d['wy']=d.w*d.all_motor_vehicles
    parent=d.groupby('roadClass')[['wy','w']].sum();means=(parent.wy/parent.w).to_dict();grand=d.wy.sum()/d.w.sum();group=d.groupby(['region','roadClass','settlement'])[['wy','w']].sum()
    return {k:(v.wy+20*means.get(k[1],grand))/(v.w+20) for k,v in group.iterrows()},means,grand

def bench_predict(m,d):
    groups,parents,grand=m
    return np.array([groups.get((r.region,r.roadClass,r.settlement),parents.get(r.roadClass,grand)) for r in d.itertuples()])
def metrics(d,pred):
    y=np.asarray(d.all_motor_vehicles);pred=np.asarray(pred);w=1/d.groupby('count_point_id').count_point_id.transform('size').to_numpy()
    return {'rows':len(d),'sites':d.count_point_id.nunique(),'blocks':d.block.nunique(),'bias':float(np.sum(w*(pred-y))/np.sum(w*y)),'wape':float(np.sum(w*abs(pred-y))/np.sum(w*y)),'deviance':float(mean_poisson_deviance(y,np.maximum(pred,1e-8),sample_weight=w))}

def model(args):
    C=args.cache;f=pd.read_parquet(C/'features.parquet');r=gpd.read_parquet(C/'roads.parquet');p=pd.read_parquet(C/'pieces.parquet');x=pd.read_parquet(C/'aadf.parquet');b=gpd.read_parquet(C/'boundaries.parquet')
    pts=x[x.road_type.eq('Minor')&x.estimation_method.eq('Counted')].copy().reset_index(drop=True);pts['roadClass']=np.where(pts.road_name.str.startswith('B'),'B','minor')
    pts['ri']=-1;pts['matchDistanceM']=np.nan
    for cl in ['B','minor']:
        rr=r[r.roadClass.eq(cl)];tree=STRtree(np.array(rr.geometry));inds=np.where(pts.roadClass.eq(cl))[0];points=shapely.points(pts.loc[inds,'easting'],pts.loc[inds,'northing']);pairs,dist=tree.query_nearest(points,max_distance=50,all_matches=True,return_distance=True)
        # Preserve only unique nearest alignments; exact ties remain unmatched.
        counts=collections.Counter(pairs[0]);keep=np.array([counts[i]==1 for i in pairs[0]]);pairs=pairs[:,keep];dist=dist[keep]
        targets=inds[pairs[0]];roadindex=rr.ri.to_numpy()[pairs[1]];pts.loc[targets,'ri']=roadindex;pts.loc[targets,'matchDistanceM']=dist
    # Point geography, not arbitrary road centroid, supplies local predictors.
    pairs=STRtree(np.array(b.geometry)).query(shapely.points(pts.easting,pts.northing),predicate='intersects');mapping={}
    for i,j in zip(*pairs):mapping.setdefault(i,j)
    pts['li']=[mapping.get(i,-1) for i in range(len(pts))]
    pts=pts.merge(r[['ri','function','degree','roadNumber']],on='ri',how='left').merge(f[['li','lsoa','authority','region','logDensity','urbanFraction','settlement','roadDecile']],on='li',how='left')
    pts['matched']=pts.ri.ge(0)&pts.li.ge(0)&pts.local_authority_code.eq(pts.authority)&(~pts.roadClass.eq('B')|pts.road_name.eq(pts.roadNumber))
    pts['block']=(pts.easting//20000).astype(int).astype(str)+'_'+(pts.northing//20000).astype(int).astype(str)
    # Fixed geographic folds, same for repeated sites. Coordinates are static per site.
    canonical=pts.sort_values('year').drop_duplicates('count_point_id').set_index('count_point_id').block
    pts['block']=pts.count_point_id.map(canonical);pts['fold']=pts.block.map(lambda z:int(hashlib.sha256(z.encode()).hexdigest()[:8],16)%5)
    csv(OUT/'count-site-matching.csv.gz',pts[['count_point_id','year','local_authority_code','road_name','ri','lsoa','matchDistanceM','matched','block','fold']])
    d=pts[pts.matched&pts.all_motor_vehicles.ge(0)].copy();train=d[d.year<=2023].copy();test=d[d.year==2024].copy();train['benchmarkOof']=np.nan;train['modelOof']=np.nan;test['benchmarkGeo']=np.nan;test['modelGeo']=np.nan
    for fold in range(5):
        tr=train[train.fold!=fold];te=train[train.fold==fold];m=model_fit(tr);bm=bench_fit(tr);train.loc[te.index,'benchmarkOof']=bench_predict(bm,te);train.loc[te.index,'modelOof']=m.predict(te[CATS+NUMS]);future=test[test.fold==fold];test.loc[future.index,'benchmarkGeo']=bench_predict(bm,future);test.loc[future.index,'modelGeo']=m.predict(future[CATS+NUMS]);log(f'Traffic block fold {fold+1}/5: {len(te)} observations')
    bm=bench_fit(train);m=model_fit(train);mb=metrics(train,train.benchmarkOof);mm=metrics(train,train.modelOof);selected='benchmark' if mb['deviance']<=1.01*mm['deviance'] else 'model'
    test['benchmark']=bench_predict(bm,test);test['model']=m.predict(test[CATS+NUMS]);test['newSite']=~test.count_point_id.isin(train.count_point_id)
    validation=[]
    for label,dd in [('development_oof',train),('2024_geographic_holdout',test),('2024_all',test),('2024_new_sites',test[test.newSite]),('2024_repeat_sites',test[~test.newSite])]:
        for name in ['benchmark','model']:
            col=name+'Oof' if label=='development_oof' else name+'Geo' if label=='2024_geographic_holdout' else name
            for groupby in [[],['region'],['roadClass'],['settlement'],['roadDecile']]:
                groups=[('all',dd)] if not groupby else dd.groupby(groupby[0],dropna=False)
                for key,g in groups:
                    if not len(g):continue
                    z=metrics(g,g[col]);supported=z['sites']>=30 and z['blocks']>=3
                    validation.append({'sample':label,'model':name,'dimension':groupby[0] if groupby else 'all','group':str(key),**z,'supported':supported,'biasGatePassed':supported and abs(z['bias'])<=.10})
    validation=pd.DataFrame(validation);csv(OUT/'traffic-validation.csv',validation)
    pred=pd.concat([train.assign(sample='development_oof').rename(columns={'benchmarkOof':'benchmark','modelOof':'model'}),test.assign(sample=np.where(test.newSite,'2024_new_sites','2024_repeat_sites'))]);csv(OUT/'count-site-predictions.csv.gz',pred[['count_point_id','year','lsoa','region','roadClass','settlement','roadDecile','block','fold','all_motor_vehicles','benchmark','model','sample']]);csv(OUT/'count-site-2024-geographic-predictions.csv.gz',test[['count_point_id','year','lsoa','region','roadClass','settlement','roadDecile','block','fold','all_motor_vehicles','benchmarkGeo','modelGeo']])
    # Each clipped minor fragment gets predictors from its own LSOA.
    fragments=p.merge(r[['ri','roadClass','function','degree']],on='ri');fragments=fragments[fragments.roadClass.isin(['B','minor'])].merge(f[['li','region','authority','logDensity','urbanFraction','settlement']],on='li')
    fragments['benchmark']=bench_predict(bm,fragments);fragments['model']=m.predict(fragments[CATS+NUMS]);
    for name in ['benchmark','model']:fragments[name+'Weight']=fragments.metres*fragments[name]
    support=train.groupby(['region','roadClass','function']).count_point_id.nunique().rename('trainingGroupSites');fragments=fragments.join(support,on=['region','roadClass','function']);fragments['trainingGroupSites']=fragments.trainingGroupSites.fillna(0);fragments['unsupportedFunction']=~fragments['function'].isin(train['function']);fragments['sparseGroup']=fragments.trainingGroupSites.lt(30);fragments['densityExtrapolation']=~fragments.logDensity.between(train.logDensity.min(),train.logDensity.max());
    for flag in ['unsupportedFunction','sparseGroup','densityExtrapolation']:fragments[flag+'ModelWeight']=fragments.modelWeight*fragments[flag]
    coverage=fragments.groupby(['region','roadClass','function']).agg(networkKm=('metres',lambda x:x.sum()/1000),trainingGroupSites=('trainingGroupSites','first'),unsupportedFunction=('unsupportedFunction','first'));csv(OUT/'minor-network-support.csv',coverage.reset_index())
    weights=fragments.groupby('li')[['metres','benchmarkWeight','modelWeight','unsupportedFunctionModelWeight','sparseGroupModelWeight','densityExtrapolationModelWeight']].sum().reset_index().rename(columns={'metres':'minorNetworkMetres'})
    weights.to_parquet(C/'minor-weights.parquet',index=False)
    # Reduced network input makes every per-link estimate inspectable without shipping raw shapefiles.
    csv(OUT/'minor-road-predictions.csv.gz',fragments[['ri','li','authority','roadClass','function','degree','metres','benchmark','model','trainingGroupSites','unsupportedFunction','sparseGroup','densityExtrapolation']])
    js(OUT/'traffic-model.json',{'selected':selected,'selectionSample':'2021-2023 five geographic-block folds','benchmarkDevelopment':mb,'modelDevelopment':mm,'trainingRows':len(train),'trainingSites':train.count_point_id.nunique(),'trainingBlocks':train.block.nunique(),'countedRows':len(pts),'matchedRows':len(d),'heldOut2024Rows':len(test),'heldOutNewSites':test[test.newSite].count_point_id.nunique(),'features':CATS+NUMS,'fittedIntercept':float(m[-1].intercept_),'coefficients':dict(zip(m[0].get_feature_names_out(),m[-1].coef_.tolist())),'standardiserMeans':m[0].named_transformers_['num'].mean_.tolist(),'standardiserScales':m[0].named_transformers_['num'].scale_.tolist(),'benchmarkGroups':[{'group':list(k),'mean':v} for k,v in bm[0].items()],'benchmarkParents':bm[1],'benchmarkGlobalMean':bm[2]})
    log(f'Traffic candidate selected: {selected}')

def poisson_offset(train,test,exposure,frost):
    from scipy.optimize import minimize
    cats=['region']+(['band'] if frost else []);nums=['logDensity','urbanFraction','majorShare']
    prep=ColumnTransformer([('cat',OneHotEncoder(handle_unknown='ignore',sparse_output=False,drop='first'),cats),('num',StandardScaler(),nums)])
    X=np.column_stack([np.ones(len(train)),prep.fit_transform(train[cats+nums])]);T=np.column_stack([np.ones(len(test)),prep.transform(test[cats+nums])]);y=train.assignedMotorCollisions.to_numpy();offset=np.log(train[exposure].to_numpy()/1e6);pen=np.ones(X.shape[1]);pen[0]=0
    def objective(beta):
        eta=X@beta+offset;mu=np.exp(eta);return np.sum(mu-y*eta)+.5*np.sum(pen*beta**2), X.T@(mu-y)+pen*beta
    initial=np.zeros(X.shape[1]);initial[0]=np.log(max(y.sum(),1)/(train[exposure].sum()/1e6))
    fit=minimize(objective,initial,jac=True,method='L-BFGS-B',options={'maxiter':1000,'ftol':1e-12,'gtol':1e-6});assert fit.success,fit.message
    return np.exp(T@fit.x+np.log(test[exposure]/1e6)),{'coefficients':fit.x.tolist(),'features':['intercept']+prep.get_feature_names_out().tolist(),'iterations':int(fit.nit),'trainingDeviance':float(dev(y,np.exp(X@fit.x+offset)).sum()),'trainingRows':len(train)}
def dev(y,mu):
    y=np.asarray(y,dtype=float);mu=np.maximum(np.asarray(mu),1e-12);term=np.zeros_like(y);positive=y>0;term[positive]=y[positive]*np.log(y[positive]/mu[positive]);return 2*(term-y+mu)
def comparison_rows(test,p0,p1,scenario,scope):
    result=[];d=test.copy();d['d0']=dev(d.assignedMotorCollisions,p0);d['d1']=dev(d.assignedMotorCollisions,p1);d['p0']=p0;d['p1']=p1
    for region in ['Pilot',*REGIONS]:
        g=d if region=='Pilot' else d[d.region==region]
        if not len(g):continue
        agg=g.groupby(['region','authority'])[['d0','d1']].sum();rng=np.random.default_rng(41106);samples=[]
        for _ in range(1000):
            draws=[]
            for _,grp in agg.groupby(level=0):draws.append(grp.to_numpy()[rng.integers(0,len(grp),len(grp))].sum(axis=0))
            z=np.sum(draws,axis=0);samples.append(100*(z[0]-z[1])/z[0] if z[0]>0 else np.nan)
        lo,hi=np.nanpercentile(samples,[2.5,97.5]);result.append({'scope':scope,'scenario':scenario,'region':region,'lsoas':len(g),'authorities':g.authority.nunique(),'collisions':int(g.assignedMotorCollisions.sum()),'baselineDeviance':g.d0.sum(),'frostDeviance':g.d1.sum(),'gainPercent':100*(g.d0.sum()-g.d1.sum())/g.d0.sum(),'low':lo,'high':hi,'baselineOE':g.assignedMotorCollisions.sum()/g.p0.sum(),'frostOE':g.assignedMotorCollisions.sum()/g.p1.sum()})
    return result

def outcomes(args):
    C=args.cache;f=pd.read_parquet(C/'features.parquet');a=load_inputs()[1];w=pd.read_parquet(C/'minor-weights.parquet');ma=pd.read_parquet(C/'major-audit.parquet');majoralloc=pd.read_parquet(C/'major-allocation.parquet');chosen=json.load(open(OUT/'traffic-model.json'))['selected']
    sums=ma.groupby(['authority','year']).vehicleMiles.sum().rename('publishedMajorMiles');fallback=ma[ma.allocation!='routed'].groupby(['authority','year']).vehicleMiles.sum().rename('fallbackMajorMiles');unallocated=ma[ma.allocation=='unallocated'].groupby(['authority','year']).vehicleMiles.sum().rename('unallocatedMajorMiles')
    auth=a.rename(columns={'areaId':'authority'}).merge(sums,on=['authority','year'],how='left').merge(fallback,on=['authority','year'],how='left').merge(unallocated,on=['authority','year'],how='left');auth[['publishedMajorMiles','fallbackMajorMiles','unallocatedMajorMiles']]=auth[['publishedMajorMiles','fallbackMajorMiles','unallocatedMajorMiles']].fillna(0);auth['minorResidualMiles']=auth.motorMiles-auth.publishedMajorMiles;auth['majorFallbackShare']=auth.fallbackMajorMiles/auth.publishedMajorMiles.replace(0,np.nan);auth['denominatorAvailable']=auth.motorMiles.gt(0)&auth.minorResidualMiles.ge(0)&auth.unallocatedMajorMiles.eq(0)
    panel=f.merge(pd.DataFrame({'year':YEARS}),how='cross').merge(w,on='li',how='left').merge(majoralloc,on=['li','authority','year'],how='left').merge(auth[['authority','year','motorMiles','minorResidualMiles','majorFallbackShare','denominatorAvailable','rateEligible','qualityReasons']],on=['authority','year'],how='left')
    for col in ['minorNetworkMetres','benchmarkWeight','modelWeight','majorMiles','routedMajorMiles']:panel[col]=panel[col].fillna(0)
    for name in ['benchmark','model']:
        weight=name+'Weight';total=panel.groupby(['authority','year'])[weight].transform('sum');panel[name+'MinorMiles']=panel[weight]/total*panel.minorResidualMiles;panel[name+'Miles']=panel.majorMiles+panel[name+'MinorMiles'];panel.loc[~panel.denominatorAvailable,name+'Miles']=np.nan
    
    for flag in ['unsupportedFunction','sparseGroup','densityExtrapolation']:panel[flag+'MinorShare']=(panel[flag+'ModelWeight']/panel.modelWeight).fillna(0)
    panel['preferredMiles']=panel[chosen+'Miles'];panel['preferredMinorMiles']=panel[chosen+'MinorMiles']
    for delta in [-.2,-.1,.1,.2]:
        name=('minus' if delta<0 else 'plus')+str(int(abs(delta)*100))+'Miles';weight=panel[chosen+'Weight']*(1+delta*(panel.roadDecile.fillna(5.5)-5.5)/4.5);total=weight.groupby([panel.authority,panel.year]).transform('sum');panel[name]=panel.majorMiles+panel.minorResidualMiles*weight/total;panel.loc[~panel.denominatorAvailable,name]=np.nan
    inventory=pd.read_csv(OLD/'lsoa-year-inventory.csv.gz',low_memory=False);panel=panel.merge(inventory[['lsoa','year','assignedMotorCollisions','nearBoundaryMotorCollisions','unassignedMotorCollisionsInRecordedAuthority','authorityReportingFlags']],on=['lsoa','year'],validate='1:1')
    panel['majorShare']=panel.majorMiles/panel.preferredMiles;panel['band']=panel.roadDecile.fillna(0).astype(int).astype(str);panel['conditionalRatePer100k']=np.where(truth(panel.rateEligible)&panel.preferredMiles.gt(0),panel.assignedMotorCollisions/panel.preferredMiles*1e5,np.nan)
    panel['validatedFrequencyEligible']=False;panel['modelCohortRow']=truth(panel.rateEligible)&truth(panel.roadEligible)&panel.preferredMiles.gt(0)&panel.majorShare.notna();keep=panel.groupby('lsoa').modelCohortRow.all();panel['balancedModelCohort']=panel.lsoa.map(keep)
    panel['qualityFlags']=panel.apply(lambda q:';'.join(z for z in ['modelled_exposure','2026_network_retrospective','public_access_proxy','resident_settlement_proxy','unsupported_road_function' if q.unsupportedFunctionMinorShare>0 else '', 'sparse_traffic_stratum' if q.sparseGroupMinorShare>0 else '', 'density_extrapolation' if q.densityExtrapolationMinorShare>0 else '', 'major_fallback_over_5pct' if q.majorFallbackShare>.05 else '', 'frost_unavailable' if not str(q.roadEligible).lower()=='true' else '', 'no_positive_exposure' if not q.preferredMiles>0 else '', 'authority_rate_gate_failed' if not str(q.rateEligible).lower()=='true' else ''] if z),axis=1)
    # Reconcile every allocation scenario independently to the fixed official total.
    variants=['benchmarkMiles','modelMiles','preferredMiles','minus20Miles','minus10Miles','plus10Miles','plus20Miles'];recon=[]
    for name in variants:
        total=panel.groupby(['authority','year'])[name].sum(min_count=1);joined=auth.set_index(['authority','year']).join(total.rename('allocated'))
        for key,row in joined.iterrows():recon.append({'authority':key[0],'year':key[1],'scenario':name,'officialMiles':row.motorMiles,'allocatedMiles':row.allocated,'differenceMiles':row.allocated-row.motorMiles,'available':row.denominatorAvailable})
    recon=pd.DataFrame(recon);assert recon[recon.available].differenceMiles.abs().max()<.01;assert (panel[variants].dropna()>=0).all().all();assert len(panel)==13972*4
    auth['allocatedLsoas']=auth.apply(lambda q:int(panel[(panel.authority==q.authority)&(panel.year==q.year)].preferredMiles.gt(0).sum()),axis=1)
    csv(OUT/'authority-reconciliation.csv',recon);csv(OUT/'authority-exposure-audit.csv',auth[['authority','areaName','region','year','motorMiles','publishedMajorMiles','minorResidualMiles','fallbackMajorMiles','majorFallbackShare','unallocatedMajorMiles','denominatorAvailable','rateEligible','qualityReasons','allocatedLsoas']])
    num=panel.groupby(['authority','year']).agg(polygonAssignedCollisions=('assignedMotorCollisions','sum'),nearBoundaryCollisions=('nearBoundaryMotorCollisions','sum'),unassignedInRecordedAuthority=('unassignedMotorCollisionsInRecordedAuthority','max')).reset_index().merge(a.rename(columns={'areaId':'authority'})[['authority','year','definiteMotorCollisions']],on=['authority','year']);num['assignedMinusRecorded']=num.polygonAssignedCollisions-num.definiteMotorCollisions;csv(OUT/'numerator-reconciliation.csv',num)
    panel.to_parquet(C/'panel.parquet',index=False)
    cols=['lsoa','name','authority','authorityName','region','year','population','density','urbanFraction','settlement','roadFrostDays','roadDecile','roadEligible','roadCoverage','minorNetworkMetres','majorMiles','routedMajorMiles','preferredMinorMiles',*variants,'assignedMotorCollisions','nearBoundaryMotorCollisions','unassignedMotorCollisionsInRecordedAuthority','authorityReportingFlags','conditionalRatePer100k','majorFallbackShare','unsupportedFunctionMinorShare','sparseGroupMinorShare','densityExtrapolationMinorShare','denominatorAvailable','rateEligible','balancedModelCohort','validatedFrequencyEligible','qualityReasons','qualityFlags']
    csv(OUT/'lsoa-exposure-panel.csv.gz',panel[cols]);csv(OUT/'lsoa-index.csv',f[['li','lsoa','name','authority','authorityName','region']])
    comparisons=[];fits=[];predictions=[]
    for scenario in ['benchmarkMiles','modelMiles','minus20Miles','minus10Miles','plus10Miles','plus20Miles','without2021']:
        exposure=chosen+'Miles' if scenario=='without2021' else scenario;d=panel[panel.balancedModelCohort].copy();d['majorShare']=d.majorMiles/d[exposure];tr=d[d.year.le(2023)&(d.year.ge(2022) if scenario=='without2021' else True)];te=d[d.year.eq(2024)];p0,fit0=poisson_offset(tr,te,exposure,False);p1,fit1=poisson_offset(tr,te,exposure,True);comparisons.extend(comparison_rows(te,p0,p1,scenario,'all_roads'));fits.append({'scenario':scenario,'M0':fit0,'M1':fit1});predictions.append(te[['lsoa','authority','region','assignedMotorCollisions']].assign(scenario=scenario,expectedBaseline=p0,expectedFrost=p1));log(f'LSOA sensitivity {scenario}')
    csv(OUT/'frost-sensitivity.csv',pd.DataFrame(comparisons));csv(OUT/'lsoa-assessment-predictions.csv.gz',pd.concat(predictions));js(OUT/'outcome-models.json',fits)
    major_subset(args,panel)
    log('Allocation and sensitivity complete')

def major_subset(args,panel):
    C=args.cache;r=gpd.read_parquet(C/'roads.parquet');accepted=pd.read_parquet(C/'accepted-routes.parquet');sa=pd.read_parquet(C/'section-allocation.parquet');f=pd.read_parquet(C/'features.parquet')
    # Reject entire sections sharing any routed OS edge with another section.
    shared=accepted.groupby(['year','ri']).countPoint.nunique();bad_edges=set(shared[shared>1].index);bad={(q.year,q.countPoint) for q in accepted.itertuples() if (q.year,q.ri) in bad_edges};unique=accepted[~pd.MultiIndex.from_frame(accepted[['year','countPoint']]).isin(bad)]
    acceptedcp=set(zip(unique.year,unique.countPoint));section_lsoas=set(zip(sa.year,sa.countPoint,sa.li));exposure=sa[pd.MultiIndex.from_frame(sa[['year','countPoint']]).isin(acceptedcp)].groupby(['li','year']).routedMajorMiles.sum().rename('subsetMiles')
    assignments=pd.read_csv(OLD/'collision-assignments.csv.gz',usecols=['collisionId','year','primaryMotor','lsoa']);assignments=assignments[assignments.year.isin(YEARS)&truth(assignments.primaryMotor)&assignments.lsoa.isin(f.lsoa)];ids=set(assignments.collisionId)
    raw=ROOT/'public/stats19-evidence/06a03c0e5504175e23ed98e3e0b4767ac61d396c073405a1bddab9339c18f948.csv.gz';need=['collision_index','collision_year','location_easting_osgr','location_northing_osgr','first_road_class','first_road_number'];coll=pd.read_csv(raw,usecols=need,low_memory=False);coll=coll[coll.collision_index.isin(ids)&coll.first_road_class.isin([1,2,3])].merge(assignments[['collisionId','lsoa']],left_on='collision_index',right_on='collisionId');coll['roadNumber']=np.where(coll.first_road_class==1,'M','A')+coll.first_road_number.astype('Int64').astype(str);coll.loc[coll.first_road_class==2,'roadNumber']=coll.loc[coll.first_road_class==2,'roadNumber']+'(M)';coll['matchCount']=0;coll['li']=coll.lsoa.map(f.set_index('lsoa').li)
    for (year,num),g in unique.groupby(['year','roadNumber']):
        ix=coll.index[(coll.collision_year==year)&(coll.roadNumber==num)&coll.location_easting_osgr.notna()&coll.location_northing_osgr.notna()]
        if not len(ix):continue
        rr=g.drop_duplicates('ri');tree=STRtree(r.geometry.iloc[rr.ri].to_numpy());pairs=tree.query(shapely.points(coll.loc[ix,'location_easting_osgr'],coll.loc[ix,'location_northing_osgr']),predicate='dwithin',distance=30);matches=collections.defaultdict(set)
        for i,j in zip(*pairs):
            cp=int(rr.iloc[j].countPoint)
            if (year,cp,int(coll.loc[ix[i],'li'])) in section_lsoas:matches[ix[i]].add(cp)
        for i,cps in matches.items():coll.loc[i,'matchCount']=len(cps)
    csv(OUT/'major-collision-matching.csv.gz',coll[['collision_index','collision_year','lsoa','roadNumber','matchCount']]);counts=coll[coll.matchCount==1].groupby(['lsoa','collision_year']).size().rename('subsetCollisions').reset_index().rename(columns={'collision_year':'year'})
    sub=panel.drop(columns='assignedMotorCollisions').merge(exposure,on=['li','year'],how='left').merge(counts,on=['lsoa','year'],how='left');sub['assignedMotorCollisions']=sub.subsetCollisions.fillna(0);sub['subsetMiles']=sub.subsetMiles.fillna(0);sub['majorShare']=0.;sub['subsetEligible']=sub.balancedModelCohort&sub.subsetMiles.gt(0);keep=sub.groupby('lsoa').subsetEligible.all();sub=sub[sub.lsoa.map(keep)].copy();comparisons=[]
    if len(sub):
        tr=sub[sub.year<=2023];te=sub[sub.year==2024];p0,m0=poisson_offset(tr,te,'subsetMiles',False);p1,m1=poisson_offset(tr,te,'subsetMiles',True);comparisons=comparison_rows(te,p0,p1,'accepted_unique_routes','major_roads')
    csv(OUT/'major-subset-sensitivity.csv',pd.DataFrame(comparisons));csv(OUT/'major-subset-panel.csv.gz',sub[['lsoa','year','authority','region','subsetMiles','assignedMotorCollisions']]);js(OUT/'major-subset-audit.json',{'routedSectionYears':accepted[['year','countPoint']].drop_duplicates().shape[0],'overlappingSectionYearsExcluded':len(bad),'uniqueSectionYears':len(acceptedcp),'majorClassCandidateCollisions':len(coll),'matchedOnce':int(coll.matchCount.eq(1).sum()),'ambiguousExcluded':int(coll.matchCount.gt(1).sum()),'unmatchedExcluded':int(coll.matchCount.eq(0).sum()),'balancedLsoas':sub.lsoa.nunique(),'scope':'Routed, non-overlapping major sections with exact road number and within 30 m; numerator misses displaced/incorrectly coded locations; exploratory subset, not a certified matched exposure study'})

def package(args):
    C=args.cache;panel=pd.read_parquet(C/'panel.parquet');auth=pd.read_csv(OUT/'authority-exposure-audit.csv');ma=pd.read_parquet(C/'major-audit.parquet');validation=pd.read_csv(OUT/'traffic-validation.csv');modelinfo=json.load(open(OUT/'traffic-model.json'));selected=modelinfo['selected'];summary=[]
    for region in REGIONS:
        p=panel[(panel.region==region)&(panel.year==2024)];au=auth[(auth.region==region)&(auth.year==2024)];v=validation[(validation.model==selected)&validation['sample'].eq('2024_geographic_holdout')&validation.dimension.eq('region')&validation.group.eq(region)]
        stress=(p[['minus20Miles','minus10Miles','plus10Miles','plus20Miles']].div(p.preferredMiles,axis=0)-1).abs().max(axis=1)
        summary.append({'region':region,'lsoas':len(p),'positiveLsoas':int(p.preferredMiles.gt(0).sum()),'modelCohortLsoas':int(p.balancedModelCohort.sum()),'authorities':len(au),'billionMiles':float(au.motorMiles.sum()/1e9),'majorShare':float(au.publishedMajorMiles.sum()/au.motorMiles.sum()),'routedShareOfMajorMiles':float(1-au.fallbackMajorMiles.sum()/au.publishedMajorMiles.sum()),'ruralResidentLsoas':int(p.settlement.eq('rural_resident').sum()),'minorNetworkKm':float(p.minorNetworkMetres.sum()/1000),'assessment':v.replace({np.nan:None}).to_dict('records'),'p95StressChange':float(stress.quantile(.95)),'maxStressChange':float(stress.max())})
    checks={'authorityScenarioConservation':float(pd.read_csv(OUT/'authority-reconciliation.csv').differenceMiles.abs().max()),'allRows':len(panel),'targetRows':13972*4,'noDuplicateLsoaYears':not panel.duplicated(['lsoa','year']).any(),'validatedLsoaRows':0,'positiveExposureRows':int(panel.preferredMiles.gt(0).sum()),'missingExposureRows':int(panel.preferredMiles.isna().sum()),'zeroExposureRows':int(panel.preferredMiles.eq(0).sum()),'nonnegativeExposures':bool(panel.preferredMiles.dropna().ge(0).all()),'frozenPriorInputsUnchanged':True,'independentAllocationValidation':False,'independentReview':False}
    register=json.load(open(OUT/'source-register.json'))
    for name,h in register['frozenInputs'].items():assert sha(ROOT/name)==h
    summarydata={'version':'mileage-v1','date':'2026-10-11','status':'implemented_conditional_not_validated','selected':selected,'regions':summary,'geometry':json.load(open(C/'geometry-audit.json')),'checks':checks,'traffic':modelinfo,'comparison':pd.read_csv(OUT/'frost-sensitivity.csv').replace({np.nan:None}).to_dict('records'),'majorSubset':json.load(open(OUT/'major-subset-audit.json')),'majorComparison':pd.read_csv(OUT/'major-subset-sensitivity.csv').replace({np.nan:None}).to_dict('records'),'limits':['All mileage allocations are modelled; no LSOA is certified for validated frequency.','Count sites are not a representative census of all minor links.','Road access, historical changes and within-authority allocation lack independent validation.','Annual frost/weather exposure remains unavailable.','Air-frost deciles use the original wider mapped network, including restricted roads; no silent reweighting.']}
    js(OUT/'summary.json',summarydata);js(OUT/'completion-audit.json',checks)
    # Compact per-LSOA 2024 lookup for the interactive pilot explorer.
    cols=['lsoa','name','authorityName','region','preferredMiles','majorMiles','routedMajorMiles','preferredMinorMiles','benchmarkMiles','modelMiles','minus20Miles','plus20Miles','assignedMotorCollisions','conditionalRatePer100k','roadDecile','qualityFlags']
    js(OUT/'explorer.json',{'version':'mileage-v1','rows':panel[panel.year.eq(2024)][cols].round(6).replace({np.nan:None}).to_dict('records')})
    js(OUT/'runtime.json',{'python':platform.python_version(),'numpy':np.__version__,'pandas':pd.__version__,'shapely':shapely.__version__,'pyogrio':pyogrio.__version__,'networkx':nx.__version__,'scikitLearn':sklearn.__version__,'scipy':scipy.__version__,'geopandas':gpd.__version__,'pyproj':pyproj.__version__,'pyarrow':pyarrow.__version__,'seed':41106})
    import shutil
    shutil.copy2(ROOT/'scripts/build-frost-mileage.py',OUT/'build-frost-mileage.py');shutil.copy2(ROOT/'scripts/build-frost-features.py',OUT/'build-frost-features.py')
    js(OUT/'output-manifest.json',{'version':'mileage-v1','files':[{'file':p.name,'bytes':p.stat().st_size,'sha256':sha(p)} for p in sorted(OUT.iterdir()) if p.is_file() and p.name!='output-manifest.json']})
    log(json.dumps(checks));log('Pilot release packaged')

if __name__=='__main__':
    parser=argparse.ArgumentParser();parser.add_argument('--sources',type=Path,required=True);parser.add_argument('--cache',type=Path,required=True);parser.add_argument('--stage',choices=['prepare','major','model','outcomes','package','all'],default='all');args=parser.parse_args();OUT.mkdir(exist_ok=True,parents=True)
    for name in ['prepare','major','model','outcomes','package']:
        if args.stage in [name,'all']:globals()[name](args)