# analysis_crossre/scripts/phase2_control_fit.py
"""Phase 2 v3: dimensionless + front-no-bias + quality-weighted SINDy fitting.

Usage::

    conda run -n pycuda_3_10 python phase2_control_fit.py \\
        --cross-re --out-dir output/analysis_crossre/sindy

    conda run -n pycuda_3_10 python phase2_control_fit.py \\
        --leave-one-out --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_v3_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:
    """Load controlled NPZ for a single Re.

    Returns (sensors, forces, actions_phys, rewards, mu).
    """
    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  # 1 / Re_D
    return sensors, forces, actions_phys, rewards, mu


def compute_trajectory_weights(rewards: np.ndarray, late_window: int = 80) -> float:
    """Compute a single quality weight for this trajectory.

    Uses the mean reward over the last ``late_window`` steps.
    Maps to weight via quantile-based scheme.
    """
    n = len(rewards)
    if n < late_window:
        late_mean = float(np.mean(rewards))
    else:
        late_mean = float(np.mean(rewards[-late_window:]))

    # Map reward to weight via sigmoid-like scheme:
    # reward 0.0 -> weight 0.1, reward 0.3 -> 0.3, reward 0.5 -> 0.6, reward 0.7 -> 0.9
    weight = 1.0 / (1.0 + np.exp(-8.0 * (late_mean - 0.4)))
    return float(np.clip(weight, 0.05, 1.0))


def build_dataset_v3(
    re_code: int,
    include_mu: bool = True,
) -> Tuple:
    """Build v3 data: dimensionless features, front-no-bias, quality-weighted.

    Returns
    -------
    Theta_front : (T, nf_f) for front model (no bias)
    Theta_other : (T, nf_o) for top/bottom model (with bias)
    Y : (T, 3) physical omegas
    W : (T,) quality weight per sample
    names : feature names (without "bias")
    """
    sensors, forces, actions_phys, rewards, mu = load_case_data(re_code)

    # Dimensionless
    dim = compute_dimensionless(sensors, forces, u0=U0, d=20.0)

    # Memory terms
    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]

    # Build v3 features
    Theta_f, Theta_top, names = compute_v3_symbols(
        dim, a_prev, a_prev2, mu=mu, include_mu=include_mu)

    # Quality weight per sample: inherit trajectory weight
    traj_weight = compute_trajectory_weights(rewards)
    W = np.full(Theta_f.shape[0], traj_weight, dtype=np.float64)

    # Remove warmup (need 2 steps of memory)
    Theta_f = Theta_f[2:]
    Theta_top = Theta_top[2:]
    Y = actions_phys[2:]
    W = W[2:]

    print(f"  Re{re_code}: {Theta_f.shape[0]} samples, {Theta_f.shape[1]} feats, "
          f"mu={mu:.6f}, traj_weight={traj_weight:.4f}")
    return Theta_f, Theta_top, Y, W, names, mu


def fit_channel_weighted(
    Theta: np.ndarray,
    y: np.ndarray,
    w: np.ndarray,
    thresholds: list,
    alpha: float = 1e-4,
    max_iter: int = 25,
) -> tuple:
    """Weighted STLSQ fit."""
    import pysindy as ps

    # Weighted normalisation
    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=alpha, max_iter=max_iter)
        opt.fit(Theta_s, y, sample_weight=w)
        coef = np.asarray(opt.coef_, dtype=np.float64).flatten() / std
        y_pred = Theta @ coef
        # Weighted R2
        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 build_cross_re_v3(
    train_re_codes: List[int],
    include_mu: bool = True,
) -> Tuple:
    """Stack multiple Re datasets with v3 features.

    Front model uses Theta_front (no bias).
    Top/Bottom models use Theta_other (with bias).

    Returns three stacked datasets.
    """
    all_ThetaF, all_ThetaO, all_Y, all_W, all_re = [], [], [], [], []
    names = None

    for rc in train_re_codes:
        tf, to, y, w, fn, mu = build_dataset_v3(rc, include_mu=include_mu)
        all_ThetaF.append(tf)
        all_ThetaO.append(to)
        all_Y.append(y)
        all_W.append(w)
        all_re.append(np.full(tf.shape[0], rc, dtype=np.int64))
        if names is None:
            names = fn

    ThetaF = np.vstack(all_ThetaF)
    ThetaO = np.vstack(all_ThetaO)
    Y = np.vstack(all_Y)
    W = np.concatenate(all_W)
    Re = np.concatenate(all_re)

    print(f"\n  Cross-Re: {ThetaF.shape[0]} samples, "
          f"front={ThetaF.shape[1]} feats (no bias), "
          f"other={ThetaO.shape[1]} feats (w/ bias)")
    return ThetaF, ThetaO, Y, W, names, Re


def run_cross_re_fit(
    train_re: List[int],
    tag: str = "",
    include_mu: bool = True,
) -> dict:
    """Fit all 3 cylinders independently.

    Front: no bias.
    Top/Bottom: with bias.
    """
    ThetaF, ThetaO, Y, W, names, re_labels = build_cross_re_v3(train_re, include_mu)

    cylinders = [
        {"name": "front", "theta": ThetaF, "label": "front (no bias)"},
        {"name": "bottom", "theta": ThetaO, "label": "bottom (w/ bias)"},
        {"name": "top", "theta": ThetaO, "label": "top (w/ bias)"},
    ]

    channels = []
    for ci, cyl in enumerate(cylinders):
        print(f"\n  --- {tag} {cyl['label']} ---")
        rows, best = fit_channel_weighted(cyl["theta"], Y[:, ci], W, THRESHOLDS)
        coef = best["coef"]
        print_control_law(names, coef, channel_label=f"{cyl['name']}")
        print(f"         R2={best['r2']:.6f}  MAE={best['mae']:.6f}")
        channels.append({
            "cylinder": cyl["name"],
            "has_bias": cyl["name"] != "front",
            "n_features": cyl["theta"].shape[1],
            "best": {k: float(v) if isinstance(v, (np.floating, float)) else v
                     for k, v in best.items() if k != "coef"},
            "best_coef": [float(c) for c in coef],
            "grid": [{k: float(v) for k, v in row.items() if k != "coef"}
                     for row in rows],
            "feature_names": names,
        })

    # Per-Re breakdown
    print(f"\n  --- {tag} per-Re breakdown ---")
    breakdown = {}
    for re_code in set(re_labels.tolist()):
        mask = re_labels == re_code
        Yr, Wr = Y[mask], W[mask]
        ch_b = []
        for ci, cyl in enumerate(cylinders):
            th = cyl["theta"][mask]
            coef = np.array(channels[ci]["best_coef"], dtype=np.float64)
            y_pred = th @ coef
            y_t = Yr[:, ci]
            y_mean = np.average(y_t, weights=Wr)
            ssr = np.sum(Wr * (y_t - y_pred) ** 2)
            sst = np.sum(Wr * (y_t - y_mean) ** 2) + 1e-12
            r2 = 1.0 - ssr / sst
            mae = float(np.average(np.abs(y_t - y_pred), weights=Wr))
            ch_b.append({"cylinder": cyl["name"], "r2": float(r2), "mae": mae})
        breakdown[f"re{int(re_code)}"] = ch_b
        r2s = ", ".join([f"{b['cylinder']}={b['r2']:.4f}" for b in ch_b])
        print(f"    Re{int(re_code)}: {r2s}")

    return {
        "tag": tag,
        "train_re": train_re,
        "n_samples": int(ThetaF.shape[0]),
        "n_features_front": int(ThetaF.shape[1]),
        "n_features_other": int(ThetaO.shape[1]),
        "channels": channels,
        "per_re_breakdown": breakdown,
    }


def main():
    ap = argparse.ArgumentParser(description="Phase 2 v3: dimensionless + constrained fitting")
    ap.add_argument("--cross-re", action="store_true")
    ap.add_argument("--leave-one-out", action="store_true")
    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")
    ap.add_argument("--no-mu", action="store_true")
    args = ap.parse_args()

    if not (args.cross_re or args.leave_one_out):
        print("ERROR: specify --cross-re and/or --leave-one-out")
        return 1

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

    results = {
        "metadata": {
            "method": "v3_dimensionless_front_nobias_weighted",
            "thresholds": THRESHOLDS,
            "include_mu": include_mu,
        }
    }

    if args.cross_re:
        print("\n" + "=" * 60)
        print("v3 Cross-Re unified (dimensionless + front no-bias + weighted)")
        print("=" * 60)
        results["cross_re"] = run_cross_re_fit(
            train_re, tag="v3-cross", include_mu=include_mu)

    if args.leave_one_out:
        print("\n" + "=" * 60)
        print("v3 Leave-one-out cross-validation")
        print("=" * 60)
        loo_results = {}
        for held_out in train_re:
            train_set = [r for r in train_re if r != held_out]
            print(f"\n--- LOO: train={train_set}, test={held_out} ---")
            loo = run_cross_re_fit(
                train_set, tag=f"v3-loo-{held_out}", include_mu=include_mu)
            loo_results[f"holdout_{held_out}"] = loo
        results["leave_one_out"] = loo_results

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


if __name__ == "__main__":
    main()