#!/usr/bin/env python3
"""Deep PySR search for Karman cross-Re: individual + joint, niterations=120.

Usage:
    conda run -n sr_env python src/SR_analysis/sindy/run_pysr_deep.py --individual
    conda run -n sr_env python src/SR_analysis/sindy/run_pysr_deep.py --joint
    conda run -n sr_env python src/SR_analysis/sindy/run_pysr_deep.py --both
"""
from __future__ import annotations
import json, os, sys
import numpy as np
from pysr import PySRRegressor

_REPO = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..", ".."))
sys.path.insert(0, _REPO); sys.path.insert(0, os.path.join(_REPO, "src"))
from SR_analysis.utils.feature_builder import (
    compute_dimensionless, compute_features, build_feature_matrix,
    CORE_FEAT_KEYS_V2, PHASE_STATE_KEYS, MU_FEAT_KEYS,
)
from SR_analysis.configs import get_scene

SINDY_DIR = os.path.join(os.path.dirname(__file__))
ALL_MU = list(MU_FEAT_KEYS) + ["mu_Cl_tot"]
RESULTS_DIR = os.path.join(SINDY_DIR, "..", "validate", "results")
os.makedirs(RESULTS_DIR, exist_ok=True)
RE_SCENES = ["karman_re50", "karman_re100", "karman_re200", "karman_re400"]

# Feature sets
PHYS = [k for k in CORE_FEAT_KEYS_V2 if not k.startswith(("aF_","aB_","aT_","daF","daB","daT"))]
PHYS_DADT = PHYS + ["daF_dt","daB_dt","daT_dt"]

DESC = {
    "phase": ("PHASE_STATE_KEYS (6)", PHASE_STATE_KEYS),
    "phys": ("physics only (8)", PHYS),
    "phys_dadt": ("phys + da/dt (11)", PHYS_DADT),
    "full": ("CORE_FEAT_V2 + mu (20)", CORE_FEAT_KEYS_V2 + ALL_MU),
}

def load_and_build(scene_name, front_keys, use_mu=True):
    cfg = get_scene(scene_name); u0=cfg["u0"]; mu=cfg["mu"]; dc=cfg.get("sample_interval",800)/2000.0
    npz=np.load(f"{SINDY_DIR}/../data/karman/{scene_name}/controlled.npz")
    s=npz["sensors"].astype(np.float64); f=npz["forces"].astype(np.float64)
    ap=(npz["actions"].astype(np.float64)*cfg["action_scale"]+np.array(cfg["action_bias"]))*u0
    T=s.shape[0]; a1=np.zeros((T,3)); a2=np.zeros((T,3))
    a1[1:]=ap[:-1]; a2[2:]=ap[:-2]
    dim=compute_dimensionless(s,f,u0=u0,d=20.0)
    sym=compute_features(dim,a1,a2,mu,alpha_mode=False,include_mu=use_mu,include_cos_sin=False,u0=u0,dt_c=dc,sensors_raw=s,forces_raw=f)
    Xf=build_feature_matrix(sym,front_keys,add_bias=False)[2:]; Ya=ap[2:]/u0
    return Xf, Ya

def run_pysr(X, y, feat_names, label, seed=42, niter=120):
    print(f"\n--- {label} ({X.shape[0]} samples, {X.shape[1]} feats) ---")
    for i,k in enumerate(feat_names):
        print(f"    {k}: std={np.std(X[:,i]):.6f}")
    m = PySRRegressor(
        binary_operators=["+","-","*","/"], unary_operators=["square"],
        niterations=niter, populations=40, maxsize=15,
        complexity_of_constants=2, parsimony=0.01, batching=False,
        random_state=seed,
    )
    m.fit(X, y, variable_names=feat_names)
    return {
        "best_sympy": str(m.sympy()),
        "best_score": float(m.score(X, y)),
        "equations": m.equations_.to_dict(orient="records"),
    }

def run_individual(run_id=""):
    """Run deep PySR on each Re independently with phys_dadt features."""
    print("="*60)
    print("INDIVIDUAL Re (each with phys_dadt + mu)")
    print("="*60)
    front_keys = PHYS_DADT + ALL_MU
    all_res = {}
    for sn in RE_SCENES:
        Xf, Ya = load_and_build(sn, front_keys, use_mu=True)
        res = run_pysr(Xf, Ya[:, 0], front_keys, f"{sn} front", niter=120)
        all_res[sn] = {"front": res}
    out = os.path.join(RESULTS_DIR, f"pysr_deep_individual{run_id}.json")
    json.dump(all_res, open(out, "w"), indent=2, default=str)
    for sn in RE_SCENES:
        r = all_res[sn]["front"]
        print(f"  {sn}: {r['best_sympy']}  (R2={r['best_score']:.4f})")

def run_joint(run_id=""):
    """Run deep joint PySR across all Re."""
    print("="*60)
    print("JOINT (phys_dadt + mu, across all Re)")
    print("="*60)
    front_keys = PHYS_DADT + ALL_MU
    X_list, Y_list = [], []
    for sn in RE_SCENES:
        Xf, Ya = load_and_build(sn, front_keys, use_mu=True)
        X_list.append(Xf); Y_list.append(Ya[:, 0])
    X_all = np.vstack(X_list); Y_all = np.concatenate(Y_list)
    res = run_pysr(X_all, Y_all, front_keys, "joint front", niter=120)
    out = os.path.join(RESULTS_DIR, f"pysr_deep_joint{run_id}.json")
    json.dump(res, open(out, "w"), indent=2, default=str)
    print(f"\nJoint: {res['best_sympy']}  (R2={res['best_score']:.4f})")

def show_pareto(path, label):
    d = json.load(open(path))
    if "front" in d:
        eqs = d["front"]["equations"]
    elif "equations" in d:
        eqs = d["equations"]
    else:
        print(f"  No equations in {path}")
        return
    print(f"\n=== Pareto {label} ===")
    for eq in eqs[:15]:
        print(f"  c={eq['complexity']:2d}  loss={eq['loss']:.6f}  {eq['equation']}")

if __name__ == "__main__":
    import argparse
    ap = argparse.ArgumentParser()
    ap.add_argument("--individual", action="store_true")
    ap.add_argument("--joint", action="store_true")
    ap.add_argument("--both", action="store_true")
    ap.add_argument("--show", type=str, default=None, help="Show pareto for saved result")
    args = ap.parse_args()
    if args.show:
        show_pareto(args.show, args.show)
    elif args.both:
        run_individual("_deep"); run_joint("_deep")
    elif args.individual:
        run_individual("_deep")
    elif args.joint:
        run_joint("_deep")
    else:
        run_individual("_deep"); run_joint("_deep")
