#!/usr/bin/env python3
"""Dr_SFC task 1f: operational pr/t2 correction, SPEI drought forecast, and web products."""
from __future__ import annotations
import argparse, calendar, html, json, logging, sys, warnings
from datetime import date, datetime
from pathlib import Path
import joblib
import numpy as np
from scipy.io import netcdf_file
from scipy.stats import norm
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
from matplotlib.colors import BoundaryNorm, ListedColormap
import shapefile

import ai_bias_correction as bc
import spei_pipeline as spei

ROOT=Path('/home/tanpv/Progs/CLI/Dr')
OPE_INPUT=Path('/home/tanpv/Data/562_VNU/0.25_VN/For_BC_Ope')
ART=ROOT/'outputs/production_output/model_artifacts_full_period/1399282'
SPEI_BASE=ROOT/'outputs/production_output/spei_members/1399302'
FILL=np.float32(-99.0); METHODS=('Ridge1','Ridge2','ANN1','ANN2')
DAY_MAP={1:1,2:6,3:11,4:16,5:21,6:26}
SHAPE_FILE=Path('/home/tanpv/GrADS_SCR/VN_SubRegs1/Vietnam7.shp')
MAP_EXTENT=(101.75,110.25,7.75,24.25)
TEMP_LEVELS=np.arange(10.,35.,2.)
TEMP_ANOM_LEVELS=np.array([-5.,-4.,-3.,-2.,-1.,-.5,.5,1.,2.,3.,4.,5.])
TEMP_ANOM_COLORS=('#08306b','#2171b5','#4292c6','#6baed6','#c6dbef','#ffffff','#fee8c8','#fdbb84','#fc8d59','#e34a33','#b30000')
RAIN_LEVELS=np.arange(0.,701.,50.)
RAIN_ANOM_LEVELS=np.array([-100.,-75.,-50.,-25.,-5.,5.,25.,50.,75.,100.])
RAIN_ANOM_COLORS=('#b30000','#e34a33','#fc8d59','#fee8c8','#ffffff','#deebf7','#9ecae1','#4292c6','#08519c')
SPEI_BOUNDARIES=np.array([-2.,-1.5,-1.,-.5,0.,.5,1.,1.5,2.])
SPEI_TICKS=SPEI_BOUNDARIES.copy()
PROB_LEVELS=np.arange(0.,101.,10.)
SPEI_COLORS=["#67000d","#cb181d","#fb6a4a","#fcae91","#deebf7","#9ecae1","#4292c6","#084594"]

def shape_segments(path=SHAPE_FILE):
 if not path.exists():raise FileNotFoundError(f'missing required map boundary: {path}')
 reader=shapefile.Reader(str(path));segments=[]
 for shape in reader.shapes():
  points=np.asarray(shape.points)
  for start,stop in zip(shape.parts,list(shape.parts[1:])+[len(points)]):segments.append(points[start:stop])
 return segments

def decorate_map(ax,boundary):
 for segment in boundary:ax.plot(segment[:,0],segment[:,1],color='black',linewidth=.45,zorder=5)
 ax.set_xlim(MAP_EXTENT[:2]);ax.set_ylim(MAP_EXTENT[2:]);ax.set_aspect(1/np.cos(np.deg2rad(np.mean(MAP_EXTENT[2:]))),adjustable='box')
 ax.set_xticks(np.arange(102,111,2));ax.set_yticks(np.arange(8,25,2))
 ax.set_xlabel('Longitude (°E)');ax.set_ylabel('Latitude (°N)');ax.grid(alpha=.18,linewidth=.35)

def discrete_cmap(name,levels):
 cmap=ListedColormap(name) if isinstance(name,(tuple,list)) else plt.get_cmap(name,len(levels)-1)
 return cmap,BoundaryNorm(levels,len(levels)-1,clip=True)

def getlog(tag):
 p=ROOT/'logs'/f'operational_{tag}.log';p.parent.mkdir(exist_ok=True)
 logging.basicConfig(level=logging.INFO,format='%(asctime)s %(levelname)s %(message)s',handlers=[logging.FileHandler(p),logging.StreamHandler(sys.stdout)])
 return logging.getLogger('operational')

def add_month(y,m,n):
 q=y*12+m-1+n;return q//12,q%12+1

def forecast_anomalies(ens,clim_obs,year,month):
 """Return precipitation percent anomalies and temperature differences."""
 pr=[];t2=[]
 for lead in range(6):
  target_month=add_month(year,month,lead+1)[1]
  pr_cli=clim_obs['pr'][...,target_month-1]
  with np.errstate(divide='ignore',invalid='ignore'):
   pr.append(np.divide(
    (ens['pr'][...,lead]-pr_cli)*100.,pr_cli,
    out=np.full_like(pr_cli,np.nan,dtype=np.float32),where=pr_cli!=0))
  t2.append(ens['t2'][...,lead]-clim_obs['t2'][...,target_month-1])
 return {'pr':np.stack(pr,axis=-1).astype(np.float32),
         't2':np.stack(t2,axis=-1).astype(np.float32)}

def read_field(ds,name):
 aliases={'t2':('tas','t2'),'olr':('olr','Olr'),'slp':('slp','Slp')}
 key=next((x for x in aliases.get(name,(name,)) if x in ds.variables),None)
 if key is None:raise KeyError(f'missing operational variable {name}')
 v=ds.variables[key];dims=[x.decode() if isinstance(x,bytes) else x for x in v.dimensions]
 a=np.ma.filled(v[:],np.nan).astype(np.float32)
 # Historical files may carry singleton vertical or member axes.
 extra=[i for i,d in enumerate(dims) if d not in ('lon','lat','time')]
 for i in reversed(extra):
  if a.shape[i] != 1:raise ValueError('%s: unsupported dimension %s=%s' % (name,dims[i],a.shape[i]))
  a=np.squeeze(a,axis=i);dims.pop(i)
 if sorted(dims) != ['lat','lon','time']:
  raise ValueError('%s: expected lon/lat/time dimensions, got %s' % (name,tuple(dims)))
 a=np.transpose(a,[dims.index(x) for x in ('lon','lat','time')])
 units=str(getattr(v,'units','')).lower()
 if name=='t2' and ('k'==units.strip() or 'kelvin' in units or np.nanmedian(a)>100):a-=273.15
 if name=='pr' and ('s-1' in units or '/s' in units):
  # Each time index is a target calendar month; convert kg m-2 s-1 to mm/month.
  # Month lengths are applied later because target month depends on initialization.
  pass
 a[(a<=-90)&np.isfinite(a)]=np.nan
 return a,units

def mask_and_grid(ds):
 lon=np.array(ds.variables['lon'][:],copy=True);lat=np.array(ds.variables['lat'][:],copy=True)
 reg=bc.read_mask();return lon,lat,reg,(reg>=1)&(reg<=7)

def preflight(year,month,cycle,log):
 src=OPE_INPUT/f'{year:04d}'/f'{month:02d}';errors=[];inventory=[];ref=None
 required=('pr','t2','h500','h850','olr','q2','rh2','slp','u200','u850','w700')
 for rcm in bc.RCMS:
  p=src/f'{rcm}_{cycle:02d}.nc';row={'rcm':rcm,'path':str(p),'exists':p.exists()}
  if not p.exists():errors.append(f'missing {p}');inventory.append(row);continue
  try:
   with netcdf_file(p,'r',mmap=False) as d:
    lon,lat,reg,mask=mask_and_grid(d);row.update(nlon=len(lon),nlat=len(lat),nlead=d.variables['time'].shape[0])
    if (len(lon),len(lat),d.variables['time'].shape[0])!=(40,68,6):errors.append(f'{p}: expected 40x68x6')
    for name in required:read_field(d,name)
    if ref is None:ref=(lon,lat)
    elif not(np.array_equal(lon,ref[0]) and np.array_equal(lat,ref[1])):errors.append('RCM grid mismatch')
    if reg.shape!=(40,68):errors.append('mask grid mismatch')
  except Exception as e:errors.append(f'{p}: {type(e).__name__}: {e}')
  inventory.append(row)
 if len(list(ART.glob('*.joblib')))!=288:errors.append('full-period artifact count is not 288')
 if not (SPEI_BASE/'SUCCESS_49_MEMBERS').exists():errors.append('SPEI 49-member source is incomplete')
 report={'passed':not errors,'checked_at':datetime.now().isoformat(),'init_date':f'{year:04d}-{month:02d}-{DAY_MAP[cycle]:02d}','inventory':inventory,'errors':errors}
 out=ROOT/'outputs/Ope'/f'{year:04d}'/f'{month:02d}'/f'{DAY_MAP[cycle]:02d}';out.mkdir(parents=True,exist_ok=True)
 (out/'preflight.json').write_text(json.dumps(report,indent=2));log.info('stage=preflight passed=%s errors=%d',not errors,len(errors))
 if errors:raise SystemExit('Preflight failed: '+' | '.join(errors))
 return out

def predict_member(model_info,X,reg,regional):
 out=np.full(reg.shape,np.nan,np.float32);flat=X.reshape(-1,X.shape[-1]);r=reg.reshape(-1)
 if regional:
  lon=np.broadcast_to(G_LON[:,None],reg.shape).reshape(-1,1);lat=np.broadcast_to(G_LAT[None,:],reg.shape).reshape(-1,1);flat=np.concatenate([flat,lon,lat],axis=1)
 for group,model in model_info['models'].items():
  take=(r==int(group)) if regional else np.ones(r.shape,bool);take&=np.all(np.isfinite(flat),axis=1)
  if take.any():out.reshape(-1)[take]=model.predict(flat[take]).astype(np.float32)
 return out

def bias_correct(year,month,cycle,out,log):
 global G_LON,G_LAT
 products={v:np.full((40,68,48,6),np.nan,np.float32) for v in bc.VARS}
 for ri,rcm in enumerate(bc.RCMS):
  p=OPE_INPUT/f'{year:04d}'/f'{month:02d}'/f'{rcm}_{cycle:02d}.nc'
  with netcdf_file(p,'r',mmap=False) as d:
   G_LON,G_LAT,reg,mask=mask_and_grid(d);raw={}
   needed=set(sum((list(x) for sets in bc.SETS.values() for x in sets.values()),[]))
   for n in needed:raw[n]=read_field(d,n)[0]
   pr_units=read_field(d,'pr')[1]
  if ('s-1' in pr_units or '/s' in pr_units) and np.nanmedian(raw['pr']) < 0.01:
   for lead in range(1,7):
    ty,tm=add_month(year,month,lead);raw['pr'][...,lead-1]*=86400*calendar.monthrange(ty,tm)[1]
  for var in bc.VARS:
   mcli,ocli=bc.climatology(rcm,var);items=list(bc.SETS[var].items())
   for lead in range(1,7):
    ty,tm=add_month(year,month,lead);mo=tm
    artifact=joblib.load(ART/f'{rcm}_{var}_LT{lead}_M{mo:02d}.joblib')
    member=0
    for method in bc.METHODS:
     for setname,names in items:
      feats=np.stack(sum((bc.neighbor_features(raw[n][...,lead-1:lead]) for n in names),[]),axis=-1)[...,0,:]
      for spatial in bc.SPATIAL:
       name=f'{method}_{setname}_{spatial}';pred=predict_member(artifact['models'][name],feats,reg,spatial=='regional')
       pred=ocli[...,mo-1]+pred if var=='t2' else ocli[...,mo-1]*pred
       if var=='pr':pred=np.maximum(pred,0)
       pred[~mask]=np.nan;products[var][:,:,ri*24+member,lead-1]=pred;member+=1
    log.info('stage=BC rcm=%s var=%s LT=%d target=%04d-%02d members=24 status=SUCCESS',rcm,var,lead,ty,tm)
 ens={v:np.nanmean(products[v],axis=2) for v in bc.VARS}
 clim_obs={};
 for var in bc.VARS:clim_obs[var]=bc.climatology('RegCM',var)[1]
 ano=forecast_anomalies(ens,clim_obs,year,month)
 write_prt2(out/'Ope_Pr_T2_and_Ano.nc',products,ens,ano,year,month,DAY_MAP[cycle])
 return products,ens,ano,clim_obs

def write_prt2(path,members,ens,ano,year,month,day):
 with netcdf_file(path,'w') as d:
  for n,k in [('lon',40),('lat',68),('lead',6),('member',48)]:d.createDimension(n,k)
  for n,x in [('lon',G_LON),('lat',G_LAT),('lead',np.arange(1,7)),('member',np.arange(1,49))]:d.createVariable(n,'f',(n,))[:]=x
  for n,a,u in [('pr_fcs',ens['pr'],'mm month-1'),('t2_fcs',ens['t2'],'degree_C'),('pr_ano',ano['pr'],'percent'),('t2_ano',ano['t2'],'degree_C')]:
   q=d.createVariable(n,'f',('lon','lat','lead'));q._FillValue=FILL;q.missing_value=FILL;q.units=u;q[:]=np.where(np.isfinite(a),a,FILL)
  for n,a,u in [('pr_members',members['pr'],'mm month-1'),('t2_members',members['t2'],'degree_C')]:
   q=d.createVariable(n,'f',('lon','lat','member','lead'));q._FillValue=FILL;q.missing_value=FILL;q.units=u;q[:]=np.where(np.isfinite(a),a,FILL)
  d.init_date=f'{year:04d}-{month:02d}-{day:02d}';d.source_model_artifact_run='1399282';d.ensemble='equal mean of 48 full-period bias-correction members';d.target_month='initial month plus lead'

def loglog_spei(x,pars):
 alpha,beta,gamma=[pars[...,k] for k in range(3)];pos=x-gamma
 with np.errstate(all='ignore'):cdf=np.where((pos>0)&(alpha>0)&(beta>0),1/(1+(alpha/pos)**beta),np.nan)
 return norm.ppf(np.clip(cdf,1e-6,1-1e-6)).astype(np.float32)

def operational_pet(temp, target_month, climatology):
 # Thornthwaite heat index requires a 12-month cycle. Use observed climatology
 # for I/a and the forecast temperature for the target month.
 positive=np.maximum(climatology,0);I=np.nansum((positive/5.)**1.514,axis=-1)
 aa=6.75e-7*I**3-7.71e-5*I**2+1.792e-2*I+.49239
 days=calendar.monthrange(2001,target_month)[1]
 factor=(spei.daylength_hours(G_LAT[None,:],target_month)/12.)*(days/30.)
 ratio=np.divide(10*np.maximum(temp,0),I,out=np.zeros_like(temp),where=I>0)
 return np.where(np.isfinite(temp),np.where(temp>0,16*factor*np.power(ratio,aa),0),np.nan)

def drought(products,year,month,out,log):
 # Operational files do not contain analyzed antecedent months. For incomplete
 # scale windows use observed climatological water balance, explicitly recorded.
 lat=G_LAT;all_spei=np.full((40,68,49,3,6),np.nan,np.float32);scales=(1,3,6)
 obs_pr=bc.climatology('RegCM','pr')[1];obs_t=bc.climatology('RegCM','t2')[1]
 history=[]
 for back in range(-5,1):
  _,mo=add_month(year,month,back);history.append(obs_pr[...,mo-1]-operational_pet(obs_t[...,mo-1],mo,obs_t))
 hist=np.stack(history,axis=-1)
 pet=np.empty_like(products['t2'])
 for lead in range(1,7):
  _,mo=add_month(year,month,lead)
  for member in range(48):pet[...,member,lead-1]=operational_pet(products['t2'][...,member,lead-1],mo,obs_t)
 balances=products['pr']-pet
 for si,scale in enumerate(scales):
  vals=np.concatenate([np.broadcast_to(hist[:,:,None,:],(40,68,48,6)),balances],axis=-1)
  for lead in range(1,7):
   _,mo=add_month(year,month,lead);acc=np.sum(vals[...,6+lead-scale:6+lead],axis=-1)
   par=np.load(SPEI_BASE/f'SPEI{scale}_distribution_parameters_LT{lead}.npz')['member']
   all_spei[:,:,:48,si,lead-1]=loglog_spei(acc,par[:,:,:48,mo-1,:])
  # ANN2 stack transforms the 48 member SPEIs into member 49.
  for lead in range(1,7):
   arts=joblib.load(SPEI_BASE/'stack_model_artifacts'/f'SPEI{scale}_AI_stack_candidates_LT{lead}.joblib')
   for key,item in arts.items():
    if not key.endswith('_ANN2'):continue
    i,j=map(int,key.split('_')[:2]);X=all_spei[i,j,:48,si,lead-1];keep=item['member_indices_zero_based'];means=item['imputation_full_period_means'];xx=np.where(np.isfinite(X[keep]),X[keep],means)
    if np.all(np.isfinite(xx)):all_spei[i,j,48,si,lead-1]=item['model'].predict(xx[None])[0]
  log.info('stage=SPEI scale=%d status=SUCCESS stack=ANN2',scale)
 dr=all_spei[:,:,48];prob=[]
 for lo,hi in [(-1.,0.),(-1.5,-1.),(-2.,-1.5),(-np.inf,-2.)]:
  event=np.where(np.isfinite(all_spei),(all_spei>=lo)&(all_spei<hi),np.nan);prob.append(np.nanmean(event,axis=2)*100)
 write_drought(out/'Dr_Prob.nc',dr,prob,year,month);return dr,prob

def write_drought(path,dr,prob,year,month):
 with netcdf_file(path,'w') as d:
  for n,k in [('lon',40),('lat',68),('timescale',3),('lead',6)]:d.createDimension(n,k)
  for n,x in [('lon',G_LON),('lat',G_LAT),('timescale',[1,3,6]),('lead',range(1,7))]:d.createVariable(n,'f',(n,))[:]=x
  for n,a in [('dr_ens',dr),('mild',prob[0]),('mord',prob[1]),('seve',prob[2]),('extr',prob[3])]:
   q=d.createVariable(n,'f',('lon','lat','timescale','lead'));q._FillValue=FILL;q.missing_value=FILL;q.units='1' if n=='dr_ens' else 'percent';q[:]=np.where(np.isfinite(a),a,FILL)
  d.stack_method='ANN2';d.probability_denominator='49 members';d.accumulation_method='PA2: same-lead operational forecasts from preceding initialization months; pre-2021 antecedents from full-period corrected members; no climatological substitution';d.drought_probability_thresholds='mild<=-0.5|mord<=-1.0|seve<=-1.5|extr<=-2.0';d.init_month=f'{year:04d}-{month:02d}'

def combine_climatology():
 path=ROOT/'outputs/Clim/Monthly_Clim.nc';path.parent.mkdir(parents=True,exist_ok=True)
 lon,lat,_=spei.coordinates_and_mask();po,to=bc.climatology('RegCM','pr')[1],bc.climatology('RegCM','t2')[1]
 pm=[];tm=[]
 for r in bc.RCMS:pm.append(bc.climatology(r,'pr')[0]);tm.append(bc.climatology(r,'t2')[0])
 with netcdf_file(path,'w') as d:
  for n,k in [('lon',40),('lat',68),('lead',6),('month',12)]:d.createDimension(n,k)
  for n,x in [('lon',lon),('lat',lat),('lead',range(1,7)),('month',range(1,13))]:d.createVariable(n,'f',(n,))[:]=x
  for n,a,u in [('pr_o',po,'mm month-1'),('t2_o',to,'degree_C'),('pr_m',np.nanmean(pm,axis=0),'mm month-1'),('t2_m',np.nanmean(tm,axis=0),'degree_C')]:
   dims=('lon','lat','month') if a.ndim==3 else ('lon','lat','lead','month');q=d.createVariable(n,'f',dims);q._FillValue=FILL;q.missing_value=FILL;q.units=u;q[:]=np.where(np.isfinite(a),a,FILL)
  d.model_climatology='mean of RegCM and clWRF climatologies; operational Stack reference';d.period='1982-01/2021-12'

def maps_and_web(out,ens,ano,clim,dr,prob,year,month,day,log):
 figs=out/"figures";figs.mkdir(exist_ok=True)
 for pattern in ("event_S*_LT*.png","prob_D*_S*_LT*.png"):
  for obsolete in figs.glob(pattern):obsolete.unlink()
 mask=bc.read_mask();boundary=shape_segments()
 specs=[("pr_fcs",ens["pr"],"Total monthly rainfall","mm/month","YlGnBu",RAIN_LEVELS),("pr_ano",ano["pr"],"Rainfall anomaly","%",RAIN_ANOM_COLORS,RAIN_ANOM_LEVELS),("t2_fcs",ens["t2"],"Mean monthly temperature","°C","turbo",TEMP_LEVELS),("t2_ano",ano["t2"],"Temperature anomaly","°C",TEMP_ANOM_COLORS,TEMP_ANOM_LEVELS)]
 links=[]
 for key,a,title,unit,cmap,levels in specs:
  cm,norm_=discrete_cmap(cmap,levels);fig,axs=plt.subplots(1,6,figsize=(21,5.4),constrained_layout=True)
  for l,ax in enumerate(axs):
   ty,tm=add_month(year,month,l+1);z=np.where((mask>=1)&(mask<=7),a[...,l],np.nan);m=ax.pcolormesh(G_LON,G_LAT,z.T,cmap=cm,norm=norm_,shading="nearest");ax.set_title(f"LT{l+1}: {tm:02d}/{ty}");decorate_map(ax,boundary)
  fig.colorbar(m,ax=axs,orientation="horizontal",fraction=.08,label=unit,ticks=levels,extend="both");fig.suptitle(f"{title} — Initial: {day:02d}/{month:02d}/{year}");fig.savefig(figs/f"{key}.png",dpi=180);plt.close(fig);links.append((title,f"{key}.png"))
 spei_cmap=ListedColormap(SPEI_COLORS);spei_norm=BoundaryNorm(SPEI_BOUNDARIES,len(SPEI_COLORS))
 for si,scale in enumerate((1,3,6)):
  fig,axs=plt.subplots(1,6,figsize=(21,5.4),constrained_layout=True)
  for l,ax in enumerate(axs):
   ty,tm=add_month(year,month,l+1);z=np.where((mask>=1)&(mask<=7),dr[...,si,l],np.nan);m=ax.pcolormesh(G_LON,G_LAT,z.T,cmap=spei_cmap,norm=spei_norm,shading="nearest");ax.set_title(f"LT{l+1}: {tm:02d}/{ty}");decorate_map(ax,boundary)
  fig.colorbar(m,ax=axs,orientation="horizontal",fraction=.08,label=f"SPEI-{scale}",ticks=SPEI_TICKS,extend="both");title=f"Drought forecast SPEI-{scale}";fig.suptitle(f"{title} — Initial: {day:02d}/{month:02d}/{year}");name=f"SPEI{scale}.png";fig.savefig(figs/name,dpi=180);plt.close(fig);links.append((title,name))
  for di,label in enumerate(("Mild drought (SPEI ≤ -0.5)","Moderate drought (SPEI ≤ -1.0)","Severe drought (SPEI ≤ -1.5)","Extreme drought (SPEI ≤ -2.0)")):
   pcmap,pnorm=discrete_cmap("YlOrRd",PROB_LEVELS);fig,axs=plt.subplots(1,6,figsize=(21,5.4),constrained_layout=True)
   for l,ax in enumerate(axs):
    ty,tm=add_month(year,month,l+1);z=np.where((mask>=1)&(mask<=7),prob[di][...,si,l],np.nan);m=ax.pcolormesh(G_LON,G_LAT,z.T,cmap=pcmap,norm=pnorm,shading="nearest");ax.set_title(f"LT{l+1}: {tm:02d}/{ty}");decorate_map(ax,boundary)
   fig.colorbar(m,ax=axs,orientation="horizontal",fraction=.08,label="Probability (%)",ticks=PROB_LEVELS);title=f"{label}, SPEI-{scale}";fig.suptitle(f"{title} — Initial: {day:02d}/{month:02d}/{year}");name=f"prob_D{di+1}_S{scale}.png";fig.savefig(figs/name,dpi=180);plt.close(fig);links.append((title,name))
 sections="".join(f"<section><h2>{html.escape(title)}</h2><img src=\"figures/{name}\"></section>" for title,name in links)
 doc=f"""<!doctype html><html><head><meta charset="utf-8"><title>Dr_SFC operational forecast</title><style>body{{font:15px system-ui;margin:20px;max-width:2200px}}img{{width:100%}}section{{border:1px solid #aaa;padding:12px;margin:16px 0}}a{{margin-right:18px}}</style></head><body><h1>Dr_SFC operational forecast</h1><p><b>Initial:</b> {day:02d}/{month:02d}/{year}. Target months are printed on every panel. Bias correction uses 48 full-period members; SPEI Stack uses ANN2.</p><p><a href="Ope_Pr_T2_and_Ano.nc">Download pr/t2/anomalies</a><a href="Dr_Prob.nc">Download SPEI/probabilities</a><a href="preflight.json">Preflight report</a></p>{sections}</body></html>"""
 (out/"index.html").write_text(doc);log.info("stage=website status=SUCCESS path=%s",out/"index.html")

def d4_maps_and_web(out,prob,year,month,day,log):
 """Add the three missing extreme-drought probability figures to an existing page."""
 page=out/"index.html";doc=page.read_text()
 if "prob_D4_S1.png" in doc:return
 figs=out/"figures";figs.mkdir(exist_ok=True);mask=bc.read_mask();boundary=shape_segments();links=[]
 pcmap,pnorm=discrete_cmap("YlOrRd",PROB_LEVELS)
 for si,scale in enumerate((1,3,6)):
  fig,axs=plt.subplots(1,6,figsize=(21,5.4),constrained_layout=True)
  for lead,ax in enumerate(axs):
   ty,tm=add_month(year,month,lead+1);z=np.where((mask>=1)&(mask<=7),prob[3][...,si,lead],np.nan)
   image_=ax.pcolormesh(G_LON,G_LAT,z.T,cmap=pcmap,norm=pnorm,shading="nearest");ax.set_title(f"LT{lead+1}: {tm:02d}/{ty}");decorate_map(ax,boundary)
  fig.colorbar(image_,ax=axs,orientation="horizontal",fraction=.08,label="Probability (%)",ticks=PROB_LEVELS)
  title=f"Extreme drought (SPEI ≤ -2.0), SPEI-{scale}";fig.suptitle(f"{title} — Initial: {day:02d}/{month:02d}/{year}")
  name=f"prob_D4_S{scale}.png";fig.savefig(figs/name,dpi=180);plt.close(fig);links.append((title,name))
 sections="".join(f'<section><h2>{html.escape(title)}</h2><img src="figures/{name}"></section>' for title,name in links)
 page.write_text(doc.replace("</body>",sections+"</body>"));log.info("stage=website_D4 status=SUCCESS path=%s",page)
def detailed_drought_maps(out,dr,prob,year,month,day):
 figs=out/"figures";mask=bc.read_mask();links=[];boundary=shape_segments()
 cmap=ListedColormap(SPEI_COLORS);norm=BoundaryNorm(np.arange(-.5,9.5),9)
 levels=[(-np.inf,-2.),(-2.,-1.5),(-1.5,-1.),(-1.,-.5),(-.5,.5),(.5,1.),(1.,1.5),(1.5,2.),(2.,np.inf)]
 for si,scale in enumerate((1,3,6)):
  for lead in range(6):
   ty,tm=add_month(year,month,lead+1);z=dr[...,si,lead];cat=np.full(z.shape,np.nan)
   for ci,(lo,hi) in enumerate(levels):cat[np.isfinite(z)&(z>=lo)&(z<hi)]=ci
   fig,ax=plt.subplots(figsize=(5,8),constrained_layout=True);mm=ax.pcolormesh(G_LON,G_LAT,np.where((mask>=1)&(mask<=7),cat,np.nan).T,cmap=cmap,norm=norm,shading="nearest");cb=fig.colorbar(mm,ax=ax,ticks=range(9));cb.ax.set_yticklabels(["≤-2.0 D4","-1.5 D3","-1.0 D2","-0.5 D1","0 D0","+0.5 W1","+1.0 W2","+1.5 W3","≥+2.0 W4"]);ax.set_title(f"SPEI-{scale} event | Initial {day:02d}/{month:02d}/{year} | LT{lead+1} | {tm:02d}/{ty}");decorate_map(ax,boundary);name=f"event_S{scale}_LT{lead+1}.png";fig.savefig(figs/name,dpi=170);plt.close(fig);links.append((f"SPEI-{scale} event LT{lead+1}",name))
   for di,label in enumerate(("D1 mild","D2 moderate","D3 severe","D4 extreme")):
    pcmap,pnorm=discrete_cmap("YlOrRd",PROB_LEVELS);fig,ax=plt.subplots(figsize=(5,8),constrained_layout=True);zz=np.where((mask>=1)&(mask<=7),prob[di][...,si,lead],np.nan);mm=ax.pcolormesh(G_LON,G_LAT,zz.T,cmap=pcmap,norm=pnorm,shading="nearest");fig.colorbar(mm,ax=ax,ticks=PROB_LEVELS,label="Probability (%)");ax.set_title(f"{label} | SPEI-{scale} | Initial {day:02d}/{month:02d}/{year} | LT{lead+1} | {tm:02d}/{ty}");decorate_map(ax,boundary);name=f"prob_D{di+1}_S{scale}_LT{lead+1}.png";fig.savefig(figs/name,dpi=170);plt.close(fig);links.append((f"{label}, SPEI-{scale}, LT{lead+1}",name))
 p=out/"index.html";doc=p.read_text();gallery="<h1>Detailed drought products: 18 event maps and 72 probability maps</h1>"+"".join(f"<section><h2>{html.escape(t)}</h2><img src=\"figures/{n}\"></section>" for t,n in links);p.write_text(doc.replace("</body>",gallery+"</body>"))

def load_ope_products(path):
 with netcdf_file(path,"r",mmap=False) as d:
  return {"pr":np.array(d.variables["pr_members"][:],copy=True),"t2":np.array(d.variables["t2_members"][:],copy=True)}

def progressive_cycle(products,year,month,cycle,log):
 prior=[c for c in range(1,cycle) if (ROOT/"outputs/Ope"/f"{year:04d}"/f"{month:02d}"/f"{DAY_MAP[c]:02d}"/"Ope_Pr_T2_and_Ano.nc").exists()]
 if not prior:return products,1
 pc=max(prior);pp=ROOT/"outputs/Ope"/f"{year:04d}"/f"{month:02d}"/f"{DAY_MAP[pc]:02d}"/"Ope_Pr_T2_and_Ano.nc"
 old=load_ope_products(pp);n=len(prior)
 # n is the number of successfully processed prior cycles; missing or invalid cycles are excluded.
 ans={v:(old[v]*n+products[v])/(n+1) for v in bc.VARS};log.info("stage=cycle_ensemble cycles=%d",n+1);return ans,n+1

FULL_BC=ROOT/"outputs/production_output/bias_corrected_full_period/1399282"

def load_full_period_initialization(year,month):
 """Load 48 corrected members for one historical initialization month."""
 products={v:np.full((40,68,48,6),np.nan,np.float32) for v in bc.VARS}
 for var in bc.VARS:
  for lead in range(1,7):
   target_year,target_month=add_month(year,month,lead)
   time_index=(target_year-1982)*12+target_month-1
   if not 0<=time_index<480:raise IndexError(f"historical target outside 1982-2021: {target_year:04d}-{target_month:02d}")
   for ri,rcm in enumerate(bc.RCMS):
    path=FULL_BC/f"{rcm}_{var}_bc_full_LT{lead}.nc"
    with netcdf_file(path,"r",mmap=False) as d:
     field=np.array(d.variables[f"{var}_bc"][:,:,:,time_index],copy=True).astype(np.float32)
    products[var][:,:,ri*24:(ri+1)*24,lead-1]=np.where(field<=-90,np.nan,field)
 return products
def latest_month_product(year,month):
 p=ROOT/"outputs/Ope"/f"{year:04d}"/f"{month:02d}"
 found=[p/f"{DAY_MAP[c]:02d}"/"Ope_Pr_T2_and_Ano.nc" for c in range(1,7)]
 found=[x for x in found if x.exists()];return found[-1] if found else None

def drought_pa2(products,year,month,out,log):
 all_spei=np.full((40,68,49,3,6),np.nan,np.float32);obs_t=bc.climatology("RegCM","t2")[1]
 scales=(1,3,6);history={(year,month):products};balances={}
 for back in range(1,max(scales)):
  iy,im=add_month(year,month,-back);path=latest_month_product(iy,im)
  history[(iy,im)]=load_ope_products(path) if path else (load_full_period_initialization(iy,im) if iy<=2020 else None)
 for (iy,im),q in history.items():
  if q is None:continue
  for lead in range(1,7):
   _,target=add_month(iy,im,lead);pet=np.empty((40,68,48),np.float32)
   for member in range(48):pet[...,member]=operational_pet(q["t2"][...,member,lead-1],target,obs_t)
   balances[(iy,im,lead)]=q["pr"][...,lead-1]-pet
 for si,scale in enumerate(scales):
  for lead in range(1,7):
   keys=[(*add_month(year,month,-back),lead) for back in range(scale-1,-1,-1)]
   if any(key not in balances for key in keys):
    raise RuntimeError(f'PA2 requires {scale-1} preceding initialization months for SPEI-{scale} LT{lead}; history is incomplete before {year:04d}-{month:02d}')
   acc=np.sum([balances[key] for key in keys],axis=0);_,mo=add_month(year,month,lead)
   par=np.load(SPEI_BASE/f"SPEI{scale}_distribution_parameters_LT{lead}.npz")["member"]
   all_spei[:,:,:48,si,lead-1]=loglog_spei(acc,par[:,:,:48,mo-1,:])
   arts=joblib.load(SPEI_BASE/"stack_model_artifacts"/f"SPEI{scale}_AI_stack_candidates_LT{lead}.joblib")
   for key,item in arts.items():
    if not key.endswith("_ANN2"):continue
    i,j=map(int,key.split("_")[:2]);X=all_spei[i,j,:48,si,lead-1];keep=item["member_indices_zero_based"];means=item["imputation_full_period_means"];xx=np.where(np.isfinite(X[keep]),X[keep],means)
    if np.all(np.isfinite(xx)):all_spei[i,j,48,si,lead-1]=item["model"].predict(xx[None])[0]
  log.info("stage=SPEI_PA2 scale=%d status=SUCCESS",scale)
 dr=all_spei[:,:,48];prob=[]
 for threshold in (-.5,-1.,-1.5,-2.):
  event=np.where(np.isfinite(all_spei),all_spei<=threshold,np.nan);prob.append(np.nanmean(event,axis=2)*100)
 write_drought(out/"Dr_Prob.nc",dr,prob,year,month);return dr,prob

def load_existing_drought(path):
 with netcdf_file(path,"r",mmap=False) as d:
  names=("dr_ens","mild","mord","seve","extr")
  arrays=[np.array(d.variables[name][:],copy=True).astype(np.float32) for name in names]
 arrays=[np.where(a<=-90,np.nan,a) for a in arrays]
 return arrays[0],arrays[1:]

def main():
 global G_LON,G_LAT
 G_LON,G_LAT,_=spei.coordinates_and_mask()
 p=argparse.ArgumentParser();p.add_argument("--year",type=int,required=True);p.add_argument("--month",type=int,required=True);p.add_argument("--cycle",type=int,choices=range(1,7),required=True);p.add_argument("--bc-only",action="store_true");p.add_argument("--drought-only",action="store_true");p.add_argument("--plot-only",action="store_true");p.add_argument("--d4-only",action="store_true");a=p.parse_args();tag=f'{a.year:04d}{a.month:02d}{DAY_MAP[a.cycle]:02d}';log=getlog(tag)
 if a.d4_only:
  out=ROOT/"outputs/Ope"/f"{a.year:04d}"/f"{a.month:02d}"/f"{DAY_MAP[a.cycle]:02d}"
  _,prob=load_existing_drought(out/"Dr_Prob.nc");d4_maps_and_web(out,prob,a.year,a.month,DAY_MAP[a.cycle],log);return
 if a.plot_only:
  out=ROOT/"outputs/Ope"/f"{a.year:04d}"/f"{a.month:02d}"/f"{DAY_MAP[a.cycle]:02d}"
  products=load_ope_products(out/"Ope_Pr_T2_and_Ano.nc")
  ens={v:np.nanmean(products[v],axis=2) for v in bc.VARS}
  clim={v:bc.climatology("RegCM",v)[1] for v in bc.VARS}
  ano=forecast_anomalies(ens,clim,a.year,a.month)
  dr,prob=load_existing_drought(out/"Dr_Prob.nc")
  maps_and_web(out,ens,ano,clim,dr,prob,a.year,a.month,DAY_MAP[a.cycle],log)
  log.info("stage=finalize status=PLOT_ONLY_SUCCESS output=%s",out)
  return
 out=preflight(a.year,a.month,a.cycle,log);
 if not a.drought_only:combine_climatology()
 if a.drought_only:
  products=load_ope_products(out/"Ope_Pr_T2_and_Ano.nc");ens={v:np.nanmean(products[v],axis=2) for v in bc.VARS};clim={v:bc.climatology("RegCM",v)[1] for v in bc.VARS};ano=forecast_anomalies(ens,clim,a.year,a.month)
 else:
  products,ens,ano,clim=bias_correct(a.year,a.month,a.cycle,out,log);products,ncycle=progressive_cycle(products,a.year,a.month,a.cycle,log);ens={v:np.nanmean(products[v],axis=2) for v in bc.VARS};ano=forecast_anomalies(ens,clim,a.year,a.month);write_prt2(out/"Ope_Pr_T2_and_Ano.nc",products,ens,ano,a.year,a.month,DAY_MAP[a.cycle])
  if a.bc_only:(out/"SUCCESS_BC").write_text(datetime.now().isoformat());return
 dr,prob=drought_pa2(products,a.year,a.month,out,log);maps_and_web(out,ens,ano,clim,dr,prob,a.year,a.month,DAY_MAP[a.cycle],log);(out/"SUCCESS").write_text(datetime.now().isoformat());log.info("stage=finalize status=SUCCESS output=%s",out)
if __name__=='__main__':main()