#!/usr/bin/env python3
"""Dr_SFC section 1.6: SPEI calculation, correction and drought verification.

The pipeline deliberately keeps evaluation and operational calibration separate.
Evaluation distribution/model parameters are fitted on the training period only;
operational parameters are refitted on 1982-2021 and stored independently.
"""
from __future__ import annotations
import argparse, csv, html, json, logging, os, sys, warnings
from datetime import datetime
from pathlib import Path

import joblib
import numpy as np
from scipy.io import netcdf_file
from scipy.special import gamma as gamma_function
from scipy.stats import norm
from sklearn.linear_model import RidgeCV
from sklearn.metrics import roc_auc_score, roc_curve
from sklearn.model_selection import GridSearchCV, KFold
from sklearn.neural_network import MLPRegressor
from sklearn.pipeline import make_pipeline
from sklearn.preprocessing import StandardScaler

ROOT = Path('/home/tanpv/Progs/CLI/Dr')
INPUT = Path('/home/tanpv/Data/562_VNU/0.25_VN/For_BC')
MASK = Path('/home/tanpv/Progs/Dr_VN_Ope/DATA/Mask_VN_0.25_Corrected.nc')
EVAL_RUN = '1399272'; FULL_RUN = '1399282'; FILL = np.float32(-99.)
LEADS = range(1, 7); SCALES = (1, 3, 6); METHODS = ('Ridge1','Ridge2','ANN1','ANN2')
ALPHAS = np.logspace(-4, 4, 17); DROUGHT = {'D1':-1., 'D2':-1.5, 'D3':-2.}

def getlog(run):
    path=ROOT/'logs'/f'spei_{run}.log'; path.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(path),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')

def nc_array(path, name):
    with netcdf_file(path,'r',mmap=False) as d:
        a=np.array(d.variables[name][:],dtype=np.float64,copy=True)
    a[a <= -90] = np.nan; return a

def coordinates_and_mask():
    with netcdf_file(INPUT/'lead1_RegCM_01_BL_Test_VN.nc','r',mmap=False) as d:
        lon=np.array(d.variables['lon'][:],copy=True); lat=np.array(d.variables['lat'][:],copy=True)
    with netcdf_file(MASK,'r',mmap=False) as d:
        v=d.variables['reg']; reg=np.array(v[:],copy=True)
        if tuple(v.dimensions)==('lat','lon'): reg=reg.T
    return lon,lat,(reg>=1)&(reg<=7)

def read_obs(lead):
    path=INPUT/f'lead{lead}_RegCM_01_BL_Test_VN.nc'
    with netcdf_file(path,'r',mmap=False) as d:
        out=[]
        for name in ('pr_o','t2_o'):
            v=d.variables[name]; a=np.array(v[:],dtype=np.float64,copy=True)
            dims=list(v.dimensions); a=np.transpose(a,[dims.index(x) for x in ('lon','lat','mon')]);a[a<=-90]=np.nan;out.append(a)
    return out

def read_members(lead, run, full=False):
    base=ROOT/'outputs/production_output'/('bias_corrected_full_period' if full else 'bias_corrected')/run
    suffix='_full' if full else ''
    pairs=[]; names=[]
    for rcm in ('RegCM','clWRF'):
        p=nc_array(base/f'{rcm}_pr_bc{suffix}_LT{lead}.nc','pr_bc')
        t=nc_array(base/f'{rcm}_t2_bc{suffix}_LT{lead}.nc','t2_bc')
        for m in range(24): pairs.append((p[:,:,m,:],t[:,:,m,:])); names.append(f'{rcm}_member_{m+1:02d}')
    p=nc_array(base/f'pr_bc_ens{suffix}_LT{lead}.nc','pr_bc_ens')[:,:,0,:]
    t=nc_array(base/f't2_bc_ens{suffix}_LT{lead}.nc','t2_bc_ens')[:,:,0,:]
    pairs.append((p,t)); names.append('Stack_Ensemble'); return pairs,names

def daylength_hours(latitude, month):
    day=np.array([15,45,74,105,135,166,196,227,258,288,319,349])[month-1]
    phi=np.deg2rad(latitude); delta=.409*np.sin(2*np.pi*day/365.-1.39)
    x=np.clip(-np.tan(phi)*np.tan(delta),-1,1); return 24/np.pi*np.arccos(x)

def thornthwaite(temp, lat):
    """Monthly PET (mm/month), including latitude/day-length correction."""
    positive=np.maximum(temp,0); monthly=np.full_like(temp,np.nan)
    for y0 in range(0,temp.shape[-1],12):
        block=positive[...,y0:y0+12]
        I=np.nansum((block/5.)**1.514,axis=-1); a=6.75e-7*I**3-7.71e-5*I**2+1.792e-2*I+.49239
        for q in range(block.shape[-1]):
            days=np.array([31,28,31,30,31,30,31,31,30,31,30,31])[q]
            factor=(daylength_hours(lat[None,:],q+1)/12.)*(days/30.)
            ratio=np.divide(10*block[...,q],I,out=np.zeros_like(I),where=I>0)
            monthly[...,y0+q]=np.where(np.isfinite(block[...,q]),np.where(block[...,q]>0,16*factor*np.power(ratio,a),0),np.nan)
    return monthly

def rolling_sum(a,k):
    out=np.full_like(a,np.nan)
    for t in range(k-1,a.shape[-1]):
        w=a[...,t-k+1:t+1]; ok=np.all(np.isfinite(w),axis=-1);out[...,t]=np.where(ok,np.sum(w,axis=-1),np.nan)
    return out

def fit_spei(balance, scale, fit_time):
    """Fit 3-parameter log-logistic by unbiased probability weighted moments."""
    acc=rolling_sum(balance,scale); spei=np.full_like(acc,np.nan); pars=np.full(balance.shape[:2]+(12,3),np.nan)
    for mo in range(12):
        ids=np.arange(mo,acc.shape[-1],12); train=ids[fit_time[ids]]
        for i,j in np.argwhere(np.any(np.isfinite(acc[...,train]),axis=-1)):
            x=acc[i,j,train];x=x[np.isfinite(x)]
            if x.size<20 or np.ptp(x)<=1e-8: continue
            x=np.sort(x);n=x.size;rank=np.arange(n,dtype=float)
            w0=np.mean(x);w1=np.sum(((n-rank-1)/(n-1))*x)/n
            w2=np.sum(((n-rank-1)*(n-rank-2)/((n-1)*(n-2)))*x)/n
            den=6*w1-w0-6*w2
            if abs(den)<1e-12:continue
            beta=(2*w1-w0)/den
            if not np.isfinite(beta) or beta<=1.01:continue
            alpha=(w0-2*w1)*beta/(gamma_function(1+1/beta)*gamma_function(1-1/beta))
            gamma=w0-alpha*np.pi/(beta*np.sin(np.pi/beta))
            if not np.isfinite(alpha) or alpha<=0:continue
            pars[i,j,mo]=(alpha,beta,gamma);z=acc[i,j,ids];positive=z-gamma
            cdf=np.where(positive>0,1/(1+(alpha/positive)**beta),np.nan)
            spei[i,j,ids]=norm.ppf(np.clip(cdf,1e-6,1-1e-6))
    return spei.astype(np.float32),pars.astype(np.float32)

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

def generate_unit(lead, scale, run, source, log):
    lon,lat,mask=coordinates_and_mask(); train=np.arange(480)<(lead+348); full=np.ones(480,bool)
    pairs,names=read_members(lead,source,False); obs=read_obs(lead)
    members=[];params=[]
    for idx,(pr,t2) in enumerate(pairs):
        s,p=fit_spei(pr-thornthwaite(t2,lat),scale,train);s[~mask]=np.nan;p[~mask]=np.nan;members.append(s);params.append(p)
        log.info('stage=baseline lead=%d scale=%d member=%s status=done',lead,scale,names[idx])
    so,po=fit_spei(obs[0]-thornthwaite(obs[1],lat),scale,train);so[~mask]=np.nan
    out=ROOT/'outputs/production_output/spei'/run; arr=np.stack(members,axis=2)
    write_nc(out/f'SPEI{scale}_members_LT{lead}.nc',f'SPEI{scale}',arr,lon,lat,names,{'calibration':'training_only','source_run':source,'lead':lead,'timescale_months':scale})
    np.savez_compressed(out/f'SPEI{scale}_parameters_evaluation_LT{lead}.npz',member=np.stack(params,axis=2),observed=po,parameter_order=np.array(['alpha_scale','beta_shape','gamma_location']),fit='3-parameter log-logistic; unbiased probability weighted moments; calendar-month specific; training only')
    np.savez_compressed(out/f'SPEI{scale}_observed_LT{lead}.npz',spei=so)
    # Operational distribution parameters use full-period refitted pr/t2 products.
    fpairs,_=read_members(lead,FULL_RUN,True); fparams=[]
    for pr,t2 in fpairs:
        _,p=fit_spei(pr-thornthwaite(t2,lat),scale,full);p[~mask]=np.nan;fparams.append(p)
    _,fpo=fit_spei(obs[0]-thornthwaite(obs[1],lat),scale,full)
    np.savez_compressed(out/f'SPEI{scale}_parameters_operational_LT{lead}.npz',member=np.stack(fparams,axis=2),observed=fpo,parameter_order=np.array(['alpha_scale','beta_shape','gamma_location']),calibration='1982-01/2021-12')

def estimator(method,seed):
    if method=='Ridge1':return make_pipeline(StandardScaler(),RidgeCV(alphas=ALPHAS,cv=KFold(5,shuffle=False)))
    if method=='Ridge2':return make_pipeline(StandardScaler(),RidgeCV(alphas=ALPHAS,cv=None,gcv_mode='auto'))
    hidden=(32,) if method=='ANN1' else (32,16)
    base=make_pipeline(StandardScaler(),MLPRegressor(hidden_layer_sizes=hidden,max_iter=600,early_stopping=True,random_state=seed,batch_size=64))
    return GridSearchCV(base,{'mlpregressor__alpha':[1e-5,1e-4,1e-3],'mlpregressor__learning_rate_init':[5e-4,1e-3]},cv=3,n_jobs=1,scoring='neg_mean_absolute_error',refit=True)

def scores(obs,pred,sel):
    ok=sel&np.isfinite(obs)&np.isfinite(pred)
    if ok.sum()<12:return np.nan,np.nan
    return float(np.mean(np.abs(obs[ok]-pred[ok]))),float(np.corrcoef(obs[ok],pred[ok])[0,1]) if np.std(obs[ok]) and np.std(pred[ok]) else np.nan

def correct_unit(lead,scale,run,rank,size,log):
    out=ROOT/'outputs/production_output/spei'/run; lon,lat,mask=coordinates_and_mask()
    data=nc_array(out/f'SPEI{scale}_members_LT{lead}.nc',f'SPEI{scale}');obs=np.load(out/f'SPEI{scale}_observed_LT{lead}.npz')['spei']
    train=np.arange(480)<(lead+348);valid=(np.arange(480)>=lead+348)&(np.arange(480)<lead+468)
    pred=np.full(obs.shape,np.nan,np.float32);choice=np.full(obs.shape[:2],-1,np.int8);records=[];artifacts={}
    points=[tuple(x) for x in np.argwhere(mask)]; mine=points[rank::size]
    for n,(i,j) in enumerate(mine):
        X=data[i,j].T;y=obs[i,j]
        # A failed distribution fit in one member must not discard the other
        # valid members at this gridpoint. Keep every member with sufficient
        # training support and archive the exact subset/imputation statistics.
        keep=np.sum(train[:,None]&np.isfinite(X),axis=0)>=30
        if keep.sum()<2:continue
        X=X[:,keep];means=np.nanmean(np.where(train[:,None],X,np.nan),axis=0)
        okx=np.any(np.isfinite(X),axis=1);Xi=np.where(np.isfinite(X),X,means)
        oktr=train&np.isfinite(y)&okx;okall=okx
        if oktr.sum()<30:continue
        candidates=[]
        for mi,method in enumerate(METHODS):
            try:
                model=estimator(method,lead*100000+scale*1000+i*100+j);model.fit(Xi[oktr],y[oktr]);p=np.full(480,np.nan);p[okall]=model.predict(Xi[okall]);mae,corr=scores(y,p,valid)
                candidates.append((mi,mae,corr,p,model));records.append({'i':i,'j':j,'lead':lead,'scale':scale,'method':method,'n_members_used':int(keep.sum()),'validation_mae':mae,'validation_corr':corr,'status':'passed'})
            except Exception as e:records.append({'i':i,'j':j,'lead':lead,'scale':scale,'method':method,'status':'failed','reason':f'{type(e).__name__}: {e}'})
        good=[x for x in candidates if np.isfinite(x[1]) and np.isfinite(x[2])]
        if good:
            # Equal-weight rank aggregation enforces both requested criteria.
            maes=sorted(x[1] for x in good);cors=sorted((x[2] for x in good),reverse=True)
            best=min(good,key=lambda x:(maes.index(x[1])+cors.index(x[2]),x[1],-x[2]));choice[i,j]=best[0];pred[i,j]=best[3]
            artifacts[f'{i}_{j}']={'method':METHODS[best[0]],'model':best[4],'member_indices_zero_based':np.flatnonzero(keep),'imputation_training_means':means,'selection':'minimum equal-weight rank(MAE asc, CORR desc) on independent validation'}
        if n%25==0:log.info('stage=AI lead=%d scale=%d rank=%d progress=%d/%d',lead,scale,rank,n,len(mine))
    shard=out/'shards';shard.mkdir(exist_ok=True);stem=f'LT{lead}_S{scale}_R{rank:02d}'
    np.savez_compressed(shard/f'{stem}.npz',pred=pred,choice=choice);joblib.dump(artifacts,shard/f'{stem}.joblib',compress=3)
    with open(shard/f'{stem}.csv','w',newline='') as f:
        keys=sorted(set().union(*(r.keys() for r in records))) if records else ['status'];w=csv.DictWriter(f,fieldnames=keys);w.writeheader();w.writerows(records)

def categorical(obs,pred,threshold,sel):
    ok=np.broadcast_to(sel,obs.shape)&np.isfinite(obs)&np.isfinite(pred);o=obs[ok]<=threshold;p=pred[ok]<=threshold
    h=np.sum(o&p);m=np.sum(o&~p);f=np.sum(~o&p);c=np.sum(~o&~p)
    div=lambda a,b:float(a/b) if b else np.nan
    ans={'n':int(ok.sum()),'hits':int(h),'misses':int(m),'false_alarms':int(f),'correct_negatives':int(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)}
    if len(np.unique(o))==2:
        prob=norm.cdf(-pred[ok]);ans['AUC']=float(roc_auc_score(o,prob));ff,tt,_=roc_curve(o,prob);ans['roc_fpr']=ff.tolist();ans['roc_tpr']=tt.tolist()
    else:ans.update(AUC=np.nan,roc_fpr=[],roc_tpr=[])
    return ans

def finalize(run,size,log):
    import matplotlib;matplotlib.use('Agg');import matplotlib.pyplot as plt
    out=ROOT/'outputs/production_output/spei'/run;report=ROOT/'outputs/evaluation_output/spei'/run;report.mkdir(parents=True,exist_ok=True)
    lon,lat,mask=coordinates_and_mask();rows=[]
    for lead in LEADS:
      for scale in SCALES:
        pred=np.full((len(lon),len(lat),480),np.nan,np.float32);choice=np.full((len(lon),len(lat)),-1,np.int8)
        for rank in range(size):
            z=np.load(out/'shards'/f'LT{lead}_S{scale}_R{rank:02d}.npz');take=z['choice']>=0;pred[take]=z['pred'][take];choice[take]=z['choice'][take]
        obs=np.load(out/f'SPEI{scale}_observed_LT{lead}.npz')['spei'];train=np.arange(480)<(lead+348);valid=(np.arange(480)>=lead+348)&(np.arange(480)<lead+468)
        write_nc(out/f'SPEI{scale}_ens_LT{lead}.nc',f'SPEI{scale}_ens',pred[:,:,None,:],lon,lat,['best_AI'],{'selection':'rank MAE ascending plus CORR descending on validation','lead':lead,'timescale_months':scale})
        np.savez_compressed(out/f'SPEI{scale}_best_model_LT{lead}.npz',method_index=choice,methods=np.array(METHODS))
        for split,sel in [('training',train),('validation',valid)]:
          mae,corr=scores(obs[mask],pred[mask],sel)
          for level,thr in DROUGHT.items():rows.append({'lead':lead,'scale':scale,'split':split,'drought_level':level,'threshold':thr,'MAE':mae,'CORR':corr,**categorical(obs[mask],pred[mask],thr,sel)})
    flat=[]
    for r in rows:
        x=dict(r);x['roc_fpr']=json.dumps(x['roc_fpr']);x['roc_tpr']=json.dumps(x['roc_tpr']);flat.append(x)
    with open(report/'drought_verification.csv','w',newline='') as f:w=csv.DictWriter(f,fieldnames=list(flat[0]));w.writeheader();w.writerows(flat)
    figs=report/'figures';figs.mkdir(exist_ok=True)
    for scale in SCALES:
      for level in DROUGHT:
        rs=[r for r in rows if r['scale']==scale and r['drought_level']==level and r['lead'] in (1,3,5)];metrics=['POD','POFD','FAR','FBI','CSI']
        fig,axs=plt.subplots(5,2,figsize=(10,13),sharex=True,constrained_layout=True)
        for c,split in enumerate(('training','validation')):
          q=[r for r in rs if r['split']==split]
          for rr,m in enumerate(metrics):
            vals=[next(x[m] for x in q if x['lead']==lead) for lead in (1,3,5)];axs[rr,c].bar(['LT1','LT3','LT5'],vals,color='#3977a8');axs[rr,c].set_ylabel(m);axs[rr,c].grid(axis='y',alpha=.3);axs[rr,c].set_title(split.title() if rr==0 else '')
        fig.suptitle(f'SPEI-{scale} {level} drought verification');fig.savefig(figs/f'SPEI{scale}_{level}_verification.png',dpi=180);plt.close(fig)
    links=''.join(f'<section><h2>SPEI-{s} {d}</h2><img src="figures/SPEI{s}_{d}_verification.png" alt="SPEI-{s} {d} verification"></section>' for s in SCALES for d in DROUGHT)
    doc=f'<!doctype html><html><head><meta charset="utf-8"><title>Dr_SFC SPEI verification</title><style>body{{font:15px system-ui;margin:20px}}img{{max-width:100%}}section{{border:1px solid #bbb;padding:12px;margin:16px 0}}</style></head><body><h1>SPEI drought verification — run {html.escape(run)}</h1><p>Independent training/validation results, continuous time series; D1≤−1, D2≤−1.5, D3≤−2. Metrics: POD, POFD, FAR, FBI, CSI, ROC and AUC. <a href="drought_verification.csv">Source metrics</a>.</p>{links}</body></html>'
    (report/'index.html').write_text(doc);(out/'SUCCESS').write_text(datetime.now().isoformat());log.info('stage=finalize status=SUCCESS report=%s',report/'index.html')

def main():
    p=argparse.ArgumentParser();p.add_argument('stage',choices=['generate','correct','finalize']);p.add_argument('--run',required=True);p.add_argument('--source-run',default=EVAL_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();log=getlog(a.run)
    if a.stage=='generate':
        units=[(l,s) for l in LEADS for s in SCALES]
        for l,s in units[a.rank::a.size]:generate_unit(l,s,a.run,a.source_run,log)
    elif a.stage=='correct':
        for l in LEADS:
          for s in SCALES:correct_unit(l,s,a.run,a.rank,a.size,log)
    else:finalize(a.run,a.size,log)
if __name__=='__main__':main()