from pathlib import Path
import pandas as pd,numpy as np,zipfile,json,io
p=Path(__file__).parent
cols=['QSTNUM','OCCURNUM','PRTAGE','PEEDUCA','PEC1Q14A','PEC1Q15A','PWSUPWGT']
spans=[(814,819),(819,821),(121,123),(136,138),(1044,1046),(1049,1051),(1483,1493)]
def load(y):
 if y==2017:
  with zipfile.ZipFile(p/'jul17pub.zip') as z:
   rows=[[int(l[a:b]) for a,b in spans] for l in z.read(z.namelist()[0]).splitlines()]
  d=pd.DataFrame(rows,columns=cols)
  with zipfile.ZipFile(p/'jul17rep.zip') as z:
   rows=[[int(l[:5]),int(l[5:7])]+[int(l[7+10*i:17+10*i]) for i in range(161)] for l in z.read(z.namelist()[0]).splitlines()]
  r=pd.DataFrame(rows,columns=cols[:2]+['R'+str(i) for i in range(161)])
 else:
  d=pd.read_csv(p/'jul22pub.csv',usecols=cols)
  r=pd.read_csv(p/'jul22rep.csv').rename(columns={'PWSUPWGT'+str(i):'R'+str(i) for i in range(161)})
 d=d[(d.PRTAGE>=18)&(d.PWSUPWGT>0)]
 d=d.merge(r,on=cols[:2],validate='one_to_one',how='left')
 assert d.R0.notna().all()
 print(y,len(d),'base/rep0 ratio',np.unique(np.round(d.PWSUPWGT/d.R0,6)))
 return d
out=[];audits={};frames={}
for y in [2017,2022]:
 d=load(y)
 audits[y]={'supplement_n':len(d),'book_fiction_response_counts':{f'{a},{b}':int(n) for (a,b),n in d.groupby(['PEC1Q14A','PEC1Q15A']).size().items()}}
 d['agegroup']=pd.cut(d.PRTAGE,[17,24,34,44,54,64,74,100],labels=['18-24','25-34','35-44','45-54','55-64','65-74','75+'])
 d['education']=pd.cut(d.PEEDUCA,[30,38,39,42,43,46],labels=['Less than HS','HS graduate','Some college/associate','Bachelor','Graduate'])
 frames[y]=d
 for dim in ['all','agegroup','education']:
  groups=[('All',d)] if dim=='all' else d.groupby(dim,observed=True)
  for label,g in groups:
   valid=g.PEC1Q14A.isin([1,2])&g.PEC1Q15A.isin([1,2])
   measures={
   'fiction_among_book_readers':(valid&(g.PEC1Q14A==1),g.PEC1Q15A==1),
   'fiction_among_nonbook_readers':(valid&(g.PEC1Q14A==2),g.PEC1Q15A==1),
   'both_among_valid':(valid,(g.PEC1Q14A==1)&(g.PEC1Q15A==1)),
   'books_among_valid':(valid,g.PEC1Q14A==1),
   'fiction_among_valid':(valid,g.PEC1Q15A==1)}
   for metric,(den,num) in measures.items():
    z=g[den];w=z.PWSUPWGT.to_numpy();a=num.loc[z.index].to_numpy();reps=z[['R'+str(i) for i in range(1,161)]].to_numpy()
    est=100*(w*a).sum()/w.sum();re=100*(reps*a[:,None]).sum(axis=0)/reps.sum(axis=0);se=np.sqrt(4/160*((re-est)**2).sum())
    out.append(dict(year=y,dimension=dim,group=str(label),metric=metric,n_den=len(z),n_num=int(a.sum()),pct=est,se_pp=se,ci95_low=est-1.96*se,ci95_high=est+1.96*se))
pd.DataFrame(out).to_csv(p/'results.csv',index=False)
(p/'audit.json').write_text(json.dumps(audits,indent=2))
f=pd.DataFrame(out);print(f[(f.dimension=='all')].to_string(index=False))
# Cross-wave independent samples, no overlapping respondents five years apart.
a=f[f.year==2017].merge(f[f.year==2022],on=['dimension','group','metric'],suffixes=('_17','_22'))
a['change_pp']=a.pct_22-a.pct_17;a['change_se_pp']=np.sqrt(a.se_pp_17**2+a.se_pp_22**2);a['change_ci95_low']=a.change_pp-1.96*a.change_se_pp;a['change_ci95_high']=a.change_pp+1.96*a.change_se_pp
a.to_csv(p/'changes.csv',index=False)

# Missing-fiction bounds among all observed book readers (no missing=No recode).
bounds=[]
for y,d in frames.items():
 g=d[d.PEC1Q14A==1]; w=g.PWSUPWGT; known=g.PEC1Q15A.isin([1,2]); yes=g.PEC1Q15A==1
 bounds.append(dict(year=y,book_yes_n=len(g),missing_fiction_n=int((~known).sum()),missing_fiction_weighted_pct=100*w[~known].sum()/w.sum(),lower_pct=100*w[yes].sum()/w.sum(),upper_pct=100*(w[yes].sum()+w[~known].sum())/w.sum()))
pd.DataFrame(bounds).to_csv(p/'missingness_bounds.csv',index=False)
# Direct standardization, separately by age and education, to the 2017
# complete-case book-reader distribution. Reference shares fixed, conditional inference.
std=[]
for dim in ['agegroup','education']:
 ref=frames[2017];ref=ref[(ref.PEC1Q14A==1)&ref.PEC1Q15A.isin([1,2])]
 shares=ref.groupby(dim,observed=True).PWSUPWGT.sum();shares/=shares.sum()
 for y,d in frames.items():
  theta=0.;rr=np.zeros(160)
  for label,share in shares.items():
   g=d[(d[dim]==label)&(d.PEC1Q14A==1)&d.PEC1Q15A.isin([1,2])];a=(g.PEC1Q15A==1).to_numpy();w=g.PWSUPWGT.to_numpy();rw=g[['R'+str(i) for i in range(1,161)]].to_numpy()
   theta+=share*100*(w*a).sum()/w.sum();rr+=share*100*(rw*a[:,None]).sum(axis=0)/rw.sum(axis=0)
  se=np.sqrt(4/160*((rr-theta)**2).sum())
  std.append(dict(year=y,standardize_by=dim,pct=theta,se_pp_reference_fixed=se,reference='2017 complete-case book readers; one dimension at a time'))
pd.DataFrame(std).to_csv(p/'standardized.csv',index=False)