#!/usr/bin/env python3
"""Section 1.5 full-period verification and navigable report."""
import argparse,csv,html,json,logging,re
from pathlib import Path
import numpy as np
import matplotlib;matplotlib.use('Agg')
import matplotlib.pyplot as plt
from scipy.io import netcdf_file
ROOT=Path('/home/tanpv/Progs/CLI/Dr');LEADS=range(1,7);MAP_LEADS=(1,3,5);MONTHS=range(1,13);RCMS=('RegCM','clWRF');METHODS=('Ridge1','Ridge2','ANN1','ANN2');SPATIAL=('national','regional');SETS={'pr':('SP1','SP2','SP3'),'t2':('ST1','ST2','ST3')}
def names(var):return [f'{r}_{m}_{s}_{p}' for r in RCMS for m in METHODS for s in SETS[var] for p in SPATIAL]+['Stack_Ensemble']
def mask_lon_lat():
 with netcdf_file(Path('/home/tanpv/Progs/Dr_VN_Ope/DATA/Mask_VN_0.25_Corrected.nc'),'r',mmap=False) as d:
  v=d.variables['reg'];reg=np.array(v[:],copy=True);dims=list(v.dimensions)
  if dims==['lat','lon']:reg=reg.T
 with netcdf_file(Path('/home/tanpv/Data/562_VNU/0.25_VN/For_BC/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)
 return reg,lon,lat
def metric_fields(obs,pred,var):
 ok=np.isfinite(obs)&np.isfinite(pred);n=ok.sum(2);o=np.where(ok,obs,np.nan);p=np.where(ok,pred,np.nan);ae=np.where(ok,np.abs(p-o),np.nan);mae=np.nanmean(ae,2);rmse=np.sqrt(np.nanmean(np.where(ok,(p-o)**2,np.nan),2));om=np.nanmean(o,2);pm=np.nanmean(p,2);cov=np.nansum((o-om[...,None])*(p-pm[...,None]),2);den=np.sqrt(np.nansum((o-om[...,None])**2,2)*np.nansum((p-pm[...,None])**2,2));corr=np.divide(cov,den,out=np.full_like(cov,np.nan),where=(den>0)&(n>=3))
 if var=='pr':mae=np.divide(mae,om,out=np.full_like(mae,np.nan),where=om>1e-6)*100
 return mae,rmse,corr,n
def scalar(field):
 x=field[np.isfinite(field)];return float(np.mean(x)) if x.size else np.nan
def load(run,var,lead,month):
 zs=[np.load(ROOT/'outputs/production_output/shards_full_period'/run/f'{r}_{var}_LT{lead}_M{month:02d}.npz') for r in RCMS];stack=np.concatenate([z['pred'] for z in zs],0)
 return zs[0]['obs'],stack,np.nanmean(stack,0)
def calculate(run,out,log):
 reg,lon,lat=mask_lon_lat();valid=(reg>=1)&(reg<=7);saved={};summary=[]
 for var in ('pr','t2'):
  nm=names(var);err=np.full((49,6,12,len(lon),len(lat)),np.nan,np.float32);corr=np.full_like(err,np.nan);rmse=np.full_like(err,np.nan)
  for li,lead in enumerate(LEADS):
   for mi,month in enumerate(MONTHS):
    obs,stack,ens=load(run,var,lead,month)
    for k,name in enumerate(nm):
     e,r,c,n=metric_fields(obs,stack[k] if k<48 else ens,var);err[k,li,mi]=np.where(valid,e,np.nan);corr[k,li,mi]=np.where(valid,c,np.nan);rmse[k,li,mi]=np.where(valid,r,np.nan);summary.append({'var':var,'lead':lead,'month':month,'member':name,'metric_name':'RMAE_percent' if var=='pr' else 'MAE','error':scalar(e[valid]),'rmse':scalar(r[valid]) if var=='t2' else '', 'corr':scalar(c[valid]),'valid_gridpoints':int(np.isfinite(e[valid]).sum())})
    log.info('full metrics var=%s LT=%s month=%02d',var,lead,month)
  np.savez_compressed(out/f'{var}_verification_fields_full_period.npz',error=err,rmse=rmse,corr=corr,names=np.array(nm),leads=np.array(list(LEADS)),lon=lon,lat=lat,period='1982-01/2021-12');saved[var]=(err,rmse,corr,nm,lon,lat)
 with open(out/'verification_summary_full_period.csv','w',newline='') as f:w=csv.DictWriter(f,fieldnames=list(summary[0]));w.writeheader();w.writerows(summary)
 return saved
def map_rows():
 rows=[]
 for var in ('pr','t2'):
  metrics=('RMAE','CORR') if var=='pr' else ('MAE','CORR','RMSE')
  for metric in metrics:
   for name in names(var):rows.append((var,metric,name,f'{var}_{metric}_{name}.png'))
 return rows
def load_saved(out):
 saved={}
 for var in ('pr','t2'):
  z=np.load(out/f'{var}_verification_fields_full_period.npz');saved[var]=(z['error'],z['rmse'],z['corr'],[str(x) for x in z['names']],z['lon'],z['lat'])
 return saved
def plot_maps(saved,out,log,rank=0,size=1):
 d=out/'figures/verification_maps';d.mkdir(parents=True,exist_ok=True);idx=[x-1 for x in MAP_LEADS];tasks=map_rows();done=[]
 for task_i,(var,metric,name,fn) in enumerate(tasks):
  if task_i%size!=rank:continue
  err,rmse,corr,nm,lon,lat=saved[var];data={'RMAE':err,'MAE':err,'CORR':corr,'RMSE':rmse}[metric];k=nm.index(name);vals=data[k,idx];finite=vals[np.isfinite(vals)];vmin,vmax=((-1,1) if metric=='CORR' else (0,float(np.nanpercentile(finite,98))))
  fig,axs=plt.subplots(3,12,figsize=(22,6),sharex=True,sharey=True,constrained_layout=True);mesh=None
  for li,lead in enumerate(MAP_LEADS):
   for mi,month in enumerate(MONTHS):
    ax=axs[li,mi];mesh=ax.pcolormesh(lon,lat,vals[li,mi].T,shading='nearest',cmap='RdBu_r' if metric=='CORR' else 'YlOrRd',vmin=vmin,vmax=vmax);ax.set_title(f'M{month:02d}',fontsize=6) if li==0 else None;ax.set_ylabel(f'LT{lead}',fontsize=7) if mi==0 else None;ax.tick_params(labelsize=4,length=2)
  unit='%' if metric=='RMAE' else ('°C' if metric in ('MAE','RMSE') else 'coefficient');fig.colorbar(mesh,ax=axs,orientation='horizontal',fraction=.035,pad=.03,label=f'{metric} ({unit})');fig.suptitle(f'{name} — full-period {var} {metric} (1982–2021)',fontsize=11);fig.savefig(d/fn,dpi=150);plt.close(fig);done.append((var,metric,name,fn));log.info('rank=%d/%d full map %s',rank,size,fn)
 return done
def load_status(run,out):
 rows=[]
 for p in sorted((ROOT/'outputs/production_output/shards_full_period'/run).glob('*.json')):
  for r in json.loads(p.read_text()):rows.append({'unit':p.stem,**r})
 keys=sorted(set().union(*(x.keys() for x in rows)))
 with open(out/'model_status_full_period.csv','w',newline='') as f:w=csv.DictWriter(f,fieldnames=keys);w.writeheader();w.writerows(rows)
 return rows
def make_html(run,out,maps,status,source_run):
 summary={}
 for r in status:key=(r['member'].split('_')[0],r.get('status','unknown'));summary[key]=summary.get(key,0)+1
 trs=''.join(f'<tr><td>{html.escape(m)}</td><td>{html.escape(s)}</td><td>{n}</td></tr>' for (m,s),n in sorted(summary.items()));indexed=[]
 for i,(var,metric,name,fn) in enumerate(maps,1):indexed.append((f'map-{i:03d}',f'{var} — {metric} — {name}',fn))
 toc=''.join(f'<li class="map-link" data-search="{html.escape(label.lower())}"><a href="#{anchor}">{i:03d}. {html.escape(label)}</a></li>' for i,(anchor,label,fn) in enumerate(indexed,1));cards=''.join(f'<section id="{anchor}"><h3>{html.escape(label)}</h3><a href="figures/verification_maps/{fn}"><img loading="lazy" src="figures/verification_maps/{fn}" alt="{html.escape(label)}"></a><p class="back"><a href="#map-list">↑ Back to map list</a></p></section>' for anchor,label,fn in indexed)
 style='body{font:14px system-ui;margin:20px;scroll-behavior:smooth}img{width:100%}section{border:1px solid #bbb;margin:15px 0;padding:10px;scroll-margin-top:12px}table{border-collapse:collapse}td,th{border:1px solid #aaa;padding:5px}.map-index{columns:3;column-gap:28px}.map-index li{break-inside:avoid;margin:3px 0}.back{text-align:right}#map-filter{width:min(720px,95%);padding:8px}@media(max-width:900px){.map-index{columns:1}}';script='function filterMaps(){const q=document.getElementById("map-filter").value.toLowerCase();document.querySelectorAll(".map-link").forEach(x=>x.hidden=!x.dataset.search.includes(q));}'
 doc=f'<!doctype html><html><head><meta charset="utf-8"><title>Dr_SFC full-period verification {run}</title><style>{style}</style></head><body><h1>Dr_SFC full-period verification — run {run}</h1><p>Models refitted on all valid 1982–2021 samples. No training/validation split. Evaluation configuration source: <code>{source_run}</code>.</p><p>Fields: <a href="pr_verification_fields_full_period.npz">pr</a>, <a href="t2_verification_fields_full_period.npz">t2</a>. Tables: <a href="verification_summary_full_period.csv">verification</a>, <a href="model_status_full_period.csv">model status/hyperparameters</a>.</p><h2>Fit status</h2><table><tr><th>Method</th><th>Status</th><th>fits</th></tr>{trs}</table><nav id="map-list"><h2>Map list</h2><label for="map-filter">Filter maps:</label><br><input id="map-filter" type="search" oninput="filterMaps()" placeholder="Example: t2 RMSE ANN2 ST3"><ol class="map-index">{toc}</ol></nav><h2>Full-period verification maps</h2>{cards}<script>{script}</script></body></html>'
 (out/'index.html').write_text(doc)
def main():
 p=argparse.ArgumentParser();p.add_argument('mode',choices=['all','calculate','render','html'],nargs='?',default='all');p.add_argument('--run',required=True);p.add_argument('--source-run',default='1399272');p.add_argument('--rank',type=int);p.add_argument('--size',type=int);a=p.parse_args();out=ROOT/'outputs/evaluation_output_full_period'/a.run;out.mkdir(parents=True,exist_ok=True);logging.basicConfig(level=logging.INFO,format='%(asctime)s %(message)s',handlers=[logging.FileHandler(ROOT/'logs'/f'full_period_verification_{a.run}.log'),logging.StreamHandler()]);log=logging.getLogger()
 if a.mode in ('all','calculate'):load_status(a.run,out);saved=calculate(a.run,out,log)
 if a.mode in ('all','render'):
  if a.mode=='render':saved=load_saved(out)
  import os
  rank=a.rank if a.rank is not None else int(os.getenv('SLURM_PROCID','0'));size=a.size if a.size is not None else int(os.getenv('SLURM_NTASKS','1'));plot_maps(saved,out,log,rank,size)
 if a.mode in ('all','html'):
  with open(out/'model_status_full_period.csv') as f:status=list(csv.DictReader(f))
  make_html(a.run,out,map_rows(),status,a.source_run);log.info('full-period verification complete %s',out/'index.html')
if __name__=='__main__':main()