#!/usr/bin/env python3
"""Stage-2 boosted-tree Tx/severity experiment and HTML verification report.

The workflow pools all Vietnam grid cells in one model (lat/lon are predictors),
uses 2003-2016 for training, 2017-2018 for feature selection, and 2019-2022
once for final testing.  Slurm ranks independently verify test initializations.
"""
from __future__ import annotations

import argparse, base64, glob, io, json, os, pickle, socket
from functools import lru_cache
from datetime import date
from pathlib import Path
import numpy as np
from netCDF4 import Dataset
import shapefile

BASE = Path(__file__).resolve().parent
DATA = Path("/home/tanpv/Data/QG.25.196")
RUN = BASE / "outputs" / "AI_Results"
BOUNDARY = Path("/home/tanpv/GrADS_SCR/VN_SubRegs1/Vietnam7.shp")
WEEKS = ((0, 11), (11, 18), (18, 25), (25, 32), (32, 39), (39, 46))
ALL = ["tp","cp","tx","tm","t2m","td2m","ts","st20","st100","sm20","sm100",
       "msl","u10","v10","w","cape","pwat","tcl","sh","le","sw","lw","swn",
       "lwn","olr","dz15","dz18","dz87","dz95","dz97","ivtu","ivtv","ivt",
       "dt85","dt98","dt75","dtd2m","srn","bowen","uv10"]
CORE = ["tx","t2m","tm","td2m","dtd2m","ts","st20","sm20","msl","u10","v10",
        "pwat","tcl","srn","sh","le","dz15","sev_mod"]
STATIC = ["lead","doy_sin","doy_cos","lat","lon"]
SETS = ("core", "full", "selected")
APPROACHES = ("tx", "sev")

def doy365(d):
    n=d.timetuple().tm_yday-1; leap=d.year%4==0 and (d.year%100!=0 or d.year%400==0)
    if d.month==2 and d.day==29: return 58
    return n-int(leap and d.month>2)

def paths(years):
    out=[]
    for y in years: out += sorted(glob.glob(str(DATA/f"EC_Z_score/{y}/Z_score_EC_*.nc")))
    return out

def stamp(path): return Path(path).stem.replace("Z_score_EC_", "")

@lru_cache(maxsize=1)
def context():
    with Dataset(DATA/"OBS/Mask_VN_S2S_EC_0.25.nc") as d:
        lon=np.asarray(d["lon"][:]); lat=np.asarray(d["lat"][:]); a=np.ma.asarray(d["vn"][:])
        mask=~np.ma.getmaskarray(a)&np.isfinite(np.ma.filled(a,np.nan))
    yy,xx=np.where(mask)
    with Dataset(BASE/"outputs/Cli/Qp_tx_mod.nc") as d: qm=np.ma.filled(d["tx_thres"][:],np.nan)
    with Dataset(BASE/"outputs/Cli/Qp_tx_obs.nc") as d: qo=np.ma.filled(d["tx_thres"][:],np.nan)
    return lon,lat,mask,yy,xx,qm,qo

def read_rows(path, point_ids=None, features=ALL):
    lon,lat,mask,yy,xx,qm,qo=context(); s=stamp(path); init=date.fromisoformat(s)
    doys=np.array([doy365(date.fromordinal(init.toordinal()+i)) for i in range(46)])
    if point_ids is not None: yy,xx=yy[point_ids],xx[point_ids]
    npnt=len(yy); lead=np.repeat(np.arange(46),npnt); py=np.tile(yy,46); px=np.tile(xx,46)
    with Dataset(path) as d:
        cols=[]
        for f in features:
            if f=="sev_mod": continue
            a=np.ma.filled(d[f][:],np.nan); cols.append(a[:,yy,xx].reshape(-1))
        txz=np.ma.filled(d["tx"][:],np.nan)[:,yy,xx]
    raw=DATA/f"EC_Combined/{init.year}/EC_all_{s}.nc"
    obs=DATA/f"OBS_4EC/{init.year}/Obs_{s}.nc"
    with Dataset(raw) as d: txm=np.ma.filled(d["mx2t6"][:46],np.nan)[:,yy,xx]
    with Dataset(obs) as d: txo=np.ma.filled(d["tasmax"][:46],np.nan)[:,yy,xx]
    sevmod=txm-qm[doys,:,yy,xx].transpose(1,0,2)[:,np.arange(npnt),np.arange(npnt)] if False else None
    # Explicit indexing avoids NumPy's mixed advanced-index axis surprises.
    thm=np.empty_like(txm); tho=np.empty_like(txo)
    for k in range(46): thm[k]=qm[doys[k],k,yy,xx]; tho[k]=qo[doys[k],yy,xx]
    sevmod=txm-thm
    fmap={f:c for f,c in zip([x for x in features if x!="sev_mod"],cols)}
    if "sev_mod" in features: fmap["sev_mod"]=sevmod.reshape(-1)
    static={"lead":lead/45.,"doy_sin":np.repeat(np.sin(2*np.pi*doys/365),npnt),
            "doy_cos":np.repeat(np.cos(2*np.pi*doys/365),npnt),
            "lat":lat[py],"lon":lon[px]}
    fmap.update(static)
    ytx=txo.reshape(-1); ysev=(txo-tho).reshape(-1)
    return fmap,ytx,ysev,lead,py,px

def matrix(fmap,names): return np.column_stack([fmap[n] for n in names]).astype("f4")

def sample_data(years, npoint, seed, features):
    _,_,_,yy,_,_,_=context(); rng=np.random.RandomState(seed); xs=[]; yt=[]; ys=[]
    for i,p in enumerate(paths(years)):
        ids=rng.choice(len(yy),min(npoint,len(yy)),replace=False)
        f,a,b,*_=read_rows(p,ids,features); xs.append(matrix(f,features+STATIC)); yt.append(a); ys.append(b)
        if (i+1)%100==0: print(f"sampled {i+1}/{len(paths(years))}",flush=True)
    X=np.concatenate(xs); a=np.concatenate(yt); b=np.concatenate(ys)
    good=np.isfinite(X).all(1)&np.isfinite(a)&np.isfinite(b)
    return X[good],a[good],b[good]

def booster(X,y,names,seed):
    from xgboost import XGBRegressor
    m=XGBRegressor(n_estimators=500,max_depth=8,learning_rate=.05,subsample=.8,
       colsample_bytree=.8,min_child_weight=20,reg_lambda=2,n_jobs=20,
       objective="reg:squarederror",tree_method="hist",random_state=seed)
    m.fit(X,y); return m

def score(y,p):
    e=p-y; return {"mae":float(np.mean(abs(e))),"rmse":float(np.sqrt(np.mean(e*e))),"me":float(np.mean(e))}

def train(args):
    RUN.mkdir(parents=True,exist_ok=True); (RUN/"parts").mkdir(exist_ok=True)
    full=ALL+["sev_mod"]+STATIC; core=CORE+STATIC
    X,y_tx,y_sev=sample_data(range(2003,2017),args.points_per_init,args.seed,ALL+["sev_mod"])
    V,vtx,vsev=sample_data(range(2017,2019),args.validation_points,args.seed+1,ALL+["sev_mod"])
    pos={n:i for i,n in enumerate(full)}; meta={"engine":"XGBoost 2.1.4 (LightGBM unavailable on cluster)",
      "train_years":"2003-2016","validation_years":"2017-2018","test_years":"2019-2022",
      "points_per_training_initialization":args.points_per_init,"seed":args.seed,"models":{}}
    for approach,y,vy in (("tx",y_tx,vtx),("sev",y_sev,vsev)):
        rankings=None
        for fs,names in (("core",core),("full",full)):
            ix=[pos[n] for n in names]; m=booster(X[:,ix],y,names,args.seed)
            pred=m.predict(V[:,ix]); imp=dict(zip(names,map(float,m.feature_importances_)))
            with open(RUN/f"model_{approach}_{fs}.pkl","wb") as f: pickle.dump(m,f)
            meta["models"][f"{approach}_{fs}"]={"features":names,"validation":score(vy,pred),"importance":imp}
            if fs=="full": rankings=sorted(names,key=imp.get,reverse=True)
        selected=list(dict.fromkeys(STATIC+["tx","sev_mod"]+rankings[:args.selected_features]))
        ix=[pos[n] for n in selected]; m=booster(X[:,ix],y,selected,args.seed)
        pred=m.predict(V[:,ix]); imp=dict(zip(selected,map(float,m.feature_importances_)))
        with open(RUN/f"model_{approach}_selected.pkl","wb") as f: pickle.dump(m,f)
        meta["models"][f"{approach}_selected"]={"features":selected,"validation":score(vy,pred),"importance":imp}
    with open(RUN/"metadata.json","w") as f: json.dump(meta,f,indent=2)
    print("training complete",flush=True)

def init_acc(shape):
    return {k:np.zeros((6,)+shape,"f8") for k in ("n","sum","abs","sq","hit","miss","fa","cn")}

def worker(args):
    rank=int(os.environ.get("SLURM_PROCID",0)); size=int(os.environ.get("SLURM_NTASKS",1))
    files=paths(range(2019,2023))[rank::size]; lon,lat,mask,yy,xx,_,_=context()
    with open(RUN/"metadata.json") as f: meta=json.load(f)
    models={}; acc={}
    for a in APPROACHES:
      for fs in SETS:
        key=f"{a}_{fs}"; models[key]=pickle.load(open(RUN/f"model_{key}.pkl","rb")); acc[key]=init_acc(mask.shape)
    needed=ALL+["sev_mod"]
    for i,p in enumerate(files):
      fmap,yt,ys,lead,py,px=read_rows(p,None,needed)
      for key,m in models.items():
        target=yt if key.startswith("tx_") else ys; names=meta["models"][key]["features"]
        pred=m.predict(matrix(fmap,names)); sev=pred-(target-ys) if key.startswith("tx_") else pred
        valid=np.isfinite(pred)&np.isfinite(target); ev=ys>=0; fv=sev>=0
        for wi,(i0,i1) in enumerate(WEEKS):
          use=valid&(lead>=i0)&(lead<i1); e=pred-target; A=acc[key]
          np.add.at(A["n"][wi],(py[use],px[use]),1); np.add.at(A["sum"][wi],(py[use],px[use]),e[use])
          np.add.at(A["abs"][wi],(py[use],px[use]),abs(e[use])); np.add.at(A["sq"][wi],(py[use],px[use]),e[use]**2)
          for n,c in (("hit",fv&ev),("miss",~fv&ev),("fa",fv&~ev),("cn",~fv&~ev)):
            z=use&c; np.add.at(A[n][wi],(py[z],px[z]),1)
      print(f"rank {rank}: {i+1}/{len(files)}",flush=True)
    flat={f"{key}__{n}":v for key,A in acc.items() for n,v in A.items()}
    np.savez_compressed(RUN/"parts"/f"part_{rank:03d}.npz",**flat)

def div(a,b): return np.divide(a,b,out=np.full_like(a,np.nan),where=b!=0)

def finalize(args):
    parts=sorted((RUN/"parts").glob("part_*.npz")); expected=int(os.environ.get("EXPECTED_TASKS",args.tasks))
    if len(parts)!=expected: raise RuntimeError(f"Expected {expected} worker parts, found {len(parts)}")
    total={}
    for p in parts:
      with np.load(p) as d:
       for k in d.files: total[k]=total.get(k,0)+d[k]
    lon,lat,mask,*_=context(); metrics={}
    for a in APPROACHES:
      for fs in SETS:
       key=f"{a}_{fs}"; A={n:total[f"{key}__{n}"] for n in ("n","sum","abs","sq","hit","miss","fa","cn")}
       h,m,f,c=A["hit"],A["miss"],A["fa"],A["cn"]
       metrics[key]={"me":div(A["sum"],A["n"]),"mae":div(A["abs"],A["n"]),"rmse":np.sqrt(div(A["sq"],A["n"])),
          "pod":div(h,h+m),"far":div(f,h+f),"csi":div(h,h+m+f),"fbi":div(h+f,h+m)}
    with Dataset(RUN/"ehd_ai_verification.nc","w") as d:
      d.createDimension("week",6); d.createDimension("lat",len(lat)); d.createDimension("lon",len(lon))
      d.createVariable("week","i1",("week",))[:]=range(1,7); d.createVariable("lat","f4",("lat",))[:]=lat; d.createVariable("lon","f4",("lon",))[:]=lon
      for key,M in metrics.items():
       for n,a in M.items():
        v=d.createVariable(f"{key}_{n}","f4",("week","lat","lon"),zlib=True,fill_value=np.float32(9.96921e36))
        v[:]=np.ma.masked_invalid(np.where(mask,a,np.nan))
    make_html(metrics,lon,lat,mask)

def make_html(metrics,lon,lat,mask):
    import matplotlib; matplotlib.use("Agg"); import matplotlib.pyplot as plt
    with open(RUN/"metadata.json") as f: meta=json.load(f)
    rows=[]; cards=[]
    for key,M in metrics.items():
      vals={n:float(np.nanmean(np.where(mask,a,np.nan))) for n,a in M.items()}
      rows.append(f"<tr><td>{key}</td>"+"".join(f"<td>{vals[n]:.3f}</td>" for n in ("mae","rmse","me","pod","far","csi","fbi"))+"</tr>")
      for wi in range(6):
       a=np.where(mask,M["csi"][wi],np.nan); fig,ax=plt.subplots(figsize=(4.5,5)); im=ax.pcolormesh(lon,lat,a,shading="auto",vmin=0,vmax=1,cmap="viridis")
       for shape in shapefile.Reader(str(BOUNDARY)).shapes():
        pts=np.asarray(shape.points); cuts=list(shape.parts)+[len(pts)]
        for i0,i1 in zip(cuts[:-1],cuts[1:]):
         ax.plot(pts[i0:i1,0],pts[i0:i1,1],color="black",linewidth=.65,zorder=3)
       ax.set(title=f"{key}: W{wi+1} CSI",xlabel="Longitude",ylabel="Latitude",xlim=(101.5,110.5),ylim=(7.5,24.0))
       ax.set_aspect("equal",adjustable="box"); fig.colorbar(im,ax=ax); b=io.BytesIO(); fig.savefig(b,format="png",dpi=110,bbox_inches="tight"); plt.close(fig)
       cards.append(f"<div class=card><img src='data:image/png;base64,{base64.b64encode(b.getvalue()).decode()}'></div>")
    imp=[]
    for key,v in meta["models"].items():
      top=sorted(v["importance"].items(),key=lambda z:z[1],reverse=True)[:15]
      imp.append(f"<h3>{key}</h3><p>{', '.join(f'{n} ({x:.3f})' for n,x in top)}</p>")
    html=f"""<!doctype html><html><head><meta charset=utf-8><title>Stage 2 EHD AI</title><style>body{{font:15px system-ui;max-width:1250px;margin:auto;padding:25px}}table{{border-collapse:collapse}}td,th{{border:1px solid #bbb;padding:6px}}.grid{{display:grid;grid-template-columns:repeat(3,1fr);gap:8px}}img{{width:100%}}.note{{background:#eef4fa;padding:12px}}</style></head><body>
    <h1>Stage 2: AI prediction of extreme-heat days</h1><div class=note><b>Design.</b> Two regressions are compared: direct Tx followed by the observed Tx threshold, and direct severity followed by zero. Feature sets are physical core, all variables, and importance-selected. Training: {meta['train_years']}; feature selection: {meta['validation_years']}; untouched test: {meta['test_years']}. Initializations—not individual rows—define the temporal split. Engine: {meta['engine']}. Training sampled {meta['points_per_training_initialization']} Vietnam points per initialization (seed {meta['seed']}); final verification uses every valid point and day. Threshold is max(local 31-day-window Q90, 35°C). Scores below are spatial means across W1-W6; MAE/RMSE/ME concern the regression target (°C), while POD/FAR/CSI/FBI concern EHD occurrence.</div>
    <h2>Test summary</h2><table><tr><th>approach_set</th><th>MAE</th><th>RMSE</th><th>ME</th><th>POD</th><th>FAR</th><th>CSI</th><th>FBI</th></tr>{''.join(rows)}</table>
    <h2>Feature screening</h2>{''.join(imp)}<p>Importance is normalized tree gain and can share credit unpredictably among correlated variables; it was computed without test years. Selected models always retain lead, annual-cycle terms, coordinates, model Tx and model severity.</p>
    <h2>Weekly CSI maps</h2><div class=grid>{''.join(cards)}</div>
    <h2>Interpretation cautions</h2><p>This first boosted-tree benchmark is a national pooled point model, not a spatial neural network. The EOF fields listed in the assignment were not supplied as per-sample modal coefficients and are therefore excluded from this first screening. A single fixed test period supports fair approach comparison but not uncertainty estimates; blocked cross-validation can be added after this benchmark.</p></body></html>"""
    (RUN/"ehd_ai_results.html").write_text(html)
    # README requires the presentation HTML at the project root as well.
    (BASE/"ehd_ai_results.html").write_text(html)

def preflight(args):
    RUN.mkdir(parents=True,exist_ok=True); (RUN/"parts").mkdir(exist_ok=True)
    for p in (BASE/"outputs/Cli/Qp_tx_mod.nc",BASE/"outputs/Cli/Qp_tx_obs.nc"): assert p.exists(),p
    assert len(paths(range(2003,2023)))==2080
    print(json.dumps({"host":socket.gethostname(),"files":2080,"output":str(RUN)},indent=2))

def main():
    p=argparse.ArgumentParser(); p.add_argument("command",choices=("preflight","train","worker","finalize")); p.add_argument("--points-per-init",type=int,default=24); p.add_argument("--validation-points",type=int,default=48); p.add_argument("--selected-features",type=int,default=18); p.add_argument("--seed",type=int,default=20260827); p.add_argument("--tasks",type=int,default=120); a=p.parse_args()
    globals()[a.command](a)
if __name__=="__main__": main()