#!/usr/bin/env python3
"""Calculate Tx90p and EHD thresholds directly from EC_Combined and OBS_4EC."""
import argparse, glob, os, warnings
from datetime import date
from pathlib import Path
import numpy as np
from netCDF4 import Dataset
BASE = Path(__file__).resolve().parent
DATA = "/home/tanpv/Data/QG.25.196"
warnings.filterwarnings("ignore", message="All-NaN slice encountered")

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 window(days,target):
    delta=np.abs(days-target); return np.minimum(delta,365-delta)<=15

def grid(root):
    path=f"{root}/OBS/Mask_VN_S2S_EC_0.25.nc"
    with Dataset(path) as d:
        lon,lat=np.asarray(d['lon'][:]),np.asarray(d['lat'][:]); raw=np.ma.asarray(d['vn'][:])
        mask=~np.ma.getmaskarray(raw)&np.isfinite(np.ma.filled(raw,np.nan))
    return lon,lat,mask,path

def load_model(files,mask):
    val=np.empty((len(files),46)+mask.shape,'f4'); days=np.empty((len(files),46),'i2')
    for i,path in enumerate(files):
        init=date.fromisoformat(os.path.basename(path)[7:17])
        with Dataset(path) as d: val[i]=np.ma.filled(d['mx2t6'][:46],np.nan)
        val[i,:,~mask]=np.nan
        days[i]=[doy365(date.fromordinal(init.toordinal()+lead)) for lead in range(46)]
        if (i+1)%200==0: print(f"Read EC {i+1}/{len(files)}",flush=True)
    return val,days

def load_obs(files,mask):
    val=np.empty((len(files),)+mask.shape,'f4'); days=np.empty(len(files),'i2')
    for i,path in enumerate(files):
        days[i]=doy365(date.fromisoformat(os.path.basename(path)[4:14]))
        with Dataset(path) as d: val[i]=np.ma.filled(d['tasmax'][0],np.nan)
        val[i,~mask]=np.nan
        if (i+1)%1000==0: print(f"Read OBS {i+1}/{len(files)}",flush=True)
    return val,days

def q90_model(val,days):
    qp=np.full((365,46)+val.shape[2:],np.nan,'f4'); count=np.empty((365,46),'i2')
    for lead in range(46):
        for doy in range(365):
            use=window(days[:,lead],doy); count[doy,lead]=use.sum()
            qp[doy,lead]=np.nanpercentile(val[use,lead],90,axis=0)
        print(f"Calculated EC Q90 lead {lead+1}/46",flush=True)
    return qp,count

def q90_obs(val,days):
    qp=np.full((365,)+val.shape[1:],np.nan,'f4'); count=np.empty(365,'i2')
    for doy in range(365):
        use=window(days,doy); count[doy]=use.sum(); qp[doy]=np.nanpercentile(val[use],90,axis=0)
        if (doy+1)%50==0: print(f"Calculated OBS Q90 day {doy+1}/365",flush=True)
    return qp,count

def write(path,qp,count,lon,lat,source,nfiles,mask_path):
    os.makedirs(os.path.dirname(path),exist_ok=True); model=qp.ndim==4
    with Dataset(path,'w') as d:
        d.createDimension('day',365)
        if model: d.createDimension('lead',46)
        d.createDimension('lat',len(lat)); d.createDimension('lon',len(lon))
        d.createVariable('day','i2',('day',))[:]=np.arange(1,366)
        if model: d.createVariable('lead','i1',('lead',))[:]=np.arange(46)
        d.createVariable('lat','f4',('lat',))[:]=lat; d.createVariable('lon','f4',('lon',))[:]=lon
        dims=('day','lead','lat','lon') if model else ('day','lat','lon')
        q=d.createVariable('qp','f4',dims,zlib=True,fill_value=-9999.)
        t=d.createVariable('tx_thres','f4',dims,zlib=True,fill_value=-9999.)
        q[:]=np.ma.masked_invalid(qp); t[:]=np.ma.masked_invalid(np.maximum(qp,35.)); q.units=t.units='degree_Celsius'
        d.createVariable('sample_count','i2',('day','lead') if model else ('day',))[:]=count
        d.source=source; d.source_file_count=nfiles; d.period='2003-2022'; d.percentile=90.
        d.calendar_window='circular day-of-year +/-15 days'; d.method='empirical percentile; no bootstrap'; d.mask=mask_path

def main():
    p=argparse.ArgumentParser(); p.add_argument('--data',default=DATA); p.add_argument('--output-dir', default=str(BASE / 'outputs' / 'Cli')); a=p.parse_args()
    lon,lat,mask,mask_path=grid(a.data)
    mf=sorted(glob.glob(f"{a.data}/EC_Combined/*/EC_all_*.nc")); of=sorted(glob.glob(f"{a.data}/OBS_4EC/*/Obs_*.nc"))
    if len(mf) != 2080 or len(of) != 7305:
        raise RuntimeError(f"Expected 2080 EC and 7305 OBS files; found {len(mf)} and {len(of)}")
    mv,md=load_model(mf,mask); mq,mc=q90_model(mv,md)
    write(f"{a.output_dir}/Qp_tx_mod.nc",mq,mc,lon,lat,'EC_Combined/mx2t6',len(mf),mask_path); del mv,mq
    ov,od=load_obs(of,mask); oq,oc=q90_obs(ov,od)
    write(f"{a.output_dir}/Qp_tx_obs.nc",oq,oc,lon,lat,'OBS_4EC/tasmax[0]',len(of),mask_path)
    print(f"Wrote threshold files to {a.output_dir}",flush=True)
if __name__=='__main__': main()