"""Reproduce 2024 ACS PUMS mean household income less housing costs.
Run explicitly with numpy/pandas installed. Downloads are pinned by the first
recorded checksum and cached outside the public site. Never runs during build.
"""
import concurrent.futures,hashlib,io,json,pathlib,urllib.request,zipfile
import numpy as np
import pandas as pd
ROOT=pathlib.Path(__file__).resolve().parents[1]
CACHE=ROOT/'.research-cache/pums-2024';CACHE.mkdir(parents=True,exist_ok=True)
STATES=[s for s in json.loads((ROOT/'src/data/registry.json').read_text())['states'] if s['id']!='72']
COLS=['STATE','TYPEHUGQ','NP','TEN','HINCP','GRNTP','SMOCP','ADJINC','ADJHSG','WGTP']+[f'WGTP{i}' for i in range(1,81)]

def weighted_quantile(values,weights,q=.5):
    order=np.argsort(values,kind='stable');v=values[order];w=weights[order]
    assert np.isfinite(v).all() and np.isfinite(w).all() and (w>=0).all() and w.sum()>0
    return float(v[np.searchsorted(np.cumsum(w),w.sum()*q,side='left')])

def summarize(values,weights,reps):
    order=np.argsort(values,kind='stable');v=values[order];w=weights[order];rw=reps[order]
    # Retain signed replicate weights, as required by Census. Ratios of weighted
    # totals are well-defined even when a replicate contains negative weights.
    estimate=float(np.dot(v,w)/w.sum())
    totals=rw.sum(axis=0);assert (totals>0).all()
    rm=(v[:,None]*rw).sum(axis=0)/totals
    moe=float(1.645*np.sqrt(4/80*np.square(rm-estimate).sum()))
    assert moe>0 and len(v)>=100, 'Do not publish unreliable or zero-SE estimates'
    direct=sum(value*weight for value,weight in zip(values.tolist(),weights.tolist()))/sum(weights.tolist())
    assert abs(direct-estimate)<1e-6
    return dict(value=estimate,moe=moe,sample=int(len(v)),weightedHouseholds=int(w.sum()),median=weighted_quantile(v,w),p25=weighted_quantile(v,w,.25),p75=weighted_quantile(v,w,.75),negativeShare=float(100*w[v<0].sum()/w.sum()))

def one(state):
    url=f"https://www2.census.gov/programs-surveys/acs/data/pums/2024/1-Year/csv_h{state['abbr'].lower()}.zip"
    path=CACHE/(state['id']+'.zip')
    if not path.exists():path.write_bytes(urllib.request.urlopen(url,timeout=120).read())
    data=path.read_bytes();sha=hashlib.sha256(data).hexdigest();pin=path.with_suffix('.sha256')
    published=ROOT/'research/pums-results.json'
    if published.exists():assert json.loads(published.read_text())['states'][state['id']]['sha256']==sha, 'Producer file changed since reviewed edition'
    if pin.exists():assert pin.read_text()==sha
    else:pin.write_text(sha)
    z=zipfile.ZipFile(io.BytesIO(data));name=next(n for n in z.namelist() if n.endswith('.csv'))
    d=pd.read_csv(z.open(name),usecols=COLS)
    assert set(d['STATE'].unique())=={int(state['id'])}
    occupied=(d.TYPEHUGQ==1)&(d.NP>0)&(d.WGTP>0)
    out={}
    for ten,label in [(3,'renters'),(1,'mortgaged-owners'),(2,'mortgage-free-owners')]:
        x=d[occupied&(d.TEN==ten)].copy();cost='GRNTP' if ten==3 else 'SMOCP'
        assert x[['HINCP',cost,'ADJINC','ADJHSG']].notna().all().all()
        value=x.HINCP.to_numpy()*x.ADJINC.to_numpy()/1e6-12*x[cost].to_numpy()*x.ADJHSG.to_numpy()/1e6
        out[label]=summarize(value,x.WGTP.to_numpy(),x[[f'WGTP{i}' for i in range(1,81)]].to_numpy())
    print(state['name'],{k:round(v['value']) for k,v in out.items()},flush=True)
    return state['id'],dict(name=state['name'],url=url,sha256=sha,occupiedHouseholds=int(d.loc[occupied,'WGTP'].sum()),groups=out)

if __name__=='__main__':
    with concurrent.futures.ThreadPoolExecutor(max_workers=4) as pool:results=dict(pool.map(one,STATES))
    output=dict(year=2024,method='Weighted mean of adjusted household income minus twelve times adjusted monthly housing costs. 90% MOE = 1.645 sqrt((4/80) sum((replicate mean - full mean)^2)). Signed replicate weights are retained.',exclusions='Vacant units, group quarters, no-cash-rent households. Zero and negative income retained. No taxes or other living expenses deducted. Respondent and Census rounding, allocation and top/bottom coding remain in the source.',quantile='Supporting percentiles are point estimates without confidence intervals. Smallest observed value whose cumulative nonnegative household weight reaches the requested share; no interpolation.',national=None,states=results)
    (ROOT/'research/pums-results.json').write_text(json.dumps(output,indent=2)+'\n')