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

Adapted from run_pysr_deep.py for illusion scenes:
- Uses ILLUSION_PHASE_KEYS (10-dim) + target_Cd/target_Cl (2-dim)
- Data loaded from illusion directories
- Supports both individual (per-scene) and joint (cross-diameter) search

Usage:
    conda run -n sr_env python src/SR_analysis/sindy/run_pysr_deep_illusion.py --individual
    conda run -n sr_env python src/SR_analysis/sindy/run_pysr_deep_illusion.py --joint
    conda run -n sr_env python src/SR_analysis/sindy/run_pysr_deep_illusion.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,
    ILLUSION_PHASE_KEYS,
)
from SR_analysis.configs import get_scene

SINDY_DIR = os.path.join(os.path.dirname(__file__))
RESULTS_DIR = os.path.join(SINDY_DIR, "..", "validate", "results")
os.makedirs(RESULTS_DIR, exist_ok=True)

# Illusion scenes (skip 1.5L - bang-bang regime)
ILLUSION_SCENES = ["illusion_0.75L", "illusion_1L"]

# Illusion phase keys + target force features (for front channel)
# ILLUSION_PHASE_KEYS has 10 dims: u_a, du_a_dt, Cl_tot, dCl_tot_dt, Cd_tot, Cd_rear,
#                                      Cd_err, Cl_err, dCd_err_dt, dCl_err_dt
# For joint search, also include target features (already part of ILLUSION_PHASE_KEYS
# via Cd_err = Cd_tot - target_Cd, etc.)
FRONT_KEYS = list(ILLUSION_PHASE_KEYS)
REAR_KEYS = list(ILLUSION_PHASE_KEYS)  # same features, bias added by build_feature_matrix

# Extended feature set with explicit target_Cd/target_Cl
FRONT_KEYS_WITH_TARGET = FRONT_KEYS + ["target_Cd", "target_Cl"]
REAR_KEYS_WITH_TARGET = REAR_KEYS + ["target_Cd", "target_Cl"]


def load_and_build(scene_name, front_keys, rear_keys, use_mu=False, dt_c=None):
    """Load controlled data and build feature matrix for one illusion scene."""
    cfg = get_scene(scene_name)
    u0 = cfg["u0"]
    mu = cfg["mu"]
    si = cfg.get("sample_interval", 600)
    t0_steps = 2000
    if dt_c is None:
        dt_c = si / t0_steps

    data_dir = os.path.join(SINDY_DIR, "..", "data", "illusion", scene_name)
    npz = np.load(os.path.join(data_dir, "controlled.npz"))
    sensors = npz["sensors"].astype(np.float64)
    forces = npz["forces"].astype(np.float64)
    actions_norm = npz["actions"].astype(np.float64)
    actions_phys = (actions_norm * cfg["action_scale"] + np.array(cfg["action_bias"])) * u0
    target_forces = npz["target_forces"].astype(np.float64) if "target_forces" in npz else None

    T = sensors.shape[0]
    a1 = np.zeros((T, 3), dtype=np.float64)
    a2 = np.zeros((T, 3), dtype=np.float64)
    a1[1:] = actions_phys[:-1]
    a2[2:] = actions_phys[:-2]

    dim = compute_dimensionless(sensors, forces, u0=u0, d=20.0)
    sym = compute_features(
        dim, a1, a2, mu,
        alpha_mode=False, include_mu=use_mu,
        include_cos_sin=False, dt_c=dt_c, u0=u0,
        target_forces=target_forces,
        sensors_raw=sensors, forces_raw=forces,
    )

    n_warmup = 2
    Xf = build_feature_matrix(sym, front_keys, add_bias=False)[n_warmup:]
    Xr = build_feature_matrix(sym, rear_keys, add_bias=True)[n_warmup:]
    Ya = actions_phys[n_warmup:] / u0  # non-dim alpha
    return Xf, Xr, Ya, {"u0": u0, "mu": mu, "cfg": cfg, "scene_name": scene_name,
                        "sensors": sensors, "forces": forces, "target_forces": target_forces,
                        "dt_c": dt_c}


def run_pysr(X, y, feat_names, label, seed=42, niter=120):
    """Run PySR on feature matrix X -> target y."""
    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="", with_target=False):
    """Run deep PySR on each illusion scene independently."""
    print("=" * 60)
    print("INDIVIDUAL Illusion (each with ILLUSION_PHASE_KEYS + target_Cd/target_Cl)")
    print("=" * 60)

    front_keys = FRONT_KEYS_WITH_TARGET if with_target else FRONT_KEYS
    rear_keys = REAR_KEYS_WITH_TARGET if with_target else REAR_KEYS

    all_res = {}
    for sn in ILLUSION_SCENES:
        print(f"\n--- {sn} ---")
        Xf, Xr, Ya, info = load_and_build(sn, front_keys, rear_keys)
        print(f"  Data loaded: Xf={Xf.shape}, Ya={Ya.shape}")

        # Front channel
        res_f = run_pysr(Xf, Ya[:, 0], front_keys, f"{sn} front", niter=120)
        # Top channel (with bias)
        res_t = run_pysr(Xr, Ya[:, 2], ["bias"] + rear_keys, f"{sn} top", niter=120)

        all_res[sn] = {"front": res_f, "top": res_t}

    out_path = os.path.join(RESULTS_DIR, f"pysr_illusion_deep_individual{run_id}.json")
    json.dump(all_res, open(out_path, "w"), indent=2, default=str)
    print(f"\nSaved: {out_path}")

    for sn in ILLUSION_SCENES:
        r = all_res[sn]
        print(f"  {sn}:")
        print(f"    front: {r['front']['best_sympy']}  (R2={r['front']['best_score']:.4f})")
        print(f"    top:   {r['top']['best_sympy']}  (R2={r['top']['best_score']:.4f})")


def run_joint(run_id=""):
    """Run deep joint PySR across both Illusion scenes.

    Method B: no target_diameter marker. Just concatenated data.
    """
    print("=" * 60)
    print("JOINT Illusion (ILLUSION_PHASE_KEYS, both scenes)")
    print("=" * 60)

    Xf_list, Xr_list, Y_list = [], [], []
    for sn in ILLUSION_SCENES:
        Xf, Xr, Ya, info = load_and_build(sn, FRONT_KEYS, REAR_KEYS)
        Xf_list.append(Xf)
        Xr_list.append(Xr)
        Y_list.append(Ya)

    Xf_all = np.vstack(Xf_list)
    Xr_all = np.vstack(Xr_list)
    Y_all = np.vstack(Y_list)
    print(f"  Joint front: {Xf_all.shape}, rear: {Xr_all.shape}")

    res_f = run_pysr(Xf_all, Y_all[:, 0], FRONT_KEYS, "joint front", niter=120)
    res_t = run_pysr(Xr_all, Y_all[:, 2], ["bias"] + REAR_KEYS, "joint top", niter=120)

    out_path = os.path.join(RESULTS_DIR, f"pysr_illusion_joint{run_id}.json")
    json.dump({"front": res_f, "top": res_t}, open(out_path, "w"), indent=2, default=str)
    print(f"\nSaved: {out_path}")
    print(f"  joint front: {res_f['best_sympy']}  (R2={res_f['best_score']:.4f})")
    print(f"  joint top:   {res_t['best_sympy']}  (R2={res_t['best_score']:.4f})")


def run_joint_marker(run_id=""):
    """Joint with target_diameter as extra feature (Method A).

    Concatenate data, add target_diameter column to feature matrix.
    If the optimal formula uses target_diameter, formulas are different-shaped.
    """
    print("=" * 60)
    print("JOINT Illusion with target_diameter marker (Method A)")
    print("=" * 60)

    Xf_list, Xr_list, Y_list = [], [], []
    diam_map = {"illusion_0.75L": 0.75, "illusion_1L": 1.0}

    for sn in ILLUSION_SCENES:
        Xf, Xr, Ya, info = load_and_build(sn, FRONT_KEYS, REAR_KEYS)
        diam = diam_map[sn]
        # Add target_diameter column to both front and rear feature matrices
        diam_col = np.full((Xf.shape[0], 1), diam, dtype=np.float64)
        Xf_d = np.hstack([Xf, diam_col])
        Xr_d = np.hstack([Xr, diam_col])
        Xf_list.append(Xf_d)
        Xr_list.append(Xr_d)
        Y_list.append(Ya)

    Xf_all = np.vstack(Xf_list)
    Xr_all = np.vstack(Xr_list)
    Y_all = np.vstack(Y_list)

    marker_keys = FRONT_KEYS + ["target_diameter"]
    marker_keys_r = ["bias"] + REAR_KEYS + ["target_diameter"]
    print(f"  Joint marker front: {Xf_all.shape}, rear: {Xr_all.shape}")
    print(f"  Front keys: {marker_keys}")

    res_f = run_pysr(Xf_all, Y_all[:, 0], marker_keys, "joint marker front", niter=120)
    res_t = run_pysr(Xr_all, Y_all[:, 2], marker_keys_r, "joint marker top", niter=120)

    out_path = os.path.join(RESULTS_DIR, f"pysr_illusion_joint_marker{run_id}.json")
    json.dump({"front": res_f, "top": res_t}, open(out_path, "w"), indent=2, default=str)
    print(f"\nSaved: {out_path}")
    print(f"  joint marker front: {res_f['best_sympy']}  (R2={res_f['best_score']:.4f})")
    print(f"  joint marker top:   {res_t['best_sympy']}  (R2={res_t['best_score']:.4f})")

    # Check if target_diameter is used in the best formula
    if "target_diameter" in res_f.get("best_sympy", "") or "target_diameter" in res_t.get("best_sympy", ""):
        print("\n  >>> target_diameter IS used: formulas are DIFFERENT-shaped")
    else:
        print("\n  >>> target_diameter NOT used: formulas may share a skeleton")


def show_pareto(path, label):
    """Display Pareto front from a saved result."""
    d = json.load(open(path))
    if "front" in d and isinstance(d["front"], dict) and "equations" in d["front"]:
        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("--joint-marker", action="store_true")
    ap.add_argument("--with-target", action="store_true",
                    help="Include target_Cd/target_Cl as explicit features (12-dim total)")
    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", with_target=args.with_target)
        run_joint("_deep")
        run_joint_marker("_deep")
    elif args.individual:
        run_individual("_deep", with_target=args.with_target)
    elif args.joint:
        run_joint("_deep")
    elif args.joint_marker:
        run_joint_marker("_deep")
    else:
        run_individual("_deep")
        run_joint("_deep")
        run_joint_marker("_deep")
