"""Reproduce the 2015/2025 ATUS reading-episode comparison.

Requires pandas/numpy and official Respondent, Activity, replicate-weight .dat
files in source/atus2015 and source/atus2025. Outputs aggregates only.
Reading is personal-interest reading (120312), not book/fiction-only.
"""
from pathlib import Path
import json
import numpy as np
import pandas as pd

ROOT = Path(__file__).resolve().parent

def estimate(y, w, rw):
    y=np.asarray(y,dtype=float); w=np.asarray(w,dtype=float)
    point=float(w@y/w.sum())
    reps=(rw*y[:,None]).sum(axis=0)/rw.sum(axis=0)
    se=float(np.sqrt((4/160)*np.square(reps-point).sum()))
    return {'estimate':point,'se':se,'ci95_low':point-1.96*se,'ci95_high':point+1.96*se}

def year_data(year):
    folder=ROOT/'source'/f'atus{year}'
    r=pd.read_csv(folder/f'atusresp_{year}.dat',dtype={'TUCASEID':str})
    a=pd.read_csv(folder/f'atusact_{year}.dat',dtype={'TUCASEID':str})
    w=pd.read_csv(folder/f'atuswgts_{year}.dat',dtype={'TUCASEID':str})
    assert r.TUCASEID.is_unique and w.TUCASEID.is_unique
    assert set(a.TUCASEID)==set(r.TUCASEID)
    assert a.TUACTDUR24.between(0,1440).all()
    assert (a.groupby('TUCASEID').TUACTDUR24.sum()==1440).all()
    reading=a.loc[a.TRCODE.eq(120312)]
    m=reading.groupby('TUCASEID').TUACTDUR24.agg(['sum','max','count'])
    d=r.merge(m,left_on='TUCASEID',right_index=True,how='left',validate='one_to_one')
    d[['sum','max','count']]=d[['sum','max','count']].fillna(0)
    reps=[f'FINLWGT{i:03d}' for i in range(1,161)]
    d=d.merge(w[['TUCASEID']+reps],on='TUCASEID',validate='one_to_one')
    summary=pd.read_csv(folder/f'atussum_{year}.dat',dtype={'TUCASEID':str})
    assert summary.TUCASEID.is_unique
    d=d.merge(summary[['TUCASEID','TEAGE','TESEX','t120312']],on='TUCASEID',validate='one_to_one')
    assert np.array_equal(d['sum'].to_numpy(),d.t120312.to_numpy())
    assert d[reps].notna().all().all() and d.TUFINLWGT.gt(0).all()
    d['any_reading']=d['sum'].gt(0).astype(float)
    for threshold in [15,30,60]:
        d[f'episode_{threshold}']=d['max'].ge(threshold).astype(float)
    return d,reps,len(a),len(reading)

def run():
    rows=[];audit=[]
    for year in [2015,2025]:
        d,reps,nactivities,nreading=year_data(year)
        groups={'all_15plus':np.ones(len(d),dtype=bool)}
        # Age is on the Respondent file in 2025 and 2015 as TEAGE.
        if 'TEAGE' in d:
            for lo,hi in [(15,34),(35,54),(55,200)]:
                groups[f'age_{lo}_{hi if hi<200 else "plus"}']=d.TEAGE.between(lo,hi)
        for code,label in [(1,'men'),(2,'women')]:
            groups[label]=d.TESEX.eq(code)
        # Only broad descriptive groups; no claim employment causes gaps.
        groups['employed']=d.TELFS.isin([1,2])
        groups['not_employed']=~d.TELFS.isin([1,2])
        for group,mask in groups.items():
            g=d.loc[mask];rw=g[reps].to_numpy(dtype=float)
            for metric in ['any_reading','episode_15','episode_30','episode_60','sum']:
                s=estimate(g[metric],g.TUFINLWGT,rw)
                scale=1 if metric=='sum' else 100
                s={k:float(v*scale) for k,v in s.items()}
                rows.append({'year':year,'group':group,'metric':metric,'unit':'minutes' if metric=='sum' else 'percent','sample_n':len(g),'event_n':int(g['any_reading' if metric=='sum' else metric].gt(0).sum()),**s})
            gr=g.loc[g.any_reading.eq(1)]
            s=estimate(gr['sum'],gr.TUFINLWGT,gr[reps].to_numpy(dtype=float))
            rows.append({'year':year,'group':group,'metric':'minutes_among_readers','unit':'minutes','sample_n':len(gr),'event_n':len(gr),**s})
        audit.append({'year':year,'respondents':len(d),'activities':nactivities,'reading_activities':nreading,'reading_respondents':int(d.any_reading.sum()),'all_diaries_1440_minutes':True,'reading_code':120312,'unique_respondents':True,'replicate_count':len(reps)})
    result=pd.DataFrame(rows)
    result.to_csv(ROOT/'results.csv',index=False,float_format='%.8f')
    comparisons=[]
    for (group,metric),g in result.groupby(['group','metric']):
        a=g.loc[g.year.eq(2015)].iloc[0];b=g.loc[g.year.eq(2025)].iloc[0]
        delta=float(b.estimate-a.estimate);se=float(np.sqrt(a.se*a.se+b.se*b.se))
        comparisons.append({'group':group,'metric':metric,'unit':'minutes' if metric in ['sum','minutes_among_readers'] else 'percentage_points','change_2025_minus_2015':delta,'se':se,'ci95_low':delta-1.96*se,'ci95_high':delta+1.96*se})
    pd.DataFrame(comparisons).to_csv(ROOT/'comparisons.csv',index=False,float_format='%.8f')
    (ROOT/'results.json').write_text(json.dumps({'audit':audit,'estimates':rows,'comparisons':comparisons},indent=2)+'\n',encoding='utf-8')
    print(result[result.group.eq('all_15plus')].to_string(index=False))
    print(pd.DataFrame(comparisons).query("group == 'all_15plus'").to_string(index=False))

if __name__=='__main__': run()