"""Poll-only proxy for 'late news': does movement between early (15-28d out) and late (0-14d out) poll averages predict the residual (actual - late avg)?
Run: python -I late_shift.py"""
import numpy as np, pandas as pd
d=pd.read_excel('data/sb_rawpolls_021325.xlsx'); d=d[d.type_simple.isin(['Pres-G','Sen-G','Gov-G','House-G'])&(d.cand1_party=='DEM')&(d.cand2_party=='REP')&d.partisan.isna()]
d['days']=(d.electiondate-d.polldate).dt.days; d['cyc']=d.electiondate.dt.year; d=d[(d.days>=0)&(d.days<=28)&(d.cyc%2==0)]
d['per']=np.where(d.days<=14,'late','early')
g=d.groupby(['race','cyc','per','pollster']).margin_poll.mean().reset_index()
a=g.groupby(['race','cyc','per']).agg(m=('margin_poll','mean'),k=('pollster','nunique')).reset_index().pivot(index=['race','cyc'],columns='per',values=['m','k']).dropna()
a.columns=['early','late','ke','kl']; a=a.reset_index(); act=d.groupby('race').margin_actual.first(); a['act']=a.race.map(act)
a=a[(a.ke>=1)&(a.kl>=1)]; a['shift']=a.late-a.early; a['resid']=a.act-a.late
print('races',len(a),'cycles',a.cyc.nunique()); print('shift sd',a['shift'].std().round(2))
def fit(x):
    X=np.c_[np.ones(len(x)),x['shift']]; b=np.linalg.lstsq(X,x.resid,rcond=None)[0]; return b
b=fit(a); print('pooled slope of resid on shift: %.3f (intercept %.2f)'%(b[1],b[0]))
rng=np.random.default_rng(3); cy=a.cyc.unique(); bs=[]
for _ in range(5000):
    s=rng.choice(cy,len(cy)); x=pd.concat([a[a.cyc==c] for c in s]); bs.append(fit(x)[1])
print('cluster-bootstrap 95%% CI for slope: [%.3f, %.3f]'%tuple(np.percentile(bs,[2.5,97.5])))
for c in sorted(cy):
    x=a[a.cyc==c]; print(c,len(x),'slope %.2f'%fit(x)[1] if len(x)>5 else '')
# LOCO MAE: late avg vs late avg + b*shift
e0=[];e1=[]
for c in cy:
    tr=a[a.cyc!=c]; te=a[a.cyc==c]; bb=fit(tr); e0+=list((te.late-te.act).abs()); e1+=list((te.late+bb[1]*te['shift']-te.act).abs())
print('LOCO MAE late-avg %.3f  late-avg+momentum %.3f'%(np.mean(e0),np.mean(e1)))