#!/usr/bin/env python3
"""Freeze Phase 4 source joins. Usage: python scripts/prepare-parking-phase4.py SOURCE_DIR.
Raw downloads are acquired separately; this script never downloads or fits outcomes.
"""
import collections,csv,gzip,hashlib,io,json,pathlib,sys,zipfile,math
import openpyxl,numpy as np,shapefile,shapely
from shapely import STRtree,points
from shapely.geometry import shape
from shapely.ops import transform
from pyproj import Transformer
from geo_ods import ods
ROOT=pathlib.Path(__file__).resolve().parents[1];OUT=ROOT/'public/research/parking/v4';YEARS=list(range(2021,2026))
def read(p):return json.loads(p.read_text())
def sha(p):return hashlib.sha256(p.read_bytes()).hexdigest()
def dump(p,d):p.write_text(json.dumps(d,separators=(',',':'),ensure_ascii=False,allow_nan=False)+'\n')
def csvs(b):return list(csv.DictReader(io.StringIO(b.decode('utf-8-sig'))))
def unique(rr,key):
 out={}
 for r in rr:
  k=key(r);assert k not in out,('Duplicate key',k);out[k]=r
 return out
def unpack(s):
 p=ROOT/'public'/s['archivePath'].lstrip('/');b=p.read_bytes()
 if s.get('archiveSha256'):assert hashlib.sha256(b).hexdigest()==s['archiveSha256']
 if p.suffix=='.gz':b=gzip.decompress(b)
 assert hashlib.sha256(b).hexdigest()==s['sha256']
 return b

def prepare(raw):
 OUT.mkdir(parents=True,exist_ok=True);(OUT/'inputs').mkdir(exist_ok=True)
 assert not (OUT/'design.json').exists(),'Frozen design already exists; use a new analysis version for changes.'
 new_sources=[]
 def archive(name,url,title,period,compress=False):
  src=raw/name;target=OUT/'inputs'/(name+'.gz' if compress else name);b=src.read_bytes();target.write_bytes(gzip.compress(b,mtime=0) if compress else b)
  s=dict(title=title,url=url,period=period,path='/'+str(target.relative_to(ROOT/'public')),sha256=sha(target),rawSha256=hashlib.sha256(b).hexdigest(),bytes=target.stat().st_size);new_sources.append(s);return target
 ruc=archive('ruc2021.xlsx','https://www.ons.gov.uk/file?uri=%2Fmethodology%2Fgeography%2Fgeographicalproducts%2Fruralurbanclassifications%2F2021ruralurbanclassification%2F260331rucallsupplementarytables.xlsx','ONS Rural Urban Classification supplementary tables','2021 classification; March 2026 release')
 lu=archive('oa21-lsoa21.csv','https://open-geography-portalx-ons.hub.arcgis.com/api/download/v1/items/821c4e07c07042e0bee806dfec5a1f77/csv?layers=0','ONS OA21 → LSOA21 lookup','April 2023 hierarchy',True)
 nt=archive('tra8903.ods','https://assets.publishing.service.gov.uk/media/6a0b8331fcae986635db9146/tra8903-miles-by-local-authority-excluding-trunk-roads.ods','DfT TRA8903 non-trunk-road mileage','2021–2025')
 vt=archive('tra8902.ods','https://assets.publishing.service.gov.uk/media/6a0b831bc75cc34a8ff8f489/tra8902-miles-by-local-authority-and-selected-vehicle-type.ods','DfT TRA8902 vehicle-class mileage','2021–2025')
 boundary=archive('authority-boundaries.geojson',(raw/'boundary-url.txt').read_text(),'ONS counties and unitary authority boundaries','December 2025',True)
 social=read(ROOT/'db/research/social-factors-panel.json');sm=read(ROOT/'db/research/social-factors-inputs.json');ss={s['id']:s for s in sm['sources']};geo=read(ROOT/'db/geo/initial.json');v3=read(ROOT/'public/research/parking/v3/analysis.json');v3panel=read(ROOT/'public/research/parking/v3/authority-year-panel.json');v1=read(ROOT/'public/research/parking/v1/audit.json')
 context=dict(features=social['features'],rows=[r for r in social['rows'] if r['year'] in YEARS],sources=social['sources'],originalPanelSha256=sha(ROOT/'db/research/social-factors-panel.json'))
 dump(OUT/'inputs/social-context.json',context)
 popz=zipfile.ZipFile(io.BytesIO(unpack(ss['ts001.zip'])));pop=unique(csvs(popz.read('census2021-ts001-oa.csv')),lambda r:r['geography code']);pk=next(k for k in next(iter(pop.values())) if k.startswith('Residence type: Total'))
 oa_lookup=unique(csvs(gzip.decompress(lu.read_bytes())),lambda r:r['OA21CD']);lookup=unique(csvs(unpack(ss['lsoa21-utla23.csv'])),lambda r:r['LSOA21CD'])
 w=openpyxl.load_workbook(ruc,read_only=True,data_only=True);agg=collections.defaultdict(collections.Counter);oa_seen=set();oa_rows=[]
 for r in w['Table 1A'].iter_rows(min_row=4,values_only=True):
  oa=r[0]
  if not str(oa).startswith('E'):continue
  assert oa not in oa_seen;oa_seen.add(oa);assert oa in oa_lookup and oa in pop,oa
  lsoa=oa_lookup[oa]['LSOA21CD'];assert lsoa in lookup;l=lookup[lsoa];a=l['UTLA23CD'];n=int(pop[oa][pk]);category=r[4];assert category in ['Urban','Larger rural','Smaller rural'];agg[a][category]+=n;agg[a]['total']+=n;agg[a]['oas']+=1
  oa_rows.append(dict(oa=oa,lsoa=lsoa,areaId=a,population=n,settlement=category))
 assert len(oa_seen)==len([k for k in pop if k.startswith('E')])
 with gzip.open(OUT/'inputs/oa-population-rurality.csv.gz','wt',newline='') as f:
  wr=csv.DictWriter(f,fieldnames=list(oa_rows[0]));wr.writeheader();wr.writerows(oa_rows)
 # Household feature shares use each table's own published total (disclosure perturbation retained).
 households=collections.defaultdict(lambda:dict(housing=np.zeros(8),cars=np.zeros(4)))
 for table,key in [('ts044','housing'),('ts045','cars')]:
  z=zipfile.ZipFile(ROOT/f'public/research/parking/v2/inputs/census2021-{table}.zip');rr=csvs(z.read(f'census2021-{table}-lsoa.csv'));fields=list(rr[0])
  for r in rr:
   lsoa=r['geography code']
   if not lsoa.startswith('E'):continue
   vals=np.array([int(r[k]) for k in fields[4:]]);assert sum(vals)==int(r[fields[3]]);households[lookup[lsoa]['UTLA23CD']][key]+=vals
 # New network context is a 2026 snapshot and therefore only a retrospective sensitivity.
 z=zipfile.ZipFile(raw/'os-open-roads.zip');assert hashlib.md5((raw/'os-open-roads.zip').read_bytes()).hexdigest()=='8c6e66a255f79e3e31845307cefee298'
 polygons=json.loads(gzip.decompress(boundary.read_bytes()))['features'];tx=Transformer.from_crs(4326,27700,always_xy=True).transform
 polygons=sorted(polygons,key=lambda f:f['properties']['CTYUA25CD']);ids=[f['properties']['CTYUA25CD'] for f in polygons];geoms=[transform(tx,shape(f['geometry'])) for f in polygons];tree=STRtree(geoms)
 nodes=collections.Counter();nodeforms=collections.Counter();node_ids={};node_rows=[];unmatched=ambiguous=duplicate_nodes=0
 for member in sorted(n for n in z.namelist() if n.endswith('_RoadNode.shp')):
  b=member[:-4];reader=shapefile.Reader(shp=io.BytesIO(z.read(member)),dbf=io.BytesIO(z.read(b+'.dbf')),shx=io.BytesIO(z.read(b+'.shx')))
  batch=[]
  for sr in reader.iterShapeRecords():
   nid,form=sr.record;nodeforms[form]+=1
   if form not in ['junction','roundabout']:continue
   x,y=sr.shape.points[0]
   if nid in node_ids:
    assert node_ids[nid]==(form,x,y),('Conflicting OS node',nid);duplicate_nodes+=1;continue
   node_ids[nid]=(form,x,y);batch.append((nid,form,x,y))
  if not batch:continue
  xy=np.array([[r[2],r[3]] for r in batch]);ix=tree.query(points(xy),predicate='intersects');matches=collections.defaultdict(list)
  for p,g in zip(*ix):matches[int(p)].append(int(g))
  for i,(nid,form,x,y) in enumerate(batch):
   candidates=matches.get(i,[])
   if len(candidates)!=1:unmatched+=len(candidates)==0;ambiguous+=len(candidates)>1;continue
   a=ids[candidates[0]];nodes[a]+=1;node_rows.append(dict(id=nid,form=form,easting=x,northing=y,areaId=a))
 with gzip.open(OUT/'inputs/assigned-junctions.csv.gz','wt',newline='') as f:
  wr=csv.DictWriter(f,fieldnames=list(node_rows[0]));wr.writeheader();wr.writerows(node_rows)
 network=dict(sourceUrl=json.loads((raw/'os-downloads.json').read_text())[0]['url'],rawSha256=sha(raw/'os-open-roads.zip'),rawMd5='8c6e66a255f79e3e31845307cefee298',sourceBytes=(raw/'os-open-roads.zip').stat().st_size,version='April 2026',allGbNodeForms=dict(nodeforms),matchedEnglishJunctions=len(node_rows),duplicateTileNodes=duplicate_nodes,unmatchedOutsideEnglandOrBoundary=unmatched,ambiguousOmitted=ambiguous,definition='Only OS junction and roundabout nodes, one unambiguous polygon intersection; no nearest-area assignment. Snapshot follows outcomes, so network effect is a retrospective sensitivity, not as-of forecasting.')
 dump(OUT/'inputs/network-source.json',network)
 features=[];fmap={r['areaId']:r for r in social['features']};areas={a['id']:a for a in social['areas']}
 cohort3={r['areaId'] for r in v3panel if r['included']}
 for a in sorted(cohort3):
  f=fmap[a];h=households[a];v=agg[a];hs=h['housing']/sum(h['housing']);cs=h['cars']/sum(h['cars']);density=f['populationDensity2021'];population=f['population2021'];area=population/density if population and density else None;urban=v['Urban']/v['total'];ruralgroup='urban' if urban>=.8 else 'rural' if urban<.5 else 'mixed'
  features.append(dict(areaId=a,name=areas[a]['name'],region=areas[a]['region'],urbanShare=urban,largerRuralShare=v['Larger rural']/v['total'],smallerRuralShare=v['Smaller rural']/v['total'],rucPopulation=v['total'],oas=v['oas'],ruralGroup=ruralgroup,logDensity=math.log(density) if density else None,ageYoung=f['ageUnder25Share2021'],ageOld=f['age65Plus2021']/f['ageTotal2021'] if f['ageTotal2021'] else None,income=f['deprivation2019']['income'] if f['deprivation2019'] else None,employment=f['deprivation2019']['employment'] if f['deprivation2019'] else None,noQualifications=f['noQualificationsShare'],level4=f['level4PlusShare'],detached=float(hs[0]),terraced=float(hs[2]),flats=float(sum(hs[3:7])),noCars=float(cs[0]),twoPlusCars=float(sum(cs[2:])),junctions=nodes[a],junctionsPerKm2=nodes[a]/area if area else None,logJunctionDensity=math.log(1+nodes[a]/area) if area else None,areaKm2=area))
 # Official yearly miles: tables use MILLION vehicle miles; exact authority/year match.
 traffic={};t=ods(nt)['TRA8903'];head=t[4]
 for r in t[5:]:
  if not r[0].startswith('E'):continue
  for y in YEARS:
   i=next(i for i,k in enumerate(head) if k[:4]==str(y));val=r[i] if len(r)>i else ''
   try:n=round(float(val)*1e6)
   except ValueError:n=None
   key=(r[0],y);assert key not in traffic;traffic[key]=dict(nonTrunkMiles=n)
 t=ods(vt)['TRA8902'];head=t[4];seen=set()
 for r in t[5:]:
  if not r[1].startswith('E'):continue
  for y in YEARS:
   i=next(i for i,k in enumerate(head) if k[:4]==str(y));val=r[i] if len(r)>i else ''
   try:n=round(float(val)*1e6)
   except ValueError:n=None
   key=(r[1],y,r[0]);assert key not in seen;seen.add(key);traffic.setdefault((r[1],y),{})[{'Cars and Taxis':'carMiles','Light Commercial Vehicles':'vanMiles','Heavy Goods Vehicles':'hgvMiles'}[r[0]]]=n
 # Recreate canonical parked states, then cross-tab outcomes with trunk flags.
 stats=read(ROOT/'db/stats19/initial.json');src={s['id']:s for s in stats['sources']};states={};known={'1','3','4','5','6','7','8','9','10','11','12','13','14','15','19','20'}
 with gzip.open(ROOT/'public'/src['vehicle']['archivePath'].lstrip('/'),'rt',encoding='utf-8-sig') as f:
  for r in csv.DictReader(f):
   k=r['collision_index'];s='parked' if r['vehicle_manoeuvre']=='2' else 'other' if r['vehicle_manoeuvre'] in known else 'unknown';prev=states.get(k)
   states[k]='parked' if 'parked' in [prev,s] else 'unknown' if 'unknown' in [prev,s] else 'other'
 counts=collections.defaultdict(collections.Counter);keys=set()
 with gzip.open(ROOT/'public'/src['collision']['archivePath'].lstrip('/'),'rt',encoding='utf-8-sig') as f:
  for r in csv.DictReader(f):
   k=r['collision_index'];assert k not in keys;keys.add(k);a=r['local_authority_highway_current'];y=int(r['collision_year'])
   if a not in cohort3:continue
   s=states.get(k,'unknown');trunk={'1':'trunk','2':'nonTrunk'}.get(r['trunk_road_flag'],'unknownTrunk');counts[a,y]['all']+=1;counts[a,y][s]+=1;counts[a,y][trunk]+=1;counts[a,y][trunk+'_'+s]+=1
 out=[];all_social={(r['areaId'],r['year']):r for r in social['rows']};fm={f['areaId']:f for f in features}
 for r in v3panel:
  if not r['included']:continue
  a,y=r['areaId'],r['year'];c=counts[a,y];t=traffic.get((a,y),{});s=all_social[a,y];m=r['motorMiles'];n=t.get('nonTrunkMiles');assert c['all']==r['collisions'] and c['parked']==r['parkedCollisions'] and c['unknown']==r['unknownCollisions']
  row={**r,**fm[a],**t,**{'nonTrunkCount':c['nonTrunk'],'nonTrunkParked':c['nonTrunk_parked'],'nonTrunkOther':c['nonTrunk_other'],'nonTrunkUnknownParking':c['nonTrunk_unknown'],'unknownTrunkCount':c['unknownTrunk']}}
  row.update(trunkShare=1-n/m if n and 0<n<=m else None,carShare=t.get('carMiles')/m if t.get('carMiles') else None,vanShare=t.get('vanMiles')/m if t.get('vanMiles') else None,hgvShare=t.get('hgvMiles')/m if t.get('hgvMiles') else None,roadKm=s['roadLengthKm'],logFlow=math.log(m/s['roadLengthKm']) if s['roadLengthKm'] else None)
  if t.get('carMiles') is not None:assert abs(t['carMiles']-s['carsAndTaxisMiles'])<=50001,(a,y,'car exposure reconciliation')
  out.append(row)
 dump(OUT/'features.json',features);dump(OUT/'panel.json',out);dump(OUT/'source-register.json',dict(createdAt='2026-10-10',newSources=new_sources,existingSources=[s for s in social['sources'] if s['id'] in ['imd2019.raw','ts001.zip','ts006.zip','ts007a.zip','ts067.zip','lsoa21-utla23.csv']],network=network,oaCount=len(oa_seen),oaPopulation=sum(r['population'] for r in oa_rows)))
 paths=[OUT/'features.json',OUT/'panel.json',OUT/'source-register.json',OUT/'inputs/social-context.json',ROOT/'public/research/parking/v3/analysis.json',ROOT/'public/research/parking/v3/authority-decile-membership.csv']
 paths+=list((OUT/'inputs').iterdir())
 dump(OUT/'input-manifest.json',dict(version='v4',files=[dict(path=str(p.relative_to(ROOT)),sha256=sha(p),bytes=p.stat().st_size) for p in sorted(set(paths))],preparationScriptSha256=sha(pathlib.Path(__file__))))
 print(json.dumps(dict(authorities=len(features),rows=len(out),ruralGroups=dict(collections.Counter(f['ruralGroup'] for f in features)),unknownTrunkAuthorityYears=sum(r['unknownTrunkCount']>0 for r in out),missingTrunkShare=sum(r['trunkShare'] is None for r in out),network=network),indent=2))
if __name__=='__main__':prepare(pathlib.Path(sys.argv[1]))