#!/usr/bin/env python3
"""Offline multi-step rollout evaluation for SINDy models.

Evaluates a SINDy model by doing recursive multi-step prediction on
offline PPO data, WITHOUT running CFD. This isolates the model's
rollout behavior from CFD initialization stochasticity.

Usage:
    python3 src/SR_analysis/validate/eval_rollout.py \\
        --sindy-results src/SR_analysis/sindy/karman/sindy_results_deriv.json \\
        --scene karman_re100

    # Compare multiple feature presets
    python3 src/SR_analysis/validate/eval_rollout.py \\
        --scene karman_re100 --preset static --preset phase --preset full_lag
"""
from __future__ import annotations

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

import numpy as np

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

from SR_analysis.utils.feature_builder import (
    compute_dimensionless, compute_features, build_feature_matrix,
    PHYSICS_FEAT_KEYS, LAG_FEAT_KEYS, DERIV_FEAT_KEYS,
    AUG_LEVEL_1_KEYS, PHASE_STATE_KEYS, ILLUSION_ERR_KEYS, ILLUSION_PHASE_KEYS,
)
from SR_analysis.utils.sindy_fitter import compute_action_deriv
from SR_analysis.configs import get_scene


def load_model_data(scene_name: str):
    """Load controlled data and scene config."""
    cfg = get_scene(scene_name)
    scene_id = cfg["scene_id"]
    data_dir = os.path.join(
        os.path.dirname(__file__), "..", "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)
    scale = cfg["action_scale"]
    bias = np.array(cfg["action_bias"], dtype=np.float64)
    u0 = cfg["u0"]
    actions_phys = (actions_norm * scale + bias) * 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 build_sindy_features(sensors, forces, actions_phys, mu, u0, dt_c,
                          target_forces, feat_keys):
    """Build Theta_f, Theta_r, Y from data using given feat_keys.
    
    Returns (Theta_f, Theta_r, Y_deriv, fn_f, fn_r, a_prev_phys, a_prev2_phys).
    """
    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,
        alpha_mode=False, include_mu=False,
        include_cos_sin=False, dt_c=dt_c, u0=u0,
        target_forces=target_forces,
        sensors_raw=sensors, forces_raw=forces,
    )
    
    Theta_f = build_feature_matrix(sym, feat_keys, add_bias=False)
    Theta_r = build_feature_matrix(sym, feat_keys, add_bias=True)
    Y_deriv = compute_action_deriv(actions_phys, dt_c, u0=u0)
    
    return (Theta_f, Theta_r, Y_deriv,
            feat_keys, ["bias"] + feat_keys)


def compute_rollout(
    sensors, forces, actions_phys, target_forces, cfg,
    front_coef, top_coef, bottom_coef,
    feat_keys_front, feat_keys_rear,
    max_steps: int = 50,
    center_diff: bool = False,
) -> dict:
    """Compute offline multi-step rollout error."""
    mu = cfg["mu"]
    u0 = cfg["u0"]
    dt_c = cfg.get("sample_interval", 800) / 2000.0
    T = sensors.shape[0]
    
    # True PPO alpha(t)
    alpha_true = actions_phys / u0
    
    # SINDy predicted alpha(t) via recursive rollout
    alpha_pred = np.zeros_like(alpha_true)
    
    # Step 0, 1: initialize from PPO data (first 2 steps use true values)
    alpha_pred[0] = alpha_true[0]
    alpha_pred[1] = alpha_true[1]
    
    # Build feature history for augmented features
    fbs = {}  # feature build state: sensors_raw, forces_raw
    
    for t in range(2, min(T, max_steps + 2)):
        # Build current obs: need sensor/force from the actual data
        obs_slice = np.hstack([sensors[t], forces[t]])
        
        # For augmented features, also need obs[t-1]
        obs_prev = np.hstack([sensors[t-1], forces[t-1]])
        
        # Check if augmented features needed
        needs_aug = any(k.endswith("_dt") or k.endswith("_lag1") for k in feat_keys_front)
        
        if needs_aug:
            sensors_raw = np.vstack([sensors[t-1:t+1]])
            forces_raw = np.vstack([forces[t-1:t+1]])
            sensors_curr = sensors[t:t+1]
            forces_curr = forces[t:t+1]
        else:
            sensors_raw = sensors[t:t+1]
            forces_raw = forces[t:t+1]
            sensors_curr = sensors[t:t+1]
            forces_curr = forces[t:t+1]
        
        # Build features for current step
        a_prev = alpha_pred[t-1:t] * u0  # use PREDICTED a(t-1)
        a_prev2 = alpha_pred[t-2:t-1] * u0  # use PREDICTED a(t-2)
        
        dim = compute_dimensionless(sensors_curr, forces_curr, u0=u0, d=20.0)
        tf = target_forces[t:t+1] if target_forces is not None else None
        sym = compute_features(dim, a_prev, a_prev2, mu,
                               alpha_mode=False, include_mu=False,
                               include_cos_sin=False, dt_c=dt_c, u0=u0,
                               target_forces=tf,
                               sensors_raw=sensors_raw, forces_raw=forces_raw)
        
        # Front: no bias
        fv_f = build_feature_matrix(sym, feat_keys_front, add_bias=False)[0]
        d_alphaF = float(np.dot(fv_f, front_coef))
        
        # Top: with bias
        fv_r = build_feature_matrix(sym, feat_keys_rear, add_bias=True)[0]
        d_alphaT = float(np.dot(fv_r, top_coef))
        
        # Bottom: from fit (independent since we're rolling out)
        d_alphaB = float(np.dot(fv_r, bottom_coef))
        
        # Integrate
        alpha_pred[t] = alpha_pred[t-1] + dt_c * np.array([d_alphaF, d_alphaB, d_alphaT])
    
    # Compute errors at various rollout horizons
    n = min(T, max_steps + 2)
    # Use full-range alpha range for all horizons (stable denominator)
    full_alpha_range = alpha_true[2:n].max(axis=0) - alpha_true[2:n].min(axis=0)
    results = {}
    for horizon in [1, 5, 20, 50]:
        if horizon + 2 > n:
            horizon = n - 2
        if horizon <= 0:
            continue
        err = alpha_pred[2:2+horizon] - alpha_true[2:2+horizon]
        rmse_per_channel = np.sqrt(np.mean(err**2, axis=0))
        rel_err = rmse_per_channel / (full_alpha_range + 1e-12)
        results[f"{horizon}step_rmse"] = rmse_per_channel.tolist()
        results[f"{horizon}step_rel"] = rel_err.tolist()
    
    results["alpha_range"] = full_alpha_range.tolist()
    results["n_samples"] = n - 2
    
    return results


def main():
    ap = argparse.ArgumentParser(description="Offline SINDy rollout evaluation")
    ap.add_argument("--sindy-results", type=str, required=True,
                    help="Path to sindy_results JSON")
    ap.add_argument("--scene", type=str, required=True,
                    help="Scene name (e.g. karman_re100)")
    ap.add_argument("--threshold", type=float, default=None,
                    help="SINDy threshold (default: best R2)")
    args = ap.parse_args()
    
    # Load SINDy model
    with open(args.sindy_results) as f:
        data = json.load(f)
    
    per = data.get("per_scene", data).get(args.scene)
    if per is None:
        # Try top-level (joint model format)
        per = data
    
    fn_f = per.get("feature_names_front", [])
    fn_r = per.get("feature_names_rear", [])
    
    def _coef(ch_name):
        ch = per.get(ch_name, {})
        if args.threshold is not None:
            for r in ch.get("results", []):
                if abs(r["threshold"] - args.threshold) < 1e-6:
                    return np.array(r.get("coef", ch["best_coef"]), dtype=np.float64)
        return np.array(ch["best_coef"], dtype=np.float64)
    
    front_coef = _coef("front")[:len(fn_f)]
    top_coef = _coef("top")[:len(fn_r)]
    bottom_coef = _coef("bottom")[:len(fn_r)]
    feat_keys_front = [k for k in fn_f if k != "bias"]
    feat_keys_rear = [k for k in fn_r if k != "bias"]
    
    sensors, forces, actions_phys, target_forces, cfg = load_model_data(args.scene)
    
    print(f"Rollout eval: {args.scene}")
    print(f"  Features: front {len(feat_keys_front)}, rear {len(feat_keys_rear)}")
    print(f"  Front keys: {feat_keys_front}")
    print()
    
    results = compute_rollout(
        sensors, forces, actions_phys, target_forces, cfg,
        front_coef, top_coef, bottom_coef,
        feat_keys_front, feat_keys_rear,
    )
    
    print(f"  alpha_range = [{results['alpha_range'][0]:.3f}, {results['alpha_range'][1]:.3f}, {results['alpha_range'][2]:.3f}]")
    print()
    print(f"  {'Horizon':>8s}  {'Front RMSE':>12s}  {'Bottom RMSE':>12s}  {'Top RMSE':>12s}  {'Front rel':>10s}  {'Bottom rel':>10s}  {'Top rel':>10s}")
    print(f"  {'-'*8}  {'-'*12}  {'-'*12}  {'-'*12}  {'-'*10}  {'-'*10}  {'-'*10}")
    for h in [1, 5, 20, 50]:
        rmse_key = f"{h}step_rmse"
        rel_key = f"{h}step_rel"
        if rmse_key in results:
            rmse = results[rmse_key]
            rel = results[rel_key]
            print(f"  {h:>8d}  {rmse[0]:>12.4f}  {rmse[1]:>12.4f}  {rmse[2]:>12.4f}  {rel[0]:>10.3f}  {rel[1]:>10.3f}  {rel[2]:>10.3f}")
    
    print()
    print("Interpretation:")
    print("  1-step = pure one-step prediction quality")
    print("  5-step = short rollout (control propagation ~1/40 of domain)")
    print("  20-step = medium rollout (~1/10 of domain)")
    print("  50-step = long rollout (~1/4 of domain)")
    print(f"  Relative error = RMSE / alpha_range (lower is better)")
    print(f"  If 5-step error >> 1-step error: model has rollout instability")


if __name__ == "__main__":
    main()
