#!/usr/bin/env python3
"""Section 1.5: refit evaluated Dr_SFC configurations on all 1982-2021 data."""
import argparse,csv,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.exceptions import ConvergenceWarning
from sklearn.model_selection import GridSearchCV
import ai_bias_correction as base
ROOT=base.ROOT; RCMS=base.RCMS; VARS=base.VARS; LEADS=base.LEADS; MONTHS=base.MONTHS; METHODS=base.METHODS; SPATIAL=base.SPATIAL; SETS=base.SETS; MEMBERS=base.MEMBERS; FILL=base.FILL

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('full_period')

def selected_hyperparameters(model,method):
 if method.startswith('Ridge'):
  ridge=model.named_steps['ridgecv']
  return {'alpha':float(ridge.alpha_)}
 return dict(model.best_params_)

def full_period_estimator(method,seed,n_features):
 if method.startswith('Ridge'):
  model,_=base.estimator(method,seed,n_features)
  return model,{'strategy':'5-fold CV' if method=='Ridge1' else 'GCV','search_space':{'alpha':base.ALPHAS.tolist()}}
 hidden=base.ann_hidden_layers(method,n_features)
 pipeline,_=base.estimator(method,seed,n_features)
 # Keep the evaluated ANN topology fixed, but redetermine optimization and
 # regularization hyperparameters from the complete-period samples. GridSearchCV
 # refits the selected configuration on every supplied sample automatically.
 search={'mlpregressor__alpha':[1e-5,1e-4,1e-3],
         'mlpregressor__learning_rate_init':[5e-4,1e-3]}
 model=GridSearchCV(pipeline,search,cv=3,scoring='neg_mean_squared_error',refit=True,n_jobs=1,error_score='raise')
 return model,{'strategy':'3-fold CV with full-sample refit','search_space':search,'fixed_hidden_layers':list(hidden)}

def fit_predict_full(X,y,method,regions,regional,seed,name):
 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);ok=gm&np.isfinite(y)&np.all(np.isfinite(X),axis=1)
  if ok.sum()<50:
   diagnostics.append({'group':int(g),'status':'failed','reason':'insufficient_full_period_samples','n_fit':int(ok.sum())});continue
  model,tuning=full_period_estimator(method,seed+g,X.shape[1])
  hidden=list(base.ann_hidden_layers(method,X.shape[1])) if method.startswith('ANN') else []
  try:
   with warnings.catch_warnings(record=True) as caught:
    warnings.simplefilter('always');model.fit(X[ok],y[ok])
   conv=[str(w.message) for w in caught if issubclass(w.category,ConvergenceWarning)]
   selected=selected_hyperparameters(model,method);pred[ok]=model.predict(X[ok]).astype(np.float32);models[str(g)]=model
   diagnostics.append({'group':int(g),'status':'completed_with_warning' if conv else 'passed','reason':'; '.join(conv),'n_fit':int(ok.sum()),'n_features':int(X.shape[1]),'hidden_layers':hidden,'selected_alpha':selected.get('alpha'),'selected_hyperparameters':selected,'tuning':tuning})
  except Exception as exc:diagnostics.append({'group':int(g),'status':'failed','reason':f'{type(exc).__name__}: {exc}','n_fit':int(ok.sum()),'n_features':int(X.shape[1]),'hidden_layers':hidden,'selected_alpha':None,'selected_hyperparameters':{},'tuning':tuning})
 return pred,models,diagnostics

def metric(obs,pred,raw,var,member):
 ok=np.isfinite(obs)&np.isfinite(pred)&np.isfinite(raw)
 if not ok.any():return {'period':'full_1982_2021','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)));metric_name='MAE'
 if var=='pr':
  den=float(np.mean(o));err=err/den*100 if den>1e-6 else np.nan;ref=ref/den*100 if den>1e-6 else np.nan;metric_name='RMAE_percent'
 corr=float(np.corrcoef(o,p)[0,1]) if len(o)>=3 and np.std(o)>0 and np.std(p)>0 else np.nan
 av=(ref-err)/max(ref,err,1e-12) if np.isfinite(err) and np.isfinite(ref) else np.nan
 return {'period':'full_1982_2021','member':member,'n':int(ok.sum()),'metric_name':metric_name,'error':err,'corr':corr,'added_value':av}

def work_unit(rcm,var,lead,month,run,source_run,log):
 path=base.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:base.read_var(d,v) for v in needed};obs=base.read_var(d,f'{var}_o')
 reg=base.read_mask();mcli,ocli=base.climatology(rcm,var);ocl=ocli[...,month-1];tids=np.arange(month-1,480,12);raw=data[var][...,tids]
 target=obs[...,tids]-ocl[...,None] if var=='t2' else 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);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((base.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);name=f'{method}_{setname}_{spatial}'
    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_full(X,y,method,regs,spatial=='regional',lead*1000+month*10+mi,name)
    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;hidden=list(base.ann_hidden_layers(method,X.shape[1])) if method.startswith('ANN') else []
    artifacts[name]={'models':models,'diagnostics':diagnostics,'n_features':int(X.shape[1]),'hidden_layers':hidden,'configuration_source_evaluation_run':source_run,'hyperparameters_retuned_on_full_period':True}
    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']);alphas=';'.join(str(d['selected_alpha']) for d in diagnostics if d.get('selected_alpha') is not None)
    selected=';'.join(json.dumps(d.get('selected_hyperparameters',{}),sort_keys=True) for d in diagnostics)
    row=metric(obs[...,tids],bc,raw,var,name);row.update(status=status,reason=reason,n_features=int(X.shape[1]),hidden_layers='x'.join(map(str,hidden)),selected_alphas=alphas,selected_hyperparameters=selected,hyperparameters_retuned_on_full_period=True);rows.append(row)
    log.info('Full fit %s %s_%s_LT%d_M%02d status=%s n_features=%d hidden=%s alphas=%s',name,rcm,var,lead,month,status,X.shape[1],hidden,alphas or 'n/a');mi+=1
 shard=ROOT/'outputs/production_output/shards_full_period'/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)
 art=ROOT/'outputs/production_output/model_artifacts_full_period'/run;art.mkdir(parents=True,exist_ok=True);joblib.dump({'rcm':rcm,'var':var,'lead':lead,'month':month,'members':MEMBERS,'models':artifacts,'training_period':'1982-01/2021-12','configuration_source_evaluation_run':source_run,'hyperparameters_retuned_on_full_period':True},art/f'{stem}.joblib',compress=3)
 (shard/f'{stem}.json').write_text(json.dumps(rows));log.info('Completed full-period %s',stem)

def worker(run,source_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('Full-period worker units=%d world=%d source=%s',len(mine),size,source_run)
 for u in mine:
  r,v,l,m=u;stem=f'{r}_{v}_LT{l}_M{m:02d}'
  shard=ROOT/'outputs/production_output/shards_full_period'/run
  artifact=ROOT/'outputs/production_output/model_artifacts_full_period'/run
  if (shard/f'{stem}.npz').exists() and (shard/f'{stem}.json').exists() and (artifact/f'{stem}.joblib').exists():
   log.info('Resume skip completed full-period %s',stem);continue
  work_unit(*u,run,source_run,log)
 log.info('Full-period worker complete')

def finalize(run,source_run,log):
 shard=ROOT/'outputs/production_output/shards_full_period'/run;out=ROOT/'outputs/production_output/bias_corrected_full_period'/run
 with netcdf_file(base.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:
   member_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)
     for x in json.loads((shard/f'{stem}.json').read_text()):x.update(rcm=rcm,var=var,lead=lead,month=month);rows.append(x)
    all_data[(rcm,var,lead)]=data;base.write_nc(out/f'{rcm}_{var}_bc_full_LT{lead}.nc',f'{var}_bc',data,lon,lat,member_names,{'rcm':rcm,'lead':lead,'run_id':run,'period':'1982-01/2021-12','configuration_source_evaluation_run':source_run,'hyperparameters_retuned_on_full_period':'true'})
 for var in VARS:
  for lead in LEADS:
   stack=np.concatenate([all_data[(r,var,lead)] for r in RCMS],axis=2)
   with warnings.catch_warnings():warnings.simplefilter('ignore',RuntimeWarning);ens=np.nanmean(stack,axis=2,keepdims=True)
   base.write_nc(out/f'{var}_bc_ens_full_LT{lead}.nc',f'{var}_bc_ens',ens,lon,lat,['ensemble_mean'],{'lead':lead,'run_id':run,'period':'1982-01/2021-12','configuration_source_evaluation_run':source_run,'hyperparameters_retuned_on_full_period':'true'})
 ev=ROOT/'outputs/evaluation_output_full_period'/run;ev.mkdir(parents=True,exist_ok=True)
 with open(ev/'member_metrics_full_period.csv','w',newline='') as f:w=csv.DictWriter(f,fieldnames=sorted(rows[0]));w.writeheader();w.writerows(rows)
 out.mkdir(parents=True,exist_ok=True);(out/'SUCCESS').write_text(datetime.now().isoformat());log.info('Full-period finalization complete')

def main():
 p=argparse.ArgumentParser();p.add_argument('mode',choices=['worker','finalize']);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();log=logger(ROOT/'logs'/f'full_period_{a.run}.log')
 if a.mode=='worker':worker(a.run,a.source_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,a.source_run,log)
if __name__=='__main__':main()