"""India PM2.5 trend analysis; offline frozen-input rerun.
Run: python analysis.py
Requires existing numpy, pandas, scipy, matplotlib. No network or installation.
See analysis_conventions.md for preregistered computational conventions.
"""
import json, hashlib, sys
from pathlib import Path
import numpy as np
import pandas as pd
from scipy import stats
SEED = 20261008
B = 1999
RHO_GRID = np.round(np.arange(-.90,.901,.02), 2)

class Family:
    def __init__(self, years, kind, errors, minimum=5):
        self.years=np.asarray(years)
        self.n=len(years)
        self.kind=kind
        self.errors=errors
        self.rhos=RHO_GRID if errors=='AR1' else np.array([0.])
        t=self.years-self.years[0]
        self.knots=np.array([k for k in years if sum(self.years<k)>=minimum and sum(self.years>k)>=minimum]) if kind=='hinge' else np.array([np.nan])
        if not len(self.knots): raise ValueError('No eligible knots')
        self.X=np.stack([np.column_stack([np.ones(self.n),t,np.maximum(0,self.years-k)]) for k in self.knots]) if kind=='hinge' else np.array([np.column_stack([np.ones(self.n),t,t*t]) if kind=='quadratic' else np.column_stack([np.ones(self.n),t])])
        self.K={'line':3,'quadratic':4,'hinge':5}[kind]+int(errors=='AR1')
        self.cache=[]
        for rho in self.rhos:
            a=rho**np.diff(self.years)
            den=np.sqrt(1-a*a)
            wx=self.X.copy()
            wx[:,1:,:]=(self.X[:,1:,:]-a[None,:,None]*self.X[:,:-1,:])/den[None,:,None]
            Q,_=np.linalg.qr(wx,mode='reduced')
            self.cache.append((a,den,Q,np.linalg.pinv(wx),float(np.log1p(-a*a).sum())))
    def fit(self,Y):
        Y=np.atleast_2d(Y)
        best=np.full((len(Y),len(self.knots)),np.inf)
        ir=np.zeros(best.shape,dtype=int)
        for ri,(a,den,Q,pinv,ld) in enumerate(self.cache):
            wy=Y.copy()
            wy[:,1:]=(Y[:,1:]-a*Y[:,:-1])/den
            proj=np.einsum('bn,knp->bkp',wy,Q,optimize=False)
            q=np.maximum(np.sum(wy*wy,axis=1)[:,None]-np.sum(proj*proj,axis=2),1e-12)
            score=self.n*(np.log(2*np.pi)+1+np.log(q/self.n))+ld
            mask=score<best
            best[mask]=score[mask]
            ir[mask]=ri
        ik=np.argmin(best,axis=1)
        return {'score':best[np.arange(len(Y)),ik],'kindex':ik,'rhoindex':ir[np.arange(len(Y)),ik],'profile':best,'profile_rhoindex':ir}
    def detail(self,y,fit=None):
        fit=self.fit(y) if fit is None else fit
        ki=int(fit['kindex'][0]); ri=int(fit['rhoindex'][0])
        a,den,Q,pinv,ld=self.cache[ri]
        wy=np.array(y,copy=True); wy[1:]=(y[1:]-a*y[:-1])/den
        beta=pinv[ki]@wy
        mu=self.X[ki]@beta
        resid=y-mu
        wr=resid.copy(); wr[1:]=(resid[1:]-a*resid[:-1])/den
        return {'beta':beta,'mu':mu,'resid':resid,'white':wr,'sigma':float(np.sqrt(wr@wr/self.n)),'rho':float(self.rhos[ri]),'knot':float(self.knots[ki]),'score':float(fit['score'][0]),'K':self.K}
    def batch_parameters(self,Y,fit):
        betas=[]
        for b,y in enumerate(Y):
            ki=int(fit['kindex'][b]); ri=int(fit['rhoindex'][b])
            a,den,_,pinv,_=self.cache[ri]
            wy=y.copy(); wy[1:]=(y[1:]-a*y[:-1])/den
            betas.append(pinv[ki]@wy)
        return np.array(betas)

def generate_series(rng,detail,years,reps):
    R=detail['rho']**np.abs(np.subtract.outer(years,years))
    return detail['mu']+detail['sigma']*(rng.standard_normal((reps,len(years)))@np.linalg.cholesky(R).T)

def compare(years,raw,scale,errors,minimum=5):
    y=np.log(raw) if scale=='log' else np.array(raw)
    jac=2*np.log(raw).sum() if scale=='log' else 0.
    rows=[]; details={}; families={}; profiles=[]
    for kind in ['line','hinge','quadratic']:
        f=Family(years,kind,errors,minimum)
        fit=f.fit(y); d=f.detail(y,fit)
        n=len(years); K=d['K']; score=d['score']+jac
        beta=d['beta']
        rows.append(dict(model=kind,scale=scale,errors=errors,n=n,start=int(years[0]),end=int(years[-1]),minimum=minimum,knot=d['knot'],rho=d['rho'],rho_boundary=abs(d['rho'])>=.899,minus2loglik=score,K=K,AICc=score+2*K+2*K*(K+1)/(n-K-1),BIC=score+K*np.log(n),slope_before=beta[1],slope_after=beta[1]+beta[2] if kind=='hinge' else np.nan,slope_change=beta[2] if kind=='hinge' else np.nan,quadratic_coefficient=beta[2] if kind=='quadratic' else np.nan,sigma=d['sigma'],residual_lag1=np.corrcoef(d['resid'][:-1],d['resid'][1:])[0,1],white_lag1=np.corrcoef(d['white'][:-1],d['white'][1:])[0,1]))
        details[kind]=d; families[kind]=f
        if kind=='hinge':
            profiles=[dict(year=int(k),minus2loglik=float(s+jac),rho=float(f.rhos[r])) for k,s,r in zip(f.knots,fit['profile'][0],fit['profile_rhoindex'][0])]
    return pd.DataFrame(rows),details,families,profiles

def primary_analysis(df):
    years=df.year.to_numpy(); raw=df.pm25_ug_m3.to_numpy()
    rng=np.random.default_rng(SEED)
    rows=[]; profiles=[]; tests=[]; simulations=[]; primary={}
    for errors in ['iid','AR1']:
        tab,details,families,prof=compare(years,raw,'raw',errors)
        rows.extend(tab.to_dict('records'))
        profiles.extend([dict(p,errors=errors) for p in prof])
        primary[errors]=(tab,details,families)
        Y=generate_series(rng,details['line'],years,B)
        fl=families['line'].fit(Y); fh=families['hinge'].fit(Y)
        simulated=np.maximum(0,fl['score']-fh['score'])
        observed=max(0,details['line']['score']-details['hinge']['score'])
        p=(1+sum(simulated>=observed))/(B+1)
        tests.append(dict(errors=errors,statistic=observed,p_value=p,MC_SE=np.sqrt(p*(1-p)/(B+1)),replicates=B))
        simulations.extend([dict(errors=errors,replicate=i,simulation='null',statistic=float(v),knot=float(families['hinge'].knots[fh['kindex'][i]])) for i,v in enumerate(simulated)])
    return pd.DataFrame(rows),pd.DataFrame(profiles),pd.DataFrame(tests),pd.DataFrame(simulations),primary

def sensitivity_analysis(df,primary):
    years=df.year.to_numpy(); raw=df.pm25_ug_m3.to_numpy()
    rows=[]; comparisons=[]; profiles=[]; loo=[]
    windows={'full':np.ones(len(years),dtype=bool),'drop_first3':years>=years[3],'drop_last3':years<=years[-4],'drop_both3':(years>=years[3])&(years<=years[-4])}
    for window,mask in windows.items():
        for scale in ['raw','log']:
            for errors in ['iid','AR1']:
                for minimum in [4,5,6]:
                    tab,details,families,prof=compare(years[mask],raw[mask],scale,errors,minimum)
                    a=tab.loc[tab.AICc.idxmin()]; b=tab.loc[tab.BIC.idxmin()]
                    h=tab[tab.model=='hinge'].iloc[0]; l=tab[tab.model=='line'].iloc[0]
                    rows.append(dict(window=window,scale=scale,errors=errors,minimum=minimum,n=int(mask.sum()),start=int(years[mask][0]),end=int(years[mask][-1]),eligible_first=int(families['hinge'].knots[0]),eligible_last=int(families['hinge'].knots[-1]),knot=h.knot,knot_boundary=h.knot in families['hinge'].knots[[0,-1]],rho=h.rho,rho_boundary=h.rho_boundary,slope_before=h.slope_before,slope_after=h.slope_after,slope_change=h.slope_change,winner_AICc=a.model,winner_BIC=b.model,hinge_minus_line_AICc=h.AICc-l.AICc,hinge_minus_line_BIC=h.BIC-l.BIC,within_2_years_of_primary=abs(h.knot-1995)<=2))
                    comparisons.extend([dict(r,window=window) for r in tab.to_dict('records')])
                    profiles.extend([dict(p,window=window,scale=scale,errors=errors,minimum=minimum) for p in prof])
    for omitted in years:
        mask=years!=omitted
        for scale in ['raw','log']:
            for errors in ['iid','AR1']:
                tab,details,families,_=compare(years[mask],raw[mask],scale,errors,5)
                h=tab[tab.model=='hinge'].iloc[0]
                loo.append(dict(omitted_year=int(omitted),scale=scale,errors=errors,knot=h.knot,slope_change=h.slope_change,winner_AICc=tab.loc[tab.AICc.idxmin(),'model'],winner_BIC=tab.loc[tab.BIC.idxmin(),'model'],knot_boundary=h.knot in families['hinge'].knots[[0,-1]],rho=h.rho))
    rng=np.random.default_rng(SEED+1)
    conditional=[]; intervals=[]
    for errors,(_,details,families) in primary.items():
        Y=generate_series(rng,details['hinge'],years,B)
        fit=families['hinge'].fit(Y)
        betas=families['hinge'].batch_parameters(Y,fit)
        ks=families['hinge'].knots[fit['kindex']]
        for i in range(B):
            conditional.append(dict(errors=errors,replicate=i,knot=int(ks[i]),rho=float(families['hinge'].rhos[fit['rhoindex'][i]]),slope_before=betas[i,1],slope_after=betas[i,1]+betas[i,2],slope_change=betas[i,2]))
        for label,values in [('knot',ks),('slope_before',betas[:,1]),('slope_after',betas[:,1]+betas[:,2]),('slope_change',betas[:,2])]:
            qs=np.quantile(values,[.025,.5,.975])
            intervals.append(dict(errors=errors,parameter=label,lower_025=qs[0],median=qs[1],upper_975=qs[2],conditional_on='fitted hinge model'))
    return tuple(pd.DataFrame(x) for x in [rows,comparisons,profiles,loo,conditional,intervals])

def make_figure(df,primary,profile,sensitivity,conditional,destination):
    import matplotlib
    matplotlib.use('Agg')
    import matplotlib.pyplot as plt
    years=df.year.to_numpy()
    with plt.rc_context({'font.size':10,'axes.spines.top':False,'axes.spines.right':False,'figure.facecolor':'white','savefig.facecolor':'white'}):
        fig,axs=plt.subplots(2,2,figsize=(12,8),layout='constrained')
        ax=axs[0,0]
        ax.plot(years,df.pm25_ug_m3,'o-',color='#555555',lw=.8,ms=3,label='Annual exposure estimate')
        for kind,col,ls in [('line','#0072B2','-'),('hinge','#D55E00','--'),('quadratic','#009E73',':')]:
            ax.plot(years,primary['iid'][1][kind]['mu'],color=col,ls=ls,lw=2,label=kind)
        ax.set(title='A  National exposure and fitted trends (iid)',xlabel='Year',ylabel='PM2.5 exposure (µg/m³)')
        ax.legend(fontsize=8)
        ax=axs[0,1]
        for errors,col,marker in [('iid','#0072B2','o'),('AR1','#D55E00','s')]:
            sub=profile[profile.errors==errors]
            line=primary[errors][0].query("model == 'line'").iloc[0]
            K=primary[errors][1]['hinge']['K']; n=len(years)
            delta=sub.minus2loglik+2*K+2*K*(K+1)/(n-K-1)-line.AICc
            ax.plot(sub.year,delta,color=col,marker=marker,ms=3,label=errors)
        ax.axhline(0,color='black',ls=':')
        ax.set(title='B  Candidate profile: hinge versus line',xlabel='Candidate knot year',ylabel='AICc(hinge) − AICc(line)')
        ax.text(.03,.94,'Positive values favor the no-break line',transform=ax.transAxes,va='top',fontsize=8)
        ax.legend(loc='lower right')
        ax=axs[1,0]
        order=['full','drop_first3','drop_last3','drop_both3']
        labels=['1990–2023','1993–2023','1990–2020','1993–2020']
        for j,w in enumerate(order):
            sub=sensitivity[sensitivity.window==w]
            for z,(_,row) in enumerate(sub.iterrows()):
                ax.plot(row.knot,j+(z-5.5)*.025,marker='o' if row.scale=='raw' else 'x',color='#0072B2' if row.errors=='iid' else '#D55E00',ms=5,ls='none')
        ax.set_yticks(range(4),labels)
        ax.invert_yaxis()
        ax.set(title='C  Selected hinge year in all 48 settings',xlabel='Candidate knot year',ylabel='Analysis window')
        ax.text(.02,.04,'o raw; × log; blue iid; orange AR(1)\nIncludes minimum segments 4, 5, 6',transform=ax.transAxes,fontsize=8)
        ax.set_ylim(3.7,-.5)
        ax=axs[1,1]
        for errors,col,ls in [('iid','#0072B2','-'),('AR1','#D55E00','--')]:
            freq=conditional[conditional.errors==errors].knot.value_counts().reindex(np.arange(1995,2019),fill_value=0)/B
            ax.step(freq.index,freq.values,where='mid',color=col,ls=ls,label=errors,lw=1.5)
        ax.set(title='D  Knot uncertainty conditional on a hinge',xlabel='Selected knot year in simulated series',ylabel='Fraction of 1,999 simulations')
        ax.set_ylim(bottom=0)
        ax.legend()
        fig.suptitle('India PM2.5, 1990–2023: insufficient evidence for a stable trend break',fontsize=14)
        fig.savefig(destination,dpi=180)
        plt.close(fig)

def numerical_checks(df,primary):
    years=df.year.to_numpy(); y=df.pm25_ug_m3.to_numpy()
    checks=[]
    for errors,(_,details,families) in primary.items():
        for kind,d in details.items():
            f=families[kind]
            ki=int(np.where(f.knots==d['knot'])[0][0]) if kind=='hinge' else 0
            X=f.X[ki]
            R=d['rho']**np.abs(np.subtract.outer(years,years))
            inv=np.linalg.inv(R)
            beta=np.linalg.solve(X.T@inv@X,X.T@inv@y)
            r=y-X@beta
            dense=len(y)*(np.log(2*np.pi)+1+np.log(r@inv@r/len(y)))+np.linalg.slogdet(R)[1]
            error=abs(dense-d['score'])
            assert error<1e-7
            assert np.max(np.abs(beta-d['beta']))<1e-7
            checks.append(dict(errors=errors,model=kind,dense_likelihood_abs_error=error,shapiro_p=stats.shapiro(d['white']).pvalue,max_abs_standardized_white=float(np.max(np.abs(d['white']/d['sigma'])))))
    # Explicit calendar-gap covariance check for leave-one-out calculations.
    mask=years!=2005
    tab,details,families,_=compare(years[mask],y[mask],'raw','AR1')
    d=details['hinge']; yy=years[mask]; r=d['resid']
    R=d['rho']**np.abs(np.subtract.outer(yy,yy))
    dense=len(yy)*(np.log(2*np.pi)+1+np.log(r@np.linalg.solve(R,r)/len(yy)))+np.linalg.slogdet(R)[1]
    assert abs(dense-d['score'])<1e-7
    return pd.DataFrame(checks)

def run_analysis(input_csv='india_pm25.csv',output_dir='rerun_results'):
    out=Path(output_dir)
    out.mkdir(exist_ok=False)
    df=pd.read_csv(input_csv)
    assert df.year.tolist()==list(range(1990,2024))
    assert df.pm25_ug_m3.notna().all() and (df.pm25_ug_m3>0).all()
    mc,bp,tests,null,primary=primary_analysis(df)
    sens,allmc,allbp,loo,cond,intervals=sensitivity_analysis(df,primary)
    checks=numerical_checks(df,primary)
    tables={'model_comparison':mc,'breakpoint_profile':bp,'null_tests':tests,'null_simulations':null,'sensitivity_results':sens,'all_model_comparisons':allmc,'all_breakpoint_profiles':allbp,'leave_one_out':loo,'conditional_simulations':cond,'conditional_intervals':intervals,'numerical_checks':checks}
    for name,table in tables.items(): table.to_csv(out/(name+'.csv'),index=False,float_format='%.12g')
    make_figure(df,primary,bp,sens,cond,out/'trend_diagnostics.png')
    return tables,primary

if __name__=='__main__':
    tables,primary=run_analysis()
    print(tables['null_tests'].to_string(index=False))
    print('Outputs saved to rerun_results; remove or rename that directory before rerunning.')
