#!/usr/bin/env python3
"""Section 1.6 full-period SPEI and four AI-stack candidates.

All distribution fitting, AI fitting, prediction and verification use the continuous
1982-01 through 2021-12 period. No train/validation split and no best selection.
"""
import argparse,csv,html,json,logging,os,sys,warnings
from datetime import datetime
from pathlib import Path
import numpy as np
from scipy.io import netcdf_file
from sklearn.metrics import roc_auc_score
import joblib
import spei_pipeline as core

ROOT=core.ROOT; FULL_RUN='1399282'; LEADS=range(1,7); SCALES=(1,3,6)
LEVELS={'D1':-1.0,'D2':-1.5,'D3':-2.0}; MAP_LEADS=(1,3,5); FILL=np.float32(-99.)

def log(run):
 p=ROOT/'logs'/f'spei_members_{run}.log';p.parent.mkdir(parents=True,exist_ok=True)
 logging.basicConfig(level=logging.INFO,format='%(asctime)s %(levelname)s rank=%(rank)s %(message)s',handlers=[logging.FileHandler(p),logging.StreamHandler(sys.stdout)])
 old=logging.getLogRecordFactory()
 def factory(*a,**k):r=old(*a,**k);r.rank=os.getenv('SLURM_PROCID','driver');return r
 logging.setLogRecordFactory(factory);return logging.getLogger('spei_members')

def write_field(path,var,data,lon,lat,third_name,third_values,attrs):
 path.parent.mkdir(parents=True,exist_ok=True)
 with netcdf_file(path,'w') as d:
  dims=[('lon',len(lon)),('lat',len(lat)),(third_name,len(third_values)),('time',480)]
  for n,k in dims:d.createDimension(n,k)
  for n,x in [('lon',lon),('lat',lat),(third_name,third_values),('time',np.arange(1,481))]:d.createVariable(n,'f',(n,))[:]=x
  q=d.createVariable(var,'f',tuple(n for n,_ in dims));q._FillValue=FILL;q.missing_value=FILL;q[:]=np.where(np.isfinite(data),data,FILL)
  d.time_mapping='1=1982-01;480=2021-12'
  for k,v in attrs.items():setattr(d,k,str(v))

def unit(lead,scale,run,source,lg):
 lon,lat,mask=core.coordinates_and_mask();pairs,names=core.read_members(lead,source,True)
 pairs=pairs[:48];names=names[:48];fit=np.ones(480,bool);members=[];params=[]
 for m,(pr,t2) in enumerate(pairs):
  spei,par=core.fit_spei(pr-core.thornthwaite(t2,lat),scale,fit);spei[~mask]=np.nan;par[~mask]=np.nan
  members.append(spei);params.append(par);lg.info('stage=SPEI LT=%d scale=%d member=%02d/48',lead,scale,m+1)
 arr48=np.stack(members,axis=2)
 obs_pair=core.read_obs(lead);obs,obspar=core.fit_spei(obs_pair[0]-core.thornthwaite(obs_pair[1],lat),scale,fit);obs[~mask]=np.nan;obspar[~mask]=np.nan
 out=ROOT/'outputs/production_output/spei_members'/run
 core.write_nc(out/f'SPEI{scale}_members_LT{lead}.nc',f'SPEI{scale}',arr48,lon,lat,names,{'source_step_1_5_run':source,'calibration':'1982-01/2021-12','members':'48 SPEI fields from 48 pr/t2 pairs','timescale_months':scale,'lead':lead})
 np.savez_compressed(out/f'SPEI{scale}_distribution_parameters_LT{lead}.npz',member=np.stack(params,axis=2),observed=obspar,parameter_order=np.array(['alpha_scale','beta_shape','gamma_location']),calendar_month=np.arange(1,13),calibration='1982-01/2021-12')
 np.savez_compressed(out/f'SPEI{scale}_observed_LT{lead}.npz',spei=obs)
 candidates=np.full((len(lon),len(lat),4,480),np.nan,np.float32);artifacts={};records=[]
 for point,(i,j) in enumerate(np.argwhere(mask)):
  X=arr48[i,j].T;y=obs[i,j];keep=np.sum(np.isfinite(X),axis=0)>=30
  if keep.sum()<2:continue
  X=X[:,keep];means=np.nanmean(X,axis=0);has=np.any(np.isfinite(X),axis=1);Xi=np.where(np.isfinite(X),X,means);ok=np.isfinite(y)&has
  if ok.sum()<30:continue
  for mi,method in enumerate(core.METHODS):
   try:
    model=core.estimator(method,lead*100000+scale*1000+i*100+j);model.fit(Xi[ok],y[ok]);candidates[i,j,mi,has]=model.predict(Xi[has]);mae,corr=core.scores(y,candidates[i,j,mi],np.ones(480,bool))
    artifacts[f'{i}_{j}_{method}']={'model':model,'member_indices_zero_based':np.flatnonzero(keep),'imputation_full_period_means':means};records.append({'i':i,'j':j,'method':method,'n_features':int(keep.sum()),'full_period_mae':mae,'full_period_corr':corr,'status':'passed'})
   except Exception as exc:records.append({'i':i,'j':j,'method':method,'status':'failed','reason':f'{type(exc).__name__}: {exc}'})
  if point%25==0:lg.info('stage=AI_stack LT=%d scale=%d progress=%d/500',lead,scale,point)
 core.write_nc(out/f'SPEI{scale}_AI_stack_candidates_LT{lead}.nc',f'SPEI{scale}_stack',candidates,lon,lat,list(core.METHODS),{'candidate_member_number':49,'selection':'none; all four retained','features':'48 SPEI members','target':f'SPEI{scale}_observed','fit_period':'1982-01/2021-12; all valid samples; no split'})
 art=out/'stack_model_artifacts';art.mkdir(exist_ok=True);joblib.dump(artifacts,art/f'SPEI{scale}_AI_stack_candidates_LT{lead}.joblib',compress=3)
 with open(art/f'SPEI{scale}_AI_stack_candidates_LT{lead}.csv','w',newline='') as f:
  keys=sorted(set().union(*(x.keys() for x in records)));w=csv.DictWriter(f,fieldnames=keys);w.writeheader();w.writerows(records)
 for mi,method in enumerate(core.METHODS):
  arr49=np.concatenate([arr48,candidates[:,:,mi:mi+1,:]],axis=2)
  prob=np.stack([np.nanmean(np.where(np.isfinite(arr49),arr49<=thr,np.nan),axis=2)*100 for thr in LEVELS.values()],axis=2);prob[~mask]=np.nan
  write_field(out/f'SPEI{scale}_drought_probability_{method}_LT{lead}.nc','drought_probability',prob,lon,lat,'level',np.arange(1,4),{'units':'percent','levels':'D1<=-1|D2<=-1.5|D3<=-2','denominator':'48 SPEI members plus named AI-stack candidate','AI_stack_candidate':method,'timescale_months':scale,'lead':lead})
 lg.info('stage=unit LT=%d scale=%d status=SUCCESS',lead,scale)

def metric_fields(obs,pred,thr):
 valid=np.isfinite(obs)&np.isfinite(pred);o=(obs<=thr)&valid;p=(pred<=thr)&valid
 h=np.sum(o&p,2);m=np.sum(o&~p&valid,2);f=np.sum(~o&p&valid,2);c=np.sum(~o&~p&valid,2);div=lambda a,b:np.divide(a,b,out=np.full(a.shape,np.nan,float),where=b>0)
 auc=np.full(obs.shape[:2],np.nan)
 for i,j in np.argwhere(np.sum(valid,2)>=3):
  q=valid[i,j];target=o[i,j,q]
  if np.unique(target).size==2:auc[i,j]=roc_auc_score(target,-pred[i,j,q])
 return {'PC':div(h+c,h+m+f+c),'POD':div(h,h+m),'POFD':div(f,f+c),'FAR':div(f,h+f),'FBI':div(h+f,h+m),'CSI':div(h,h+m+f),'AUC':auc,'n':h+m+f+c}

def finalize(run,lg):
 import matplotlib;matplotlib.use('Agg');import matplotlib.pyplot as plt
 out=ROOT/'outputs/production_output/spei_members'/run;rep=ROOT/'outputs/evaluation_output/spei_members'/run;figs=rep/'figures';figs.mkdir(parents=True,exist_ok=True)
 lon,lat,mask=core.coordinates_and_mask();rows=[];fields={}
 for scale in SCALES:
  for lead in LEADS:
   candidates=core.nc_array(out/f'SPEI{scale}_AI_stack_candidates_LT{lead}.nc',f'SPEI{scale}_stack');obs=np.load(out/f'SPEI{scale}_observed_LT{lead}.npz')['spei']
   for mi,method in enumerate(core.METHODS):
    for level,thr in LEVELS.items():
     f=metric_fields(obs,candidates[:,:,mi],thr);fields[(scale,lead,level,method)]=f;row={'period':'1982-01/2021-12','scale':scale,'lead':lead,'AI_stack_candidate':method,'level':level,'threshold':thr}
     for k,v in f.items():row[k]=float(np.nanmean(v[mask]))
     rows.append(row)
 with open(rep/'drought_verification_summary.csv','w',newline='') as f:w=csv.DictWriter(f,fieldnames=list(rows[0]));w.writeheader();w.writerows(rows)
 metrics=('PC','POD','POFD','FAR','FBI','CSI','AUC')
 for scale in SCALES:
  for method in core.METHODS:
   for level in LEVELS:
    fig,axs=plt.subplots(len(metrics),3,figsize=(11,20),sharex=True,sharey=True,constrained_layout=True)
    for r,m in enumerate(metrics):
     vmax=2 if m=='FBI' else 1
     for c,lead in enumerate(MAP_LEADS):
      ax=axs[r,c];z=fields[(scale,lead,level,method)][m];mesh=ax.pcolormesh(lon,lat,z.T,cmap='viridis',vmin=0,vmax=vmax,shading='nearest');ax.set_title(f'Full period LT{lead}' if r==0 else '');ax.set_ylabel(m if c==0 else '');ax.set_xticks(np.arange(102,111,2));ax.set_yticks(np.arange(8,25,4));ax.tick_params(labelsize=6);fig.colorbar(mesh,ax=ax,fraction=.05,pad=.02,ticks=np.linspace(0,vmax,5))
    fig.suptitle(f'SPEI-{scale} {level}: {method} AI-stack verification');fig.savefig(figs/f'SPEI{scale}_{level}_{method}_verification_maps.png',dpi=170);plt.close(fig)
 for scale in SCALES:
  for method in core.METHODS:
   fig,axs=plt.subplots(3,3,figsize=(13,13),sharex=True,sharey=True,constrained_layout=True);mesh=None
   for r,level in enumerate(LEVELS):
    for c,lead in enumerate(MAP_LEADS):
     prob=core.nc_array(out/f'SPEI{scale}_drought_probability_{method}_LT{lead}.nc','drought_probability')[:,:,r];z=np.nanmean(prob,2);ax=axs[r,c];mesh=ax.pcolormesh(lon,lat,z.T,cmap='YlOrRd',vmin=0,vmax=100,shading='nearest');ax.set_title(f'LT{lead}');ax.set_ylabel(f'{level} probability (%)');ax.set_xticks(np.arange(102,111,2));ax.set_yticks(np.arange(8,25,4));ax.tick_params(labelsize=7)
   fig.colorbar(mesh,ax=axs,orientation='horizontal',fraction=.035,pad=.03,label='Drought probability: 48 SPEI + AI candidate (%)',ticks=np.arange(0,101,10));fig.suptitle(f'SPEI-{scale}: {method} candidate drought probability');fig.savefig(figs/f'SPEI{scale}_{method}_drought_probability_maps.png',dpi=180);plt.close(fig)
 links=''.join(f'<section><h2>{html.escape(f.stem)}</h2><a href="figures/{f.name}"><img src="figures/{f.name}" alt="{html.escape(f.stem)}"></a></section>' for f in sorted(figs.glob('*.png')))
 (rep/'index.html').write_text(f'<!doctype html><html><head><meta charset="utf-8"><title>SPEI AI-stack candidates</title><style>body{{font:14px system-ui;margin:20px}}img{{width:100%}}section{{border:1px solid #aaa;padding:10px;margin:15px 0}}</style></head><body><h1>SPEI AI-stack candidates — run {run}</h1><p>Continuous full period 1982-01–2021-12 with no training/validation split. 48 SPEI members; Ridge1, Ridge2, ANN1 and ANN2 independently produce four candidate member-49 fields. No best model selected. Maps only. <a href="drought_verification_summary.csv">Metrics</a>.</p>{links}</body></html>')
 (out/'SUCCESS').write_text(datetime.now().isoformat());lg.info('stage=finalize status=SUCCESS report=%s',rep/'index.html')

def main():
 p=argparse.ArgumentParser();p.add_argument('stage',choices=['calculate','finalize']);p.add_argument('--run',required=True);p.add_argument('--source-run',default=FULL_RUN);p.add_argument('--rank',type=int,default=int(os.getenv('SLURM_PROCID','0')));p.add_argument('--size',type=int,default=int(os.getenv('SLURM_NTASKS','1')));a=p.parse_args();lg=log(a.run)
 if a.stage=='calculate':
  units=[(l,s) for l in LEADS for s in SCALES]
  for l,s in units[a.rank::a.size]:unit(l,s,a.run,a.source_run,lg)
 else:finalize(a.run,lg)
if __name__=='__main__':main()