#!/usr/bin/env python3
"""Distributed Stage-1 AI bias correction for Dr_SFC (real data only)."""
from __future__ import annotations
import argparse, 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 sklearn.linear_model import RidgeCV
from sklearn.neural_network import MLPRegressor
from sklearn.exceptions import ConvergenceWarning
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")
RCMS=("RegCM","clWRF"); VARS=("pr","t2"); LEADS=range(1,7); MONTHS=range(1,13); FILL=np.float32(-99)
ALPHAS=np.array([.01,.1,1.,10.,100.])
SETS={"pr":{"SP1":("pr","q2","rh2","w700","olr"),"SP2":("pr","q2","w700","olr","u850","h500"),"SP3":("pr","olr","q2","rh2","slp","u850","u200","w700","h500","h850")},
"t2":{"ST1":("t2","h500","h850","slp","q2"),"ST2":("t2","h500","h850","w700","q2","rh2","slp"),"ST3":("t2","olr","q2","rh2","slp","u850","u200","w700","h500","h850")}}
METHODS=("Ridge1","Ridge2","ANN1","ANN2"); SPATIAL=("national","regional")
MEMBERS=[(m,s,p) for m in METHODS for s in range(3) for p in SPATIAL]
def logger(path):
 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("ai")
def read_var(ds,name):
 v=ds.variables[name];dims=[x.decode() if isinstance(x,bytes) else x for x in v.dimensions]
 a=np.ma.filled(v[:],np.nan).astype(np.float32);a[np.isclose(a,-99)]=np.nan
 if name == "pr": a[(a < 0) & np.isfinite(a)] = 0.0
 return np.transpose(a,[dims.index(x) for x in ("lon","lat","mon")])
def read_mask():
 with netcdf_file(MASK,"r",mmap=False) as d:
  v=d.variables["reg"];dims=[x.decode() if isinstance(x,bytes) else x for x in v.dimensions];a=np.array(v[:],copy=True)
 return (a.T if dims==["lat","lon"] else a).astype(np.int16)
def preflight(log):
 out=ROOT/"outputs/evaluation_output/preflight";out.mkdir(parents=True,exist_ok=True);rows=[];ref=None;obs=None;errors=[]
 required={"pr_o","t2_o","pr","t2","q2","rh2","w700","olr","u850","u200","h500","h850","slp"}
 for lead in LEADS:
  for rcm in RCMS:
   p=INPUT/f"lead{lead}_{rcm}_01_BL_Test_VN.nc";item={"lead":lead,"rcm":rcm,"path":str(p),"exists":p.exists()}
   if not p.exists():errors.append(f"missing {p}");rows.append(item);continue
   try:
    with netcdf_file(p,"r",mmap=False) as d:
     missing=sorted(required-set(d.variables));item["missing_variables"]=";".join(missing)
     if missing:errors.append(f"{p}: missing {missing}")
     lon=np.array(d.variables["lon"][:],copy=True);lat=np.array(d.variables["lat"][:],copy=True);shape=read_var(d,"pr").shape
     item.update(nlon=len(lon),nlat=len(lat),nmon=shape[-1])
     if shape[-1]!=480:errors.append(f"{p}: {shape[-1]} months")
     if ref is None:ref=(lon,lat)
     elif not(np.array_equal(lon,ref[0]) and np.array_equal(lat,ref[1])):errors.append(f"{p}: grid mismatch")
     oo=(read_var(d,"pr_o"),read_var(d,"t2_o"))
     if obs is None:obs=oo
     elif any(not np.allclose(a,b,equal_nan=True) for a,b in zip(obs,oo)):errors.append(f"{p}: observations differ")
     pr=read_var(d,"pr");rh=read_var(d,"rh2");item["negative_pr_replaced_by_zero"]=int(np.sum(pr<0));item["rh_outside_0_100"]=int(np.sum((rh<0)|(rh>100)))
     
   except Exception as e:errors.append(f"{p}: {e}")
   rows.append(item)
 import csv
 with open(out/"input_inventory.csv","w",newline="") as f:
  w=csv.DictWriter(f,fieldnames=sorted(set().union(*(x.keys() for x in rows))));w.writeheader();w.writerows(rows)
 if ref and read_mask().shape!=(len(ref[0]),len(ref[1])):errors.append("mask grid mismatch")
 report={"passed":not errors,"checked_at":datetime.now().isoformat(),"errors":errors,"files":len(rows)}
 (out/"input_validation.json").write_text(json.dumps(report,indent=2));log.info("Preflight passed=%s errors=%d",not errors,len(errors))
 if errors:raise SystemExit("Preflight failed: "+" | ".join(errors[:5]))
def climatology(rcm,var):
 with netcdf_file(ROOT/"Clim"/f"Monthly_Clim_{rcm}.nc","r",mmap=False) as d:
  m=np.ma.filled(d.variables[f"{var}_cli_{rcm}"][:],np.nan).astype(np.float32);o=np.ma.filled(d.variables[f"{var}_o_cli"][:],np.nan).astype(np.float32)
 m[np.isclose(m,-99)]=np.nan;o[np.isclose(o,-99)]=np.nan;return m,o
def neighbor_features(a):
 z=[]
 for di,dj in ((0,0),(-2,0),(2,0),(0,-2),(0,2)):
  b=np.full_like(a,np.nan);si=slice(max(0,-di),min(a.shape[0],a.shape[0]-di));sj=slice(max(0,-dj),min(a.shape[1],a.shape[1]-dj))
  b[si,sj,:]=a[slice(si.start+di,si.stop+di),slice(sj.start+dj,sj.stop+dj),:];z.append(b)
 return z
def ann_hidden_layers(method,n_features):
 first=max(2,int(np.ceil(n_features/2)))
 if first>=n_features:first=n_features-1
 if method=="ANN1":return (first,)
 second=max(1,int(np.ceil(first/2)))
 if second>=first:second=first-1
 return (first,second)
def estimator(method,seed,n_features):
 if method=="Ridge1":return make_pipeline(StandardScaler(),RidgeCV(alphas=ALPHAS,cv=5)),None
 if method=="Ridge2":return make_pipeline(StandardScaler(),RidgeCV(alphas=ALPHAS)),None
 hidden=ann_hidden_layers(method,n_features)
 model=MLPRegressor(hidden_layer_sizes=hidden,early_stopping=True,max_iter=500,validation_fraction=.15,n_iter_no_change=20,random_state=seed,batch_size=512)
 return make_pipeline(StandardScaler(),model),hidden
def fit_predict(X,y,train,test,method,regions,regional,seed):
 pred=np.full(y.shape,np.nan,np.float32);models={};diagnostics=[]
 for g in (range(1,8) if regional else (0,)):
  gm=(regions==g) if regional else np.ones(regions.shape,bool);tr=train&gm&np.isfinite(y)&np.all(np.isfinite(X),axis=1);te=test&gm&np.all(np.isfinite(X),axis=1)
  if tr.sum()<50:
   diagnostics.append({"group":int(g),"status":"failed","reason":"insufficient_training_samples","n_train":int(tr.sum()),"n_predict":int(te.sum())});continue
  model,hidden=estimator(method,seed+g,X.shape[1])
  try:
   with warnings.catch_warnings(record=True) as caught:
    warnings.simplefilter("always")
    model.fit(X[tr],y[tr])
   conv=[str(w.message) for w in caught if issubclass(w.category,ConvergenceWarning)]
   pred[te]=model.predict(X[te]).astype(np.float32);models[str(g)]=model
   diagnostics.append({"group":int(g),"status":"completed_with_warning" if conv else "passed","reason":"; ".join(conv),"n_train":int(tr.sum()),"n_predict":int(te.sum()),"n_features":int(X.shape[1]),"hidden_layers":list(hidden) if hidden else []})
  except Exception as exc:
   diagnostics.append({"group":int(g),"status":"failed","reason":f"{type(exc).__name__}: {exc}","n_train":int(tr.sum()),"n_predict":int(te.sum()),"n_features":int(X.shape[1]),"hidden_layers":list(hidden) if hidden else []})
 return pred,models,diagnostics
def metric(obs,pred,raw,var,split,member):
 ok=np.isfinite(obs)&np.isfinite(pred)&np.isfinite(raw)
 if not ok.any():return {"split":split,"member":member,"n":0,"error":np.nan,"corr":np.nan,"added_value":np.nan}
 o,p,r=obs[ok],pred[ok],raw[ok];err=float(np.mean(np.abs(p-o)));ref=float(np.mean(np.abs(r-o)))
 if var=="pr":den=max(float(np.mean(np.abs(o))),1e-6);err/=den;ref/=den
 return {"split":split,"member":member,"n":int(ok.sum()),"error":err,"corr":float(np.corrcoef(o,p)[0,1]),"added_value":(ref-err)/max(ref,err,1e-12)}
def work_unit(rcm,var,lead,month,run,log):
 path=INPUT/f"lead{lead}_{rcm}_01_BL_Test_VN.nc";items=list(SETS[var].items());needed=sorted(set(sum((list(v) for _,v in items),[])))
 with netcdf_file(path,"r",mmap=False) as d:
  lon=np.array(d.variables["lon"][:],copy=True);lat=np.array(d.variables["lat"][:],copy=True);data={v:read_var(d,v) for v in needed};obs=read_var(d,f"{var}_o")
 reg=read_mask();mcli,ocli=climatology(rcm,var);cli=mcli[...,lead-1,month-1];ocl=ocli[...,month-1];tids=np.arange(month-1,480,12)
 train_t=(tids>=lead)&(tids<lead+348);valid_t=(tids>=lead+348)&(tids<lead+468);raw=data[var][...,tids]
 if var=="t2":target=obs[...,tids]-ocl[...,None]
 else:target=np.divide(obs[...,tids],ocl[...,None],out=np.full_like(raw,np.nan),where=ocl[...,None]>1e-6)
 region3=np.broadcast_to(reg[...,None],raw.shape);train3=np.broadcast_to(train_t,raw.shape);valid3=np.broadcast_to(valid_t,raw.shape)
 preds=np.full((24,)+raw.shape,np.nan,np.float32);artifacts={};rows=[];mi=0
 for method in METHODS:
  for setname,names in items:
   feats=np.stack(sum((neighbor_features(data[n][...,tids]) for n in names),[]),axis=-1)
   for spatial in SPATIAL:
    X=feats.reshape(-1,feats.shape[-1]);y=target.reshape(-1);regs=region3.reshape(-1)
    if spatial=="regional":
     X=np.concatenate([X,np.broadcast_to(lon[:,None,None],raw.shape).reshape(-1,1),np.broadcast_to(lat[None,:,None],raw.shape).reshape(-1,1)],axis=1)
    pred,models,diagnostics=fit_predict(X,y,train3.reshape(-1),(train3|valid3).reshape(-1),method,regs,spatial=="regional",lead*1000+month*10+mi)
    bc=pred.reshape(raw.shape);bc=ocl[...,None]+bc if var=="t2" else ocl[...,None]*bc
    if var=="pr":bc=np.maximum(bc,0)
    preds[mi]=bc;name=f"{method}_{setname}_{spatial}";artifacts[name]={"models":models,"diagnostics":diagnostics,"n_features":int(X.shape[1]),"hidden_layers":list(ann_hidden_layers(method,X.shape[1])) if method.startswith("ANN") else []}
    status="failed" if any(d["status"]=="failed" for d in diagnostics) else ("completed_with_warning" if any(d["status"]=="completed_with_warning" for d in diagnostics) else "passed")
    reason=" | ".join(f"group={d['group']}: {d['reason']}" for d in diagnostics if d["reason"])
    log.info("Fit %s status=%s n_features=%d hidden=%s reason=%s",name,status,X.shape[1],artifacts[name]["hidden_layers"],reason or "none")
    for label,mask in (("training",train3),("validation",valid3)):
     row=metric(obs[...,tids][mask],bc[mask],raw[mask],var,label,name);row.update(status=status,reason=reason,n_features=int(X.shape[1]),hidden_layers="x".join(map(str,artifacts[name]["hidden_layers"]))) ;rows.append(row)
    mi+=1
 shard=ROOT/"outputs/production_output/shards"/run;shard.mkdir(parents=True,exist_ok=True);stem=f"{rcm}_{var}_LT{lead}_M{month:02d}"
 np.savez_compressed(shard/f"{stem}.npz",pred=preds,tids=tids,obs=obs[...,tids],raw=raw,train=train_t,valid=valid_t)
 art=ROOT/"outputs/production_output/model_artifacts"/run;art.mkdir(parents=True,exist_ok=True);joblib.dump({"rcm":rcm,"var":var,"lead":lead,"month":month,"members":MEMBERS,"models":artifacts},art/f"{stem}.joblib",compress=3)
 (shard/f"{stem}.json").write_text(json.dumps(rows));log.info("Completed %s",stem)
def worker(run,rank,size,log):
 units=[(r,v,l,m) for r in RCMS for v in VARS for l in LEADS for m in MONTHS];mine=units[rank::size];log.info("Worker units=%d world=%d",len(mine),size)
 for u in mine:work_unit(*u,run,log)
 log.info("Worker complete")
def write_nc(path,varname,data,lon,lat,members,attrs):
 path.parent.mkdir(parents=True,exist_ok=True)
 with netcdf_file(path,"w",version=1) as d:
  for n,k in (("lon",len(lon)),("lat",len(lat)),("member",len(members)),("time",480)):d.createDimension(n,k)
  for n,x in (("lon",lon),("lat",lat),("member",np.arange(1,len(members)+1)),("time",np.arange(1,481))):
   q=d.createVariable(n,"f",(n,));q[:]=x
  q=d.createVariable(varname,"f",("lon","lat","member","time"));q._FillValue=FILL;q.missing_value=FILL;q[:]=np.where(np.isfinite(data),data,FILL)
  d.member_names=";".join(members)
  for k,v in attrs.items():setattr(d,k,v)
def finalize(run,log):
 shard=ROOT/"outputs/production_output/shards"/run;out=ROOT/"outputs/production_output/bias_corrected"/run
 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)
 all_data={};rows=[]
 for rcm in RCMS:
  for var in VARS:
   names=[f"{m}_{list(SETS[var].keys())[s]}_{p}" for m,s,p in MEMBERS]
   for lead in LEADS:
    data=np.full((40,68,24,480),np.nan,np.float32)
    for month in MONTHS:
     stem=f"{rcm}_{var}_LT{lead}_M{month:02d}";z=np.load(shard/f"{stem}.npz");data[:,:,:,z["tids"]]=np.moveaxis(z["pred"],0,2)
     rr=json.loads((shard/f"{stem}.json").read_text())
     for x in rr:x.update(rcm=rcm,var=var,lead=lead,month=month);rows.append(x)
    all_data[(rcm,var,lead)]=data;write_nc(out/f"{rcm}_{var}_bc_LT{lead}.nc",f"{var}_bc",data,lon,lat,names,{"rcm":rcm,"lead":lead,"run_id":run})
 for var in VARS:
  for lead in LEADS:
   stack=np.concatenate([all_data[(r,var,lead)] for r in RCMS],axis=2);ens=np.nanmean(stack,axis=2,keepdims=True)
   write_nc(out/f"{var}_bc_ens_LT{lead}.nc",f"{var}_bc_ens",ens,lon,lat,["ensemble_mean"],{"lead":lead,"run_id":run})
 import csv
 ev=ROOT/"outputs/evaluation_output"/run;ev.mkdir(parents=True,exist_ok=True)
 with open(ev/"member_metrics.csv","w",newline="") as f:w=csv.DictWriter(f,fieldnames=sorted(rows[0]));w.writeheader();w.writerows(rows)
 (out/"SUCCESS").write_text(datetime.now().isoformat());log.info("Finalization complete")
def main():
 p=argparse.ArgumentParser();p.add_argument("mode",choices=["preflight","worker","finalize"]);p.add_argument("--run",required=True);p.add_argument("--rank",type=int);p.add_argument("--size",type=int);a=p.parse_args();log=logger(ROOT/"logs"/f"ai_bias_correction_{a.run}.log")
 if a.mode=="preflight":preflight(log)
 elif a.mode=="worker":worker(a.run,a.rank if a.rank is not None else int(os.getenv("SLURM_PROCID","0")),a.size if a.size is not None else int(os.getenv("SLURM_NTASKS","1")),log)
 else:finalize(a.run,log)
if __name__=="__main__":main()