#!/usr/bin/env python3
"""Restricted PySR on phase-state features.

Usage:
    conda run -n sr_env python src/SR_analysis/sindy/run_pysr.py --scene illusion_1L --phase
    conda run -n sr_env python src/SR_analysis/sindy/run_pysr.py --scene karman_re100 --phase
    conda run -n sr_env python src/SR_analysis/sindy/run_pysr.py --scene karman_re100 --phase --include-mu
    conda run -n sr_env python src/SR_analysis/sindy/run_pysr.py --scene all --phase
"""
from __future__ import annotations

import argparse
import json
import os
import sys
import traceback
from typing import Dict, List, Optional

import numpy as np

# PySR
from pysr import PySRRegressor

_REPO = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..", ".."))
sys.path.insert(0, _REPO)
_SRC = os.path.join(_REPO, "src")
sys.path.insert(0, _SRC)

from SR_analysis.utils.feature_builder import (
    compute_dimensionless, compute_features, build_feature_matrix,
    PHASE_STATE_KEYS, ILLUSION_PHASE_KEYS, MU_FEAT_KEYS,
)
from SR_analysis.configs import get_scene

SINDY_DIR = os.path.join(os.path.dirname(__file__))

# All mu keys produced by compute_features(include_mu=True)
# (MU_FEAT_KEYS has 5, but compute_features generates an additional mu_Cl_tot)
ALL_MU_KEYS = list(MU_FEAT_KEYS) + ["mu_Cl_tot"]


def load_controlled_data(scene_name: str):
    """Load controlled data for a scene."""
    cfg = get_scene(scene_name)
    scene_id = cfg["scene_id"]
    data_dir = os.path.join(SINDY_DIR, "..", "data", scene_id, 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"])) * cfg["u0"]
    target_forces = npz["target_forces"].astype(np.float64) if "target_forces" in npz else None
    return sensors, forces, actions_phys, target_forces, cfg


def get_feature_keys(scene_id: str, use_phase: bool, include_mu: bool) -> List[str]:
    """Determine feature keys for a scene."""
    if not use_phase:
        # Fallback: use CORE_FEAT_KEYS_V2 style (old approach)
        from SR_analysis.utils.feature_builder import CORE_FEAT_KEYS_V2
        keys = list(CORE_FEAT_KEYS_V2)
    elif scene_id == "illusion":
        keys = list(ILLUSION_PHASE_KEYS)
    else:
        # karman, steady, vortex
        keys = list(PHASE_STATE_KEYS)

    if include_mu:
        keys += ALL_MU_KEYS
    return keys


def run_pysr_scene(scene_name: str, out_dir: str, use_phase: bool, include_mu: bool):
    """Run PySR on a single scene."""
    cfg = get_scene(scene_name)
    scene_id = cfg["scene_id"]
    u0 = cfg["u0"]
    sample_interval = cfg.get("sample_interval", 800)
    t0_steps = 2000  # T0 = D/U0 = 20/0.01 = 2000 LBM steps
    dt_c = sample_interval / t0_steps

    # Determine feature keys
    front_keys = get_feature_keys(scene_id, use_phase, include_mu)
    rear_keys = list(front_keys)  # rear shares same physics features; bias added by build_feature_matrix

    print(f"\n{'='*60}")
    print(f"PySR: {scene_name}  (scene_id={scene_id}, phase={use_phase}, mu={include_mu})")
    print(f"  Front keys ({len(front_keys)}): {front_keys}")
    print(f"  Rear keys ({len(rear_keys)}): {rear_keys}")
    print(f"{'='*60}")

    # Load data
    sensors, forces, actions_phys, target_forces, cfg = load_controlled_data(scene_name)
    mu_val = cfg["mu"]

    # Build features with PROPER lags and sensors_raw/forces_raw for derivatives
    T = sensors.shape[0]
    a_prev = np.zeros((T, 3), dtype=np.float64)
    a_prev2 = np.zeros((T, 3), dtype=np.float64)
    a_prev[1:] = actions_phys[:-1]
    a_prev2[2:] = actions_phys[:-2]

    dim = compute_dimensionless(sensors, forces, u0=u0, d=20.0)
    sym = compute_features(
        dim, a_prev, a_prev2, mu_val,
        alpha_mode=False, include_mu=include_mu,
        include_cos_sin=False, u0=u0,
        target_forces=target_forces,
        dt_c=dt_c,
        sensors_raw=sensors,     # CRITICAL: enables du_a_dt, dCl_tot_dt, etc.
        forces_raw=forces,       # CRITICAL: enables derivative features
    )

    n_warmup = 2
    X_front = build_feature_matrix(sym, front_keys, add_bias=False)[n_warmup:]
    X_rear = build_feature_matrix(sym, rear_keys, add_bias=True)[n_warmup:]
    Y = actions_phys[n_warmup:]  # physical omega (target)

    # Also compute non-dimensional alpha for reference
    Y_alpha = Y / u0

    print(f"  Samples: {X_front.shape[0]}")
    print(f"  Y (omega) range: [{Y.min():.6f}, {Y.max():.6f}]")
    print(f"  alpha range: [{Y_alpha.min():.6f}, {Y_alpha.max():.6f}]")

    # Quick sanity check on derivative features before running PySR
    for key in front_keys:
        val = sym.get(key, None)
        if val is not None:
            print(f"    {key}: mean={np.mean(val):.6f}, std={np.std(val):.6f}")
        else:
            print(f"    {key}: MISSING from sym dict!")

    # Suffix for output filename
    suffix = "_mu" if include_mu else ""

    # --- Front channel (no bias, predict alpha_F) ---
    print(f"\n=== PySR: {scene_name} Front (predict alpha_F) ===")
    model_f = PySRRegressor(
        binary_operators=["+", "-", "*", "/"],
        unary_operators=["square"],
        niterations=40,
        populations=30,
        maxsize=15,
        complexity_of_constants=2,
        parsimony=0.01,
        extra_sympy_mappings={},
        batching=False,
    )
    model_f.fit(X_front, Y_alpha[:, 0], variable_names=front_keys)
    results = {
        "scene": scene_name,
        "scene_id": scene_id,
        "channel": "front",
        "output": "alpha",
        "feature_names": front_keys,
        "equations": model_f.equations_.to_dict(orient="records"),
        "best_sympy": str(model_f.sympy()),
        "best_score": float(model_f.score(X_front, Y_alpha[:, 0])),
    }
    out_path = os.path.join(out_dir, f"pysr_{scene_name}_front{suffix}.json")
    with open(out_path, "w") as f:
        json.dump(results, f, indent=2, default=str)
    print(f"  Best sympy: {model_f.sympy()}")
    print(f"  Score (R2): {results['best_score']:.4f}")
    print(f"  Saved: {out_path}")

    # --- Rear/Top channel (with bias, predict alpha_T) ---
    print(f"\n=== PySR: {scene_name} Top (predict alpha_T, v23 shared-head) ===")
    model_t = PySRRegressor(
        binary_operators=["+", "-", "*", "/"],
        unary_operators=["square"],
        niterations=40,
        populations=30,
        maxsize=15,
        complexity_of_constants=2,
        parsimony=0.01,
        extra_sympy_mappings={},
        batching=False,
    )
    model_t.fit(X_rear, Y_alpha[:, 2], variable_names=["bias"] + rear_keys)
    results_t = {
        "scene": scene_name,
        "scene_id": scene_id,
        "channel": "top",
        "output": "alpha",
        "feature_names": ["bias"] + rear_keys,
        "equations": model_t.equations_.to_dict(orient="records"),
        "best_sympy": str(model_t.sympy()),
        "best_score": float(model_t.score(X_rear, Y_alpha[:, 2])),
    }
    out_path = os.path.join(out_dir, f"pysr_{scene_name}_top{suffix}.json")
    with open(out_path, "w") as f:
        json.dump(results_t, f, indent=2, default=str)
    print(f"  Best sympy: {model_t.sympy()}")
    print(f"  Score (R2): {results_t['best_score']:.4f}")
    print(f"  Saved: {out_path}")

    # Verify front formula is no-bias (action ≈ 0 when features ≈ 0)
    zero_vec = np.zeros((1, len(front_keys)), dtype=np.float64)
    pred_zero = float(model_f.predict(zero_vec)[0])
    print(f"\n  Sanity check: front prediction at zero-input = {pred_zero:.6f} (should be ~0)")


def main():
    ap = argparse.ArgumentParser(
        description="PySR on phase-state features (uses phase-state by default)"
    )
    ap.add_argument("--scene", type=str, default="all",
                    help='Scene name like "karman_re100" or "illusion_1L", or "all"')
    ap.add_argument("--out", type=str, default=None)
    ap.add_argument("--phase", action="store_true", default=True,
                    help="Use phase-state features (default: True)")
    ap.add_argument("--no-phase", action="store_false", dest="phase",
                    help="Use CORE_FEAT_KEYS_V2 instead of phase-state")
    ap.add_argument("--include-mu", action="store_true", default=False,
                    help="Include mu (1/Re) modulation features (Karman cross-Re)")
    args = ap.parse_args()

    if args.out is None:
        args.out = os.path.join(SINDY_DIR, "..", "validate", "results")
    os.makedirs(args.out, exist_ok=True)

    # Default scenes (skip 1.5L — bang-bang regime)
    if args.scene == "all":
        scenes = ["illusion_0.75L", "illusion_1L", "karman_re100"]
    else:
        scenes = [args.scene]

    for sn in scenes:
        try:
            run_pysr_scene(sn, args.out, args.phase, args.include_mu)
        except Exception as e:
            print(f"\n  ERROR on {sn}: {e}")
            traceback.print_exc()

    print("\nDone.")


if __name__ == "__main__":
    main()
