"""Vugol Tensor v0 backtest. Leave-one-cycle-out on Silver Bulletin race-level poll data (1998-2024).
Run: python -I backtest.py   (reads data/sb_rawpolls_021325.xlsx, writes out/*.csv and out/results.txt)
Sign convention: margin = DEM - REP (points). err = poll_avg - actual; err>0 means polls too favourable to DEM.
"""
import sys, json
import numpy as np, pandas as pd
from scipy.stats import norm

WINDOW = 21          # days before election
MIN_POLLS = 2        # distinct pollsters per race (sensitivity: 1, 3)
PRES_PARTY = {1998:1,1999:1,2000:1,2002:-1,2004:-1,2006:-1,2008:-1,2010:1,2012:1,2014:1,2016:1,2018:-1,2020:-1,2022:1,2024:1}  # party of WH during the election (+1 D, -1 R)

def load(min_polls=MIN_POLLS, window=WINDOW):
    d = pd.read_excel('data/sb_rawpolls_021325.xlsx')
    d = d[d.type_simple.isin(['Pres-G','Sen-G','Gov-G','House-G'])]
    d = d[(d.cand1_party=='DEM')&(d.cand2_party=='REP')&d.partisan.isna()]   # head-to-head D vs R, nonpartisan sponsors only
    d['days'] = (d.electiondate - d.polldate).dt.days
    d = d[(d.days>=0)&(d.days<=window)]
    d['cyc'] = d.electiondate.dt.year
    d = d[d.cyc%2==0]
    d['state'] = d.location.str[:2]
    # one value per (race, pollster): mean margin; sample size summed
    g = d.groupby(['race','cyc','type_simple','state','location','electiondate','pollster'],as_index=False).agg(m=('margin_poll','mean'),n=('samplesize','sum'),act=('margin_actual','first'))
    r = g.groupby(['race','cyc','type_simple','state','location','electiondate'],as_index=False).agg(avg=('m','mean'),k=('pollster','nunique'),N=('n','sum'),act=('act','first'))
    r = r[r.k>=min_polls].copy()
    r['err'] = r.avg - r.act
    r['wh'] = r.cyc.map(PRES_PARTY)
    # state-partisanship prior: most recent earlier presidential margin in same state (US for national not used here)
    pres = r[r.type_simple=='Pres-G'].groupby(['state','cyc']).act.first().reset_index()
    # include pres results even for states where the pres race fell below the poll filters: use full data
    full = pd.read_excel('data/sb_rawpolls_021325.xlsx'); full = full[(full.type_simple=='Pres-G')&(full.cand1_party=='DEM')&(full.cand2_party=='REP')]
    full['cyc']=full.electiondate.dt.year; pres = full.groupby([full.location.str[:2],'cyc']).margin_actual.first().reset_index(); pres.columns=['state','pcyc','plean']
    def lean(row):
        p = pres[(pres.state==row.state)&(pres.pcyc<row.cyc)]
        return p.sort_values('pcyc').plean.iloc[-1] if len(p) else np.nan
    r['lean'] = r.apply(lean,axis=1)
    return r

def ols(X,y):
    b,*_ = np.linalg.lstsq(X,y,rcond=None); return b

def run(r, tag):
    r = r.dropna(subset=['lean']).copy()
    cycles = sorted(r.cyc.unique())
    rows=[]
    prev_bias = r.groupby('cyc').err.mean()
    for c in cycles:
        tr, te = r[r.cyc!=c], r[r.cyc==c]
        if len(te)==0: continue
        a = te.avg.values; y = te.act.values
        out = {}
        out['M0 plain average'] = (a, tr.err.std(ddof=1))
        # M1 constant shift
        mu = tr.err.mean(); out['M1 constant shift (global bias)'] = (a-mu, (tr.err-mu).std(ddof=1))
        # M1b: cycle-level shrunk shift, using cycle means with empirical-Bayes shrink
        cm = tr.groupby('cyc').err.agg(['mean','count'])
        s2_between = max(cm['mean'].var(ddof=1) - (tr.groupby('cyc').err.var().mean()/cm['count'].mean()), 0)
        # M2 constant + WH-party term
        X = np.c_[np.ones(len(tr)), tr.wh]; b = ols(X,tr.err.values); res = tr.err.values - X@b
        out['M2 shift + WH-party term'] = (a - (b[0]+b[1]*te.wh.values), res.std(ddof=1))
        # M3 shrink toward state lean: act = b*avg+(1-b)*lean  <=> act-lean = b*(avg-lean)
        X = (tr.avg-tr.lean).values[:,None]; b = ols(X,(tr.act-tr.lean).values); res=(tr.act-tr.lean).values - X@b
        p3 = te.lean.values + b[0]*(te.avg.values-te.lean.values); out['M3 shrink to state lean'] = (p3, res.std(ddof=1))
        # M4 M3 + intercept + wh
        X = np.c_[np.ones(len(tr)), (tr.avg-tr.lean).values, tr.wh]; b = ols(X,(tr.act-tr.lean).values); res=(tr.act-tr.lean).values-X@b
        p4 = te.lean.values + b[0] + b[1]*(te.avg.values-te.lean.values) + b[2]*te.wh.values; out['M4 shrink + intercept + WH'] = (p4, res.std(ddof=1))
        # M5 persistence: subtract previous cycle's mean bias (nearest earlier cycle in data), half weight
        prior = [k for k in prev_bias.index if k<c]
        if prior:
            pb = prev_bias[max(prior)]
            out['M5 persistence (half of last cycle bias)'] = (a-0.5*pb, tr.err.std(ddof=1))
        for name,(pred,sd) in out.items():
            e = pred - y
            for z,lab in ((1.2816,'cov80'),(1.6449,'cov90')):
                pass
            rows.append(dict(cyc=c,model=name,n=len(te),mae=np.mean(np.abs(e)),bias=e.mean(),sq=np.mean(e**2),
                             cov80=np.mean(np.abs(e)<=1.2816*sd),cov90=np.mean(np.abs(e)<=1.6449*sd),
                             wrongcall=np.mean(np.sign(pred)!=np.sign(y)), width90=2*1.6449*sd))
    R = pd.DataFrame(rows)
    # naive sampling-only interval for M0
    sd_n = 100*np.sqrt(0.95/r.N.clip(lower=300)) # approx SD of margin in pts from pooled sample
    nv = (np.abs(r.err)<=1.6449*sd_n).mean(); nv80=(np.abs(r.err)<=1.2816*sd_n).mean()
    tab = R.groupby('model').apply(lambda g: pd.Series(dict(
        MAE=np.average(g.mae,weights=g.n), mean_error=np.average(g.bias,weights=g.n), RMSE=np.sqrt(np.average(g.sq,weights=g.n)),
        cov80=np.average(g.cov80,weights=g.n), cov90=np.average(g.cov90,weights=g.n), wrongcall=np.average(g.wrongcall,weights=g.n),
        width90=g.width90.mean(), races=g.n.sum(), cycles=g.cyc.nunique())),include_groups=False).round(3)
    return R, tab, nv, nv80, r

def summarize(r, f):
    P=lambda *a: print(*a, file=f)
    P('\n== Cycle-level polling error (poll avg minus actual, DEM-REP pts; + = polls too pro-DEM) ==')
    t = r.groupby('cyc').agg(races=('err','size'),mean_err=('err','mean'),sd_err=('err','std'),mae=('err',lambda x: x.abs().mean())).round(2)
    t['se_mean']=(t.sd_err/np.sqrt(t.races)).round(2); P(t.to_string())
    P('\n== By cycle x office: mean error ==')
    P(r.pivot_table(index='cyc',columns='type_simple',values='err',aggfunc=['mean','count']).round(2).to_string())
    # variance components (one-way random effects, ICC)
    g = r.groupby('cyc').err; k=g.ngroup().nunique(); N=len(r)
    msb = (g.size()*(g.mean()-r.err.mean())**2).sum()/(k-1); msw = ((r.err-g.transform('mean'))**2).sum()/(N-k)
    n0 = (N-(g.size()**2).sum()/N)/(k-1); var_b=max((msb-msw)/n0,0)
    P(f'\nVariance components: between-cycle sd={np.sqrt(var_b):.2f} pts, within-cycle (race) sd={np.sqrt(msw):.2f}, ICC={var_b/(var_b+msw):.2f}')
    # presidential-years vs midterm
    r['pres_year']=(r.cyc%4==0)
    P('\nMean error pres-year vs midterm:'); P(r.groupby('pres_year').err.agg(['mean','std','count']).round(2).to_string())
    # sign of cycle bias
    cm=r.groupby('cyc').err.mean()
    P('\nCycle signs: '+', '.join(f'{c}:{"D" if v>0 else "R"}' for c,v in cm.items()))
    P(f'Cycles polls too pro-DEM: {(cm>0).sum()} of {len(cm)}; mean of cycle means {cm.mean():.2f}, sd {cm.std():.2f}, t-stat of mean {cm.mean()/(cm.std()/np.sqrt(len(cm))):.2f}')
    # lag-1 autocorrelation of cycle means (all cycles in sequence, midterm and pres alternate)
    P(f'Lag-1 corr of cycle mean error: {np.corrcoef(cm.values[:-1],cm.values[1:])[0,1]:.2f}; lag-2 (same cycle type): {np.corrcoef(cm.values[:-2],cm.values[2:])[0,1]:.2f}')
    # count-of-wrong-calls test, competitive races
    P('\n== Wrong-call counts in competitive races (|avg|<=5): observed vs expected if race errors were INDEPENDENT ==')
    s = r.err.std(); rows=[]
    rr = np.random.default_rng(1)
    for c,gp in r[r.avg.abs()<=5].groupby('cyc'):
        p = pd.Series(norm.cdf(-gp.avg.abs().values/s)); obs=(np.sign(gp.avg)!=np.sign(gp.act)).sum()
        sims = (rr.random((20000,len(p)))<p.values).sum(1)
        rows.append((c,len(gp),round(p.sum(),1),obs,round(np.mean(sims>=obs),3),round(np.mean(sims<=obs),3)))
    P(pd.DataFrame(rows,columns=['cyc','races','expected_wrong','observed_wrong','P(sim>=obs)','P(sim<=obs)']).to_string(index=False))

if __name__=='__main__':
    open('out/results.txt','w').close()
    with open('out/results.txt','a') as f:
        for mp in (2,1,3):
            r = load(min_polls=mp)
            r.to_csv(f'out/races_minpolls{mp}.csv',index=False)
            R,tab,nv,nv80,rr_ = run(r,mp)
            R.to_csv(f'out/loco_by_cycle_minpolls{mp}.csv',index=False)
            print(f'\n######## min distinct pollsters per race = {mp}; races (with lean available)={len(rr_)} ########',file=f)
            print(tab.to_string(),file=f)
            print(f'Naive sampling-error-only 90% interval coverage (M0): {nv:.3f}; 80%: {nv80:.3f}',file=f)
            if mp==2:
                summarize(r.copy(), f)
                # also all races without lean requirement (incl. 2000)
                print('\n(Above uses races with a prior presidential margin available, so 2000 is excluded from backtest tables.)',file=f)
                print('\nAll-race cycle table incl. 2000/2002 first cycles:',file=f)
                print(r.groupby('cyc').err.agg(['count','mean']).round(2).T.to_string(),file=f)
                print('\nPer-cycle LOCO MAE: ',file=f)
                print(R.pivot_table(index='cyc',columns='model',values='mae').round(2).to_string(),file=f)
    print(open('out/results.txt').read())