# analysis_crossre/scripts/phase2_ablation.py
"""Ablation runner: v2 baseline -> v2.1 -> v2.2 -> v2.3 -> v2.4.

Each version differs by exactly one change from the previous.
Run specific versions via --mode.

Usage::

    conda run -n pycuda_3_10 python phase2_ablation.py \\
        --mode all --out-dir output/analysis_crossre/sindy

    conda run -n pycuda_3_10 python phase2_ablation.py \\
        --mode v21 --out-dir output/analysis_crossre/sindy
"""
from __future__ import annotations

import argparse
import json
import os
import sys
from typing import Dict, List, Tuple

import numpy as np

from utils import (
    action_to_physical,
    compute_dimensionless,
    compute_physical_symbols,
    fit_channel,
    print_control_law,
)
from cfg import (
    OUTPUT_DIR,
    RE_CASES_TRAIN,
    ACTION_SCALE,
    ACTION_BIAS,
    U0,
)

THRESHOLDS = [0.0, 0.001, 0.002, 0.005, 0.01, 0.015, 0.02, 0.03, 0.05, 0.1]


def load_case_data(re_code: int) -> Tuple:
    case_dir = os.path.join(OUTPUT_DIR, f"re{re_code}")
    npz_path = os.path.join(case_dir, "controlled.npz")
    if not os.path.isfile(npz_path):
        raise FileNotFoundError(f"Missing {npz_path}")
    data = np.load(npz_path)
    sensors = data["sensors"].astype(np.float64)
    forces = data["forces"].astype(np.float64)
    actions_norm = data["actions"].astype(np.float64)
    rewards = data.get("rewards", np.zeros(sensors.shape[0])).astype(np.float64)
    actions_phys = action_to_physical(
        actions_norm, scale=ACTION_SCALE, bias=ACTION_BIAS, u0=U0)
    mu = 2.0 / re_code
    return sensors, forces, actions_phys, rewards, mu


def make_features_v2(sensors, forces, actions_prev, actions_prev2, include_raw_lattice=True):
    """v2 / v2.1 features: raw lattice physical symbols."""
    if include_raw_lattice:
        sym = compute_physical_symbols(sensors, forces, actions_prev, actions_prev2)
    else:
        sym = {}

    T = sensors.shape[0]

    # Always build v2 physical symbols even for lattice version
    s = sensors.astype(np.float64)
    f = forces.astype(np.float64)
    u0, u1, u2 = s[:, 0], s[:, 2], s[:, 4]
    v0, v1, v2 = s[:, 1], s[:, 3], s[:, 5]

    # Add derived symbols (v2 style)
    sym["u_m"] = (u0 + u1 + u2) / 3.0
    sym["u_a"] = (u2 - u0) / 2.0
    sym["u_c"] = u1.copy()
    sym["u_curv"] = u0 - 2.0 * u1 + u2
    sym["v_m"] = (v0 + v1 + v2) / 3.0
    sym["v_a"] = (v2 - v0) / 2.0
    sym["v_c"] = v1.copy()
    sym["v_curv"] = v0 - 2.0 * v1 + v2
    sym["sin_ua"] = np.sin(np.pi * sym["u_a"])
    sym["cos_ua"] = np.cos(np.pi * sym["u_a"])

    fx0, fy0 = f[:, 0], f[:, 1]
    fx1, fy1 = f[:, 2], f[:, 3]
    fx2, fy2 = f[:, 4], f[:, 5]
    sym["Fx_tot"] = fx0 + fx1 + fx2
    sym["Fx_rear"] = fx1 + fx2
    sym["Fx_diff"] = fx2 - fx1
    sym["Fy_tot"] = fy0 + fy1 + fy2
    sym["Fy_rear"] = fy1 + fy2
    sym["Fy_diff"] = fy2 - fy1

    sym["a0_lag1"] = actions_prev[:, 0]
    sym["a1_lag1"] = actions_prev[:, 1]
    sym["a2_lag1"] = actions_prev[:, 2]
    sym["da0"] = actions_prev[:, 0] - actions_prev2[:, 0]
    sym["da1"] = actions_prev[:, 1] - actions_prev2[:, 1]
    sym["da2"] = actions_prev[:, 2] - actions_prev2[:, 2]

    return sym


def make_features_dimensionless(sensors, forces, actions_prev, actions_prev2, mu):
    """v2.2 features: fully dimensionless."""
    dim = compute_dimensionless(sensors, forces, u0=U0, d=20.0)
    T = actions_prev.shape[0]

    # Nondim actions: alpha = omega_phys / U0
    T = actions_prev.shape[0]
    a_prev = np.zeros((T, 3), dtype=np.float64)
    a_prev2 = np.zeros((T, 3), dtype=np.float64)
    a_prev[1:] = actions_prev[1:] / U0
    a_prev2[2:] = actions_prev2[2:] / U0
    da = a_prev - a_prev2

    # Sensor (nondim)
    u_B, u_C, u_T = dim["u_hat_B"], dim["u_hat_C"], dim["u_hat_T"]
    v_B, v_C, v_T = dim["v_hat_B"], dim["v_hat_C"], dim["v_hat_T"]

    sym = {}
    sym["u_m"] = (u_B + u_C + u_T) / 3.0
    sym["u_a"] = (u_T - u_B) / 2.0
    sym["u_c"] = u_C.copy()
    sym["u_curv"] = u_B - 2.0 * u_C + u_T
    sym["v_m"] = (v_B + v_C + v_T) / 3.0
    sym["v_a"] = (v_T - v_B) / 2.0
    sym["v_c"] = v_C.copy()
    sym["v_curv"] = v_B - 2.0 * v_C + v_T
    sym["sin_ua"] = np.sin(np.pi * sym["u_a"])
    sym["cos_ua"] = np.cos(np.pi * sym["u_a"])

    # Force (nondim Cd/Cl)
    sym["Cd_tot"] = dim["Cd_F"] + dim["Cd_T"] + dim["Cd_B"]
    sym["Cd_rear"] = dim["Cd_T"] + dim["Cd_B"]
    sym["Cd_diff"] = dim["Cd_T"] - dim["Cd_B"]
    sym["Cl_tot"] = dim["Cl_F"] + dim["Cl_T"] + dim["Cl_B"]
    sym["Cl_rear"] = dim["Cl_T"] + dim["Cl_B"]
    sym["Cl_diff"] = dim["Cl_T"] - dim["Cl_B"]

    # Memory (nondim alpha)
    sym["a0_lag1"] = a_prev[:, 0]  # front
    sym["a1_lag1"] = a_prev[:, 1]  # bottom
    sym["a2_lag1"] = a_prev[:, 2]  # top
    sym["da0"] = da[:, 0]
    sym["da1"] = da[:, 1]
    sym["da2"] = da[:, 2]

    # Mu modulation
    sym["mu"] = np.full(T, mu, dtype=np.float64)
    sym["mu_u_a"] = sym["u_a"] * mu
    sym["mu_v_a"] = sym["v_a"] * mu
    sym["mu_Cd_tot"] = sym["Cd_tot"] * mu
    sym["mu_Cl_diff"] = sym["Cl_diff"] * mu

    return sym


def build_theta(sym, feature_keys, add_bias=True):
    """Build feature matrix from symbol dict."""
    T = sym[feature_keys[0]].shape[0]
    cols = []
    if add_bias:
        cols.append(np.ones(T, dtype=np.float64))
    for k in feature_keys:
        cols.append(sym[k])
    return np.column_stack(cols)


# Feature set definitions
V2_BASE_KEYS = [
    "u_m", "u_a", "u_c", "u_curv", "v_a", "v_curv",
    "Fx_tot", "Fx_rear", "Fx_diff", "Fy_tot", "Fy_diff",
    "sin_ua", "cos_ua",
    "a0_lag1", "a1_lag1", "a2_lag1",
    "da0", "da1", "da2",
]

V2_WITH_MU = V2_BASE_KEYS + [
    "mu", "mu_u_a", "mu_v_a", "mu_Cd_tot", "mu_Cl_diff",
]

V2DIM_KEYS = [
    "u_m", "u_a", "u_c", "u_curv", "v_a", "v_curv",
    "Cd_tot", "Cd_rear", "Cd_diff", "Cl_tot", "Cl_diff",
    "sin_ua", "cos_ua",
    "a0_lag1", "a1_lag1", "a2_lag1",
    "da0", "da1", "da2",
    "mu", "mu_u_a", "mu_v_a", "mu_Cd_tot", "mu_Cl_diff",
]


def build_dataset(re_codes, mode, use_mu=True, use_mu_nondim=True):
    """Build dataset for given mode.

    Modes:
    - "v2_baseline": raw lattice + all 3 with bias
    - "v21": same but front no-bias
    - "v22": dimensionless + front no-bias  
    - "v23": v22 + rear shared head
    - "v24": v23 + mild weighting
    """
    all_Theta = []
    all_Y = []
    all_W = []
    all_re = []

    for rc in re_codes:
        sensors, forces, actions_phys, rewards, mu = load_case_data(rc)

        if mode in ("v2_baseline", "v21"):
            # raw lattice features
            a_prev = np.zeros_like(actions_phys)
            a_prev2 = np.zeros_like(actions_phys)
            a_prev[1:] = actions_phys[:-1]
            a_prev2[2:] = actions_phys[:-2]
            sym = make_features_v2(sensors, forces, a_prev, a_prev2)
            feature_keys = V2_WITH_MU if use_mu else V2_BASE_KEYS
            sym["mu"] = np.full(sensors.shape[0], mu, dtype=np.float64)
            sym["mu_u_a"] = sym["u_a"] * mu
            sym["mu_v_a"] = sym["v_a"] * mu
            if use_mu:
                # Add mu modulated forces
                sym["mu_Cd_tot"] = sym["Fx_tot"] * mu
                sym["mu_Cl_diff"] = sym["Fy_diff"] * mu
            Y = actions_phys.copy()
        else:
            # dimensionless features
            a_prev_d = np.zeros_like(actions_phys)
            a_prev2_d = np.zeros_like(actions_phys) 
            a_prev_d[1:] = actions_phys[:-1]
            a_prev2_d[2:] = actions_phys[:-2]
            sym = make_features_dimensionless(sensors, forces, a_prev_d, a_prev2_d, mu)
            feature_keys = V2DIM_KEYS
            # Y is nondim alpha = omega/U0
            Y = actions_phys / U0

        # Compute quality weight if needed
        if mode == "v24":
            late_mean = float(np.mean(rewards[-80:]))
            weight = np.clip(0.3 + 0.7 * late_mean / 0.7, 0.2, 1.0)
            W = np.full(sensors.shape[0], weight, dtype=np.float64)
        else:
            W = np.ones(sensors.shape[0], dtype=np.float64)

        # Store for stacking (will trim warmup later)
        all_Theta.append((sym, feature_keys, Y, W, rc))

    # Stack all Re data with warmup removed
    Theta_list = []
    Y_list = []
    W_list = []
    re_list = []

    for sym, feature_keys, Y, W, rc in all_Theta:
        T = Y.shape[0]
        theta = build_theta(sym, feature_keys, add_bias=True)
        # Remove first 2 warmup steps
        theta = theta[2:]
        Y_t = Y[2:]
        W_t = W[2:]

        Theta_list.append(theta)
        Y_list.append(Y_t)
        W_list.append(W_t)
        re_list.append(np.full(theta.shape[0], rc, dtype=np.int64))

    Theta_stacked = np.vstack(Theta_list)
    Y_stacked = np.vstack(Y_list)
    W_stacked = np.concatenate(W_list)
    Re_stacked = np.concatenate(re_list)

    # For front no-bias versions: remove bias column (column 0)
    front_bias = mode not in ("v21", "v22", "v23", "v24")
    
    if front_bias:
        Theta_front = Theta_stacked
        Theta_other = Theta_stacked
    else:
        Theta_front = Theta_stacked[:, 1:]  # remove bias column for front
        Theta_other = Theta_stacked  # keep bias for bottom/top

    return Theta_front, Theta_other, Y_stacked, W_stacked, Re_stacked


def fit_weighted(Theta, y, w, thresholds):
    """Weighted STLSQ fit."""
    import pysindy as ps
    std = np.sqrt(np.average((Theta - np.average(Theta, axis=0, weights=w))**2, 
                              axis=0, weights=w))
    std = np.where(std < 1e-8, 1.0, std)
    Theta_s = Theta / std
    best = None
    rows = []
    for th in thresholds:
        opt = ps.STLSQ(threshold=th, alpha=1e-4, max_iter=25)
        opt.fit(Theta_s, y, sample_weight=w)
        coef = np.asarray(opt.coef_, dtype=np.float64).flatten() / std
        y_pred = Theta @ coef
        y_mean = np.average(y, weights=w)
        ssr = np.sum(w * (y - y_pred)**2)
        sst = np.sum(w * (y - y_mean)**2) + 1e-12
        r2 = 1.0 - ssr / sst
        mae = float(np.average(np.abs(y - y_pred), weights=w))
        nz = int(np.sum(np.abs(coef) > 1e-8))
        entry = {"threshold": float(th), "nz": nz, "r2": r2, "mae": mae, "coef": coef}
        rows.append(entry)
        if best is None or r2 > best["r2"]:
            best = entry
    return rows, best


def run_ablation(mode, train_re, out_dir):
    """Run full ablation for given mode."""
    print(f"\n{'='*60}")
    print(f"Mode: {mode}")
    print(f"{'='*60}")

    ThetaF, ThetaO, Y, W, Re = build_dataset(train_re, mode)

    # Determine which cylinders use which feature matrix
    if mode == "v23":
        # rear shared-head: only fit front and top. bottom = -top(Gx)
        # For simplicity, still fit all 3 but check rear consistency separately
        cylinders = [
            ("front", ThetaF, False),  # front: no bias
            ("bottom", ThetaO, True),  # bottom: has bias  
            ("top", ThetaO, True),     # top: has bias
        ]
    elif mode in ("v21", "v22", "v24"):
        cylinders = [
            ("front", ThetaF, False),  # front: no bias
            ("bottom", ThetaO, True),  # bottom: has bias
            ("top", ThetaO, True),     # top: has bias
        ]
    else:  # v2_baseline
        cylinders = [
            ("front", ThetaO, True),   # front: has bias
            ("bottom", ThetaO, True),  # bottom: has bias
            ("top", ThetaO, True),     # top: has bias
        ]

    channels = []
    for name, theta, has_bias in cylinders:
        ci = {"front": 0, "bottom": 1, "top": 2}[name]
        print(f"\n  --- {name} ---")
        rows, best = fit_weighted(theta, Y[:, ci], W, THRESHOLDS)
        coef = best["coef"]
        nz = int(np.sum(np.abs(coef) > 1e-8))
        print(f"  {name}: R2={best['r2']:.6f}  MAE={best['mae']:.6f}  nz={nz}")
    # Get feature names for this mode
    if mode in ("v2_baseline", "v21"):
        feat_names = [
            "bias", "u_m", "u_a", "u_c", "u_curv", "v_a", "v_curv",
            "Fx_tot", "Fx_rear", "Fx_diff", "Fy_tot", "Fy_diff",
            "sin_ua", "cos_ua",
            "a0_lag1", "a1_lag1", "a2_lag1",
            "da0", "da1", "da2",
            "mu", "mu_u_a", "mu_v_a", "mu_Cd_tot", "mu_Cl_diff",
        ]
        has_mu = True
    else:
        feat_names = [
            "bias", "u_m", "u_a", "u_c", "u_curv", "v_a", "v_curv",
            "Cd_tot", "Cd_rear", "Cd_diff", "Cl_tot", "Cl_diff",
            "sin_ua", "cos_ua",
            "a0_lag1", "a1_lag1", "a2_lag1",
            "da0", "da1", "da2",
            "mu", "mu_u_a", "mu_v_a", "mu_Cd_tot", "mu_Cl_diff",
        ]
        has_mu = True

    # Trim feat_names to match actual theta dimensions
    actual_nf = theta.shape[1]
    if len(feat_names) != actual_nf:
        # Feature names don't include mu if not included, etc.
        # Just use generic names
        feat_names = [f"f{i}" for i in range(actual_nf)]

    # Per-Re breakdown
    print(f"\n  --- Per-Re breakdown ---")
    breakdown = {}
    for rc in set(Re.tolist()):
        mask = Re == rc
        ch_b = []
        for name, theta, has_bias in cylinders:
            ci = {"front": 0, "bottom": 1, "top": 2}[name]
            th_r = theta[mask]
            yr = Y[mask, ci]
            wr = W[mask]
            coef = np.array([ch["best_coef"][ci] for ch in channels], dtype=np.float64).flatten()
            # Actually get the right coefficient for this cylinder
            coef_c = np.array(channels[ci]["best_coef"], dtype=np.float64)
            y_pred = th_r @ coef_c
            y_mean = np.average(yr, weights=wr)
            ssr = np.sum(wr * (yr - y_pred)**2)
            sst = np.sum(wr * (yr - y_mean)**2) + 1e-12
            r2 = 1.0 - ssr / sst
            mae = float(np.average(np.abs(yr - y_pred), weights=wr))
            ch_b.append({"cylinder": name, "r2": float(r2), "mae": mae})
        breakdown[f"re{int(rc)}"] = ch_b
        r2s = ", ".join([f"{b['cylinder']}={b['r2']:.4f}" for b in ch_b])
        print(f"    Re{int(rc)}: {r2s}")

    return {
        "mode": mode,
        "train_re": train_re,
        "channels": channels,
        "per_re_breakdown": breakdown,
    }


def main():
    ap = argparse.ArgumentParser(description="Ablation runner v2->v2.4")
    ap.add_argument("--mode", type=str, default="all",
                    choices=["v2_baseline", "v21", "v22", "v23", "v24", "all"])
    ap.add_argument("--out-dir", type=str, default=os.path.join(OUTPUT_DIR, "sindy"))
    ap.add_argument("--train-re", type=str, default="50,100,200")
    args = ap.parse_args()

    train_re = [int(r) for r in args.train_re.split(",")]
    os.makedirs(args.out_dir, exist_ok=True)

    modes = ["v2_baseline", "v21", "v22", "v23", "v24"] if args.mode == "all" else [args.mode]

    results = {"metadata": {"thresholds": THRESHOLDS, "train_re": train_re}}
    for mode in modes:
        results[mode] = run_ablation(mode, train_re, args.out_dir)

    out_path = os.path.join(args.out_dir, "ablation_results.json")
    with open(out_path, "w") as f:
        json.dump(results, f, indent=2)
    print(f"\nSaved: {out_path}")


if __name__ == "__main__":
    main()