"""Leakage-safe POD, PLS-type OID, and CCA sensitivity on reduced corrections."""
from __future__ import annotations

import argparse,json
from pathlib import Path
import numpy as np
from .analysis import fit_pod,fit_pls_type,fit_regularized_cca,physical_pls_modes,projector_similarity
from .artifacts import atomic_save_npz,atomic_write_json


def _r2(y,p):
 y=np.asarray(y,float); p=np.asarray(p,float); den=np.sum((y-y.mean(axis=0))**2); return float(1-np.sum((y-p)**2)/den) if den>0 else 0.0

def _linear(train_x,train_y,test_x):
 x=np.column_stack((np.ones(len(train_x)),train_x)); coef=np.linalg.lstsq(x,train_y,rcond=None)[0]
 return np.column_stack((np.ones(len(test_x)),test_x))@coef

def run(run_dir:Path,rank:int=20,train_fraction:float=.7):
 out=run_dir/"analysis"; out.mkdir(exist_ok=False)
 with np.load(run_dir/"aligned/correction_reduced.npz",allow_pickle=False) as d:
  q=d["correction"].astype(float); forces=d["forces"].astype(float); sensors=d["sensors"].astype(float)
 split=int(len(q)*train_fraction); train=np.arange(split); test=np.arange(split,len(q))
 pod=fit_pod(q[train],min(rank,len(train)-1)); a_train=pod.coefficients; a_test=pod.transform(q[test])
 recon_train=pod.mean+a_train@pod.modes.T; recon_test=pod.mean+a_test@pod.modes.T
 report={"samples":len(q),"train_samples":len(train),"test_samples":len(test),"split":"contiguous development/validation",
  "rank":pod.modes.shape[1],"pod_energy_retained":float(pod.cumulative_energy[-1]),
  "pod_reconstruction_r2_train":_r2(q[train],recon_train),"pod_reconstruction_r2_test":_r2(q[test],recon_test),"observables":{}}
 arrays={"pod_mean":pod.mean.astype(np.float32),"pod_modes":pod.modes.astype(np.float32),"pod_singular_values":pod.singular_values.astype(np.float32)}
 for name,y in (("forces",forces),("signature",sensors)):
  pls=fit_pls_type(a_train,y[train]); dims=min(3,pls["U"].shape[1]); ztr=(a_train-pls["a_mean"])/pls["a_scale"]@pls["U"][:,:dims]; zte=(a_test-pls["a_mean"])/pls["a_scale"]@pls["U"][:,:dims]
  pred=_linear(ztr,y[train],zte); pod_pred=_linear(a_train[:,:dims],y[train],a_test[:,:dims])
  cca={}
  for eps in (1e-6,1e-4,1e-2): cca[str(eps)]=fit_regularized_cca(a_train,y[train],eps)["S"][:3].tolist()
  report["observables"][name]={"pls_test_r2":_r2(y[test],pred),"pod_coordinate_test_r2":_r2(y[test],pod_pred),"cca_correlations":cca,"pls_singular_values":pls["S"].tolist()}
  arrays[f"{name}_pls_modes"]=physical_pls_modes(pod.modes,pls["a_scale"],pls["U"]).astype(np.float32)
 # Strict valid-pair delayed signature; no padded tail.
 delay=1; delayed_train=train[train+delay<split]; delayed_test=test[test+delay<len(q)]
 pls_delay=fit_pls_type(pod.transform(q[delayed_train]),sensors[delayed_train+delay]); dims=min(3,pls_delay["U"].shape[1])
 atr=pod.transform(q[delayed_train]); ate=pod.transform(q[delayed_test]); ztr=(atr-pls_delay["a_mean"])/pls_delay["a_scale"]@pls_delay["U"][:,:dims]; zte=(ate-pls_delay["a_mean"])/pls_delay["a_scale"]@pls_delay["U"][:,:dims]
 report["delayed_signature"]={"delay_fields":delay,"tail_policy":"discard","train_pairs":len(delayed_train),"test_pairs":len(delayed_test),"test_r2":_r2(sensors[delayed_test+delay],_linear(ztr,sensors[delayed_train+delay],zte))}
 half=len(train)//2; p1=fit_pod(q[train[:half]],min(10,half-1)); p2=fit_pod(q[train[half:]],min(10,len(train)-half-1)); report["development_subspace_similarity_rank10"]=projector_similarity(p1.modes,p2.modes)
 atomic_save_npz(out/"decomposition.npz",**arrays); atomic_write_json(out/"decomposition_report.json",report); return report

def main(argv=None):
 p=argparse.ArgumentParser();p.add_argument("run_dir",type=Path);p.add_argument("--rank",type=int,default=20);a=p.parse_args(argv);print(json.dumps(run(a.run_dir,a.rank),indent=2));return 0
if __name__=="__main__":raise SystemExit(main())
