#!/usr/bin/env python3
"""Stage 2 V5: PySR symbolic regression on V5-trained PPO models.

Loads V5 controlled.npz from SR_analysis/data/v5/, uses V5 action decoder
(action_scale=12, action_bias=[0,0,0]), and outputs formulas to
SR_analysis/results/formulas_v5/.

Supports per-scene individual and cross-scene joint fitting.

Usage:
    # Per-scene (fast, niter=40)
    PYTHONPATH=src conda run -n sr_env python -m SR_analysis.experiments.v5.stage_2_fit_v5 --scene kar_re100

    # Per-scene deep search (niter=120)
    PYTHONPATH=src conda run -n sr_env python -m SR_analysis.experiments.v5.stage_2_fit_v5 --scene kar_re100 --deep

    # Joint: all effective Karman cases
    PYTHONPATH=src conda run -n sr_env python -m SR_analysis.experiments.v5.stage_2_fit_v5 --group karman_v5 --mode joint

    # Joint: all Illusion cases
    PYTHONPATH=src conda run -n sr_env python -m SR_analysis.experiments.v5.stage_2_fit_v5 --group illusion_v5 --mode joint

    # Joint: mixed V5 + legacy (cross-solver)
    PYTHONPATH=src conda run -n sr_env python -m SR_analysis.experiments.v5.stage_2_fit_v5 --scenes kar_re100,kar_d075,kar_d15,karman_re100,karman_re200 --mode joint --output-label mixed_karman
"""
from __future__ import annotations

import argparse
import json
import os
import sys
from pathlib import Path

import numpy as np

_SR_ROOT = Path(__file__).resolve().parents[2]
_REPO = _SR_ROOT.parents[1]
if str(_REPO) not in sys.path:
    sys.path.insert(0, str(_REPO))
_SRC = _REPO / "src"
if str(_SRC) not in sys.path:
    sys.path.insert(0, str(_SRC))

from SR_analysis.utils.feature_builder import (
    compute_dimensionless, compute_features, build_feature_matrix,
    PHASE_STATE_KEYS, ILLUSION_PHASE_KEYS, CORE_FEAT_KEYS_V2, MU_FEAT_KEYS,
)
from SR_analysis.configs import get_scene  # for legacy scenes
from drl_pinball.case_registry import CASE_REGISTRY

ALL_MU = list(MU_FEAT_KEYS) + ["mu_Cl_tot"]
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"]

# V5 action decoder (all cases: no-bias, scale=12)
V5_ACTION_SCALE = 12.0
V5_ACTION_BIAS = (0.0, 0.0, 0.0)
V5_U0 = 0.01

# V5 scene configs (for scenes without legacy configs.py entry)
V5_KARMAN_CFG = {
    "scene_id": "karman", "u0": V5_U0, "mu": 2.0 / 100,  # Re=100
    "sample_interval": 800, "action_scale": V5_ACTION_SCALE,
    "action_bias": V5_ACTION_BIAS,
}
V5_ILLUSION_CFG = {
    "scene_id": "illusion", "u0": V5_U0, "mu": 2.0 / 100,  # Re=100
    "sample_interval": 600, "action_scale": V5_ACTION_SCALE,
    "action_bias": V5_ACTION_BIAS,
}

# V5 case attributes come from drl_pinball.case_registry.

def is_v5_scene(name: str) -> bool:
    return name in CASE_REGISTRY


def load_controlled_v5(case_name: str):
    """Load controlled.npz from V5 data directory with V5 action decoder."""
    data_dir = _SR_ROOT / "data" / "v5" / case_name
    config_path = data_dir / "config.json"

    npz = np.load(data_dir / "controlled.npz", allow_pickle=True)
    s = npz["sensors"].astype(np.float64)
    f = npz["forces"].astype(np.float64)
    a = npz["actions"].astype(np.float64)  # normalized [-1,1]

    # Load config to get SI
    cfg = V5_KARMAN_CFG.copy() if "ill_" not in case_name else V5_ILLUSION_CFG.copy()
    if config_path.exists():
        with open(config_path) as fp:
            v5_cfg = json.load(fp)
        cfg["si_actual"] = v5_cfg.get("SI", 800)
        cfg["scene_type"] = v5_cfg.get("scene_type", "")

    # Apply canonical training/evaluation SI.
    si = CASE_REGISTRY[case_name].si
    cfg["sample_interval"] = si
    cfg["si_actual"] = si  # for feature builder

    # Convert normalized action to physical surface velocity
    ap = (a * V5_ACTION_SCALE + np.array(V5_ACTION_BIAS)) * V5_U0
    return s, f, ap, None, cfg  # No target_forces for V5 karman


def load_controlled_any(scene_name: str):
    """Load from either V5 or legacy based on scene name."""
    if is_v5_scene(scene_name):
        return load_controlled_v5(scene_name)
    else:
        # Legacy scene - use original configs.py
        cfg = get_scene(scene_name)
        sid = cfg["scene_id"]
        data_dir = f"src/SR_analysis/data/{sid}/{scene_name}"
        npz = np.load(os.path.join(data_dir, "controlled.npz"), allow_pickle=True)
        s = npz["sensors"].astype(np.float64)
        f = npz["forces"].astype(np.float64)
        a = npz["actions"].astype(np.float64)
        ap = (a * cfg["action_scale"] + np.array(cfg["action_bias"])) * cfg["u0"]
        tf = npz["target_forces"].astype(np.float64) if "target_forces" in npz else None
        return s, f, ap, tf, cfg


def build_features_v5(s, f, ap, cfg, feat_keys):
    """Build feature matrix from raw V5 data."""
    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=cfg["u0"], d=20.0)
    si = cfg.get("sample_interval", 800)
    dt_c = si / 2000.0
    mu = cfg["mu"]
    sym = compute_features(dim, a1, a2, mu, alpha_mode=False,
                           include_mu="mu" in feat_keys,
                           include_cos_sin=False, u0=cfg["u0"],
                           target_forces=None, dt_c=dt_c,
                           sensors_raw=s, forces_raw=f)
    X = build_feature_matrix(sym, feat_keys, add_bias=False)[2:]
    Y = ap[2:] / cfg["u0"]
    return X, Y


FORMULA_DIR = _SR_ROOT / "results" / "formulas_v5"
os.makedirs(FORMULA_DIR, exist_ok=True)

FIT_GROUPS_V5 = {
    "karman_v5": ["kar_re100", "kar_d075", "kar_d15", "kar_d2"],
    "illusion_v5": ["ill_075L", "ill_1L", "ill_15L", "ill_2L"],
}


def get_feature_keys_v5(scene_name: str):
    if "ill_" in scene_name:
        return ILLUSION_PHASE_KEYS
    return PHYS_DADT + ALL_MU


def fit_v5(scene_name, feat_keys, niter, output_label):
    s, f, ap, tf, cfg = load_controlled_any(scene_name)
    X, Y = build_features_v5(s, f, ap, cfg, feat_keys)
    from pysr import PySRRegressor
    m = PySRRegressor(binary_operators=["+","-","*","/"], unary_operators=["square"],
                      niterations=niter, populations=30, maxsize=15,
                      complexity_of_constants=2, parsimony=0.01, batching=False)
    m.fit(X, Y[:, 0], variable_names=feat_keys)
    front = {"scene": output_label, "channel": "front", "output": "alpha",
             "feature_keys": feat_keys, "best_sympy": str(m.sympy()),
             "best_score": float(m.score(X, Y[:, 0])), "niterations": niter}
    json.dump(front, open(os.path.join(FORMULA_DIR, f"{output_label}_front.json"), "w"), indent=2)
    print(f"  Front: {front['best_sympy']}  (R2={front['best_score']:.4f})")

    # Top (with bias)
    m2 = PySRRegressor(binary_operators=["+","-","*","/"], unary_operators=["square"],
                       niterations=niter, populations=30, maxsize=15,
                       complexity_of_constants=2, parsimony=0.01, batching=False)
    Xt = build_feature_matrix(
        dict(zip(feat_keys, [X[:, i] for i in range(len(feat_keys))])),
        feat_keys, add_bias=True)
    if isinstance(Xt, list):
        Xt = Xt[0]
    if Xt.ndim == 1:
        Xt = Xt.reshape(-1, 1)
    m2.fit(Xt, Y[:, 2], variable_names=["bias"] + feat_keys)
    top = {"scene": output_label, "channel": "top", "output": "alpha",
           "feature_keys": ["bias"] + feat_keys, "best_sympy": str(m2.sympy()),
           "best_score": float(m2.score(Xt, Y[:, 2])), "niterations": niter}
    json.dump(top, open(os.path.join(FORMULA_DIR, f"{output_label}_top.json"), "w"), indent=2)
    print(f"  Top: {top['best_sympy']}  (R2={top['best_score']:.4f})")
    return front, top


def fit_joint_v5(scene_names, feat_keys, niter, output_label):
    Xl, Yl = [], []
    for sn in scene_names:
        s, f, ap, tf, cfg = load_controlled_any(sn)
        X, Y = build_features_v5(s, f, ap, cfg, feat_keys)
        Xl.append(X); Yl.append(Y[:, 0])
    X_all = np.vstack(Xl); Y_all = np.concatenate(Yl)
    from pysr import PySRRegressor
    m = PySRRegressor(binary_operators=["+","-","*","/"], unary_operators=["square"],
                      niterations=niter, populations=40, maxsize=15,
                      complexity_of_constants=2, parsimony=0.01, batching=False)
    m.fit(X_all, Y_all, variable_names=feat_keys)
    front = {"scene": output_label, "channel": "front", "output": "alpha",
             "feature_keys": feat_keys, "best_sympy": str(m.sympy()),
             "best_score": float(m.score(X_all, Y_all)), "niterations": niter,
             "n_samples": int(X_all.shape[0])}
    json.dump(front, open(os.path.join(FORMULA_DIR, f"{output_label}_front.json"), "w"), indent=2)
    print(f"  Joint Front: {front['best_sympy']}  (R2={front['best_score']:.4f})")
    return front


def main():
    ap = argparse.ArgumentParser(description="Stage 2 V5: PySR fitting on V5 models")
    ap.add_argument("--scene", type=str, default=None)
    ap.add_argument("--scenes", type=str, default=None)
    ap.add_argument("--group", type=str, default=None)
    ap.add_argument("--mode", type=str, default="per-scene", choices=["per-scene", "joint"])
    ap.add_argument("--deep", action="store_true", help="Deep search (niter=120)")
    ap.add_argument("--output-label", type=str, default=None)
    ap.add_argument("--feat-keys", type=str, default=None,
                    help="Comma-separated feat keys override (e.g. PHYS_DADT)")
    args = ap.parse_args()

    niter = 120 if args.deep else 40

    if args.group and args.group in FIT_GROUPS_V5:
        scene_names = FIT_GROUPS_V5[args.group]
    elif args.scenes:
        scene_names = [s.strip() for s in args.scenes.split(",")]
    elif args.scene:
        scene_names = [args.scene]
    else:
        ap.error("Need --scene, --scenes, or --group")

    label = args.output_label or ("_".join(scene_names) if len(scene_names) > 1 else scene_names[0])

    if args.mode == "per-scene":
        for sn in scene_names:
            keys = get_feature_keys_v5(sn)
            if args.feat_keys:
                keys = [k.strip() for k in args.feat_keys.split(",")]
            print(f"\n=== {sn} (niter={niter}, {len(keys)} features) ===")
            fit_v5(sn, keys, niter, sn)
    elif args.mode == "joint":
        keys = get_feature_keys_v5(scene_names[0])
        if args.feat_keys:
            keys = [k.strip() for k in args.feat_keys.split(",")]
        print(f"\n=== Joint {label} (niter={niter}, {len(keys)} features, {len(scene_names)} scenes) ===")
        fit_joint_v5(scene_names, keys, niter, label)

    print("\nDone.")


if __name__ == "__main__":
    main()
