"""Audit the fast changepoint implementation against direct whitening and scalar ML optimization.
The alternative fit does not use the bundle's covariance polynomials or cofactor inverse.
"""
import json,sys,math,io,zipfile
from pathlib import Path
import numpy as np,pandas as pd
from scipy.optimize import minimize_scalar
root=Path(sys.argv[1]);sys.path.insert(0,str(root/'code'))
import fr,tmax
params=json.loads((root/'data/parameters.json').read_text())
with zipfile.ZipFile(root/'data/foster-rahmstorf/AdjustData.zip') as z:
 obs=pd.read_csv(io.BytesIO(z.read('ObservedData.csv')))
pub=pd.read_csv(root/'data/foster-rahmstorf/GMST_adjusted.csv',comment='#')
output={}
for name in params['datasets']:
 ok=np.isfinite(obs[name]);tm=obs.t.to_numpy()[ok];x=obs[name].to_numpy()[ok]
 adj=fr.Adjuster(tm,np.nan_to_num(obs.volc.to_numpy(),nan=0)[ok],obs.NINO34.to_numpy()[ok],obs.ssn.to_numpy()[ok],fit=tuple(params['fit_window']),nback=params['nback']).adjust(x)[0]
 published=pub.set_index(np.round(pub.time,4))[name+'.adj'].reindex(np.round(tm,4)).to_numpy()
 delta=published-adj;dev=float(np.nanmax(np.abs(delta-np.nanmean(delta))))
 years,y=fr.annual(tm,adj,*params['test_years']);n=len(y);t=np.arange(1,n+1)
 manual=[]
 for k in range(math.ceil(.1*n),math.floor(.9*n)+1):
  X=np.column_stack([np.ones(n),t,np.maximum(t-k,0)])
  def fit(phi,full=False):
   Xw=np.vstack([np.sqrt(1-phi*phi)*X[:1],X[1:]-phi*X[:-1]])
   yw=np.r_[np.sqrt(1-phi*phi)*y[0],y[1:]-phi*y[:-1]]
   beta=np.linalg.lstsq(Xw,yw,rcond=None)[0];r=yw-Xw@beta;rss=r@r
   objective=n*np.log(rss/n)/2-np.log(1-phi*phi)/2
   if full:return float(beta[2]/np.sqrt((rss/n)*np.linalg.inv(Xw.T@Xw)[2,2]))
   return objective
  grid=np.linspace(-.995,.995,201);i=int(np.argmin([fit(p) for p in grid]))
  opt=minimize_scalar(fit,bounds=(grid[max(i-1,0)],grid[min(i+1,200)]),method='bounded',options={'xatol':1e-11})
  manual.append(fit(opt.x,True))
 test=tmax.TmaxTest(n);fast,k,all_fast=test.stats(y[None])
 diff=float(np.max(np.abs(np.array(manual)-all_fast[0])))
 j=int(np.argmax(np.abs(manual)));direct=float(abs(manual[j]));year=int(years[0]-1+math.ceil(.1*n)+j)
 output[name]={'published_adjustment_max_deviation_after_offset':dev,'fast_tmax':float(fast[0]),'direct_tmax':direct,'direct_change_year':year,'max_statistic_difference':diff}
 assert abs(direct-fast[0])<.02,(name,direct,fast)
 assert dev<=.0002,(name,dev)
print(json.dumps(output,indent=2))