import json
import os
import pickle
from typing import Dict, List, Tuple

import numpy as np
import pysindy as ps


CURRENT_DIR = os.path.dirname(os.path.abspath(__file__))
ROOT = os.path.abspath(os.path.join(CURRENT_DIR, os.pardir))

DATA_PATH = os.path.join(ROOT, "output", "d1a3o12_250421_forces02_sindy_dataset.pkl")
OUT_DIR = os.path.join(ROOT, "output", "report_dante_v3_v4")
OUT_JSON = os.path.join(OUT_DIR, "ppo_sindy_control_fit.json")

WARMUP_STEPS = 0
MAX_LAG = 1
THRESHOLDS = [0.0, 0.005, 0.01, 0.015, 0.02, 0.03, 0.05]
NZ_RATIO_THRESHOLD = 0.70


def episode_metric(rewards: np.ndarray, warmup: int = 150) -> float:
    r = np.asarray(rewards, dtype=np.float64).reshape(-1)
    if r.size == 0:
        return 0.0
    eff = r[warmup:] if r.size > warmup else r[-1:]
    return float(np.mean(eff[-100:])) if eff.size >= 100 else float(np.mean(eff))


def load_episodes(path: str) -> Tuple[List[Dict], Dict]:
    with open(path, "rb") as f:
        data = pickle.load(f)

    if isinstance(data, dict) and "episodes" in data:
        return list(data["episodes"]), dict(data.get("meta", {}))
    if isinstance(data, list):
        return data, {}
    raise RuntimeError(f"Unsupported dataset format: {type(data)}")


def extract_episode_arrays(ep: Dict) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
    actions = np.asarray(ep.get("actions", []), dtype=np.float64)
    observations = np.asarray(ep.get("observations", []), dtype=np.float64)
    rewards = np.asarray(ep.get("rewards", []), dtype=np.float64)

    if actions.ndim != 2:
        actions = actions.reshape(actions.shape[0], -1)
    actions = actions[:, :3]

    if observations.ndim != 2:
        observations = observations.reshape(observations.shape[0], -1)
    observations = observations[:, :2]

    n = min(actions.shape[0], observations.shape[0], rewards.shape[0])
    return actions[:n], observations[:n], rewards[:n]


def build_dataset(episodes: List[Dict], warmup: int = 150) -> Tuple[np.ndarray, np.ndarray, List[str], Dict]:
    x_rows = []
    y_rows = []

    feat_names = [
        "bias1",
        "obs0",
        "obs1",
        "dobs0",
        "dobs1",
        "sin_obs0",
        "sin_obs1",
        "cos_obs0",
        "cos_obs1",
        "tanh_obs0",
        "tanh_obs1",
        "act0_l1",
        "act1_l1",
        "act2_l1",
    ]

    stats = {"episodes_used": 0, "samples_used": 0}

    for ep in episodes:
        actions, obs2, rewards = extract_episode_arrays(ep)
        t_len = min(actions.shape[0], obs2.shape[0], rewards.shape[0])
        if t_len <= (MAX_LAG + warmup + 1):
            continue

        for t in range(MAX_LAG, t_len):
            if t < warmup:
                continue

            o = obs2[t]
            o1 = obs2[t - 1]
            a_prev = actions[t - 1]

            x_rows.append(
                [
                    1.0,
                    o[0],
                    o[1],
                    o[0] - o1[0],
                    o[1] - o1[1],
                    np.sin(np.pi * o[0]),
                    np.sin(np.pi * o[1]),
                    np.cos(np.pi * o[0]),
                    np.cos(np.pi * o[1]),
                    np.tanh(o[0]),
                    np.tanh(o[1]),
                    a_prev[0],
                    a_prev[1],
                    a_prev[2],
                ]
            )
            y_rows.append(actions[t])

        stats["episodes_used"] += 1

    if len(x_rows) < 512:
        raise RuntimeError(f"Too few samples for fitting: {len(x_rows)}")

    x = np.asarray(x_rows, dtype=np.float64)
    y = np.asarray(y_rows, dtype=np.float64)
    stats["samples_used"] = int(x.shape[0])
    return x, y, feat_names, stats


def r2(y: np.ndarray, yp: np.ndarray) -> float:
    ssr = float(np.sum((y - yp) ** 2))
    sst = float(np.sum((y - np.mean(y)) ** 2) + 1e-12)
    return float(1.0 - ssr / sst)


def fit_channel_grid(x: np.ndarray, y: np.ndarray, thresholds: List[float]):
    std = np.std(x, axis=0)
    std = np.where(std < 1e-8, 1.0, std)
    xs = x / std

    rows = []
    for th in thresholds:
        opt = ps.STLSQ(threshold=th, alpha=1e-4, max_iter=25)
        opt.fit(xs, y)
        coef = np.asarray(opt.coef_, dtype=np.float64).reshape(-1) / std
        yp = x @ coef
        rows.append(
            {
                "threshold": float(th),
                "nz": int(np.sum(np.abs(coef) > 1e-8)),
                "r2": r2(y, yp),
                "mae": float(np.mean(np.abs(y - yp))),
                "coef": coef,
            }
        )

    best = max(rows, key=lambda z: z["r2"])
    return rows, best


def top_terms(feat_names: List[str], coef: np.ndarray, topk: int = 10) -> List[Dict]:
    idx = np.argsort(np.abs(coef))[::-1]
    out = []
    for i in idx:
        c = float(coef[i])
        if abs(c) < 1e-8:
            continue
        out.append({"term": feat_names[i], "coef": c, "abs_coef": float(abs(c))})
        if len(out) >= topk:
            break
    return out


def main() -> None:
    os.makedirs(OUT_DIR, exist_ok=True)
    episodes_all, meta = load_episodes(DATA_PATH)

    ppo_eps = [ep for ep in episodes_all if str(ep.get("source", "unknown")) == "ppo_eval"]
    if len(ppo_eps) < 10:
        # Fallback only when ppo episodes are too few.
        ppo_eps = episodes_all

    metrics = []
    for i, ep in enumerate(ppo_eps):
        _, _, rewards = extract_episode_arrays(ep)
        metrics.append((i, episode_metric(rewards, warmup=WARMUP_STEPS)))
    metrics_sorted = sorted(metrics, key=lambda x: x[1], reverse=True)

    x, y, feat_names, data_stats = build_dataset(ppo_eps, warmup=WARMUP_STEPS)

    channel_models = []
    votes = np.zeros(len(feat_names), dtype=np.int64)
    coef_abs_sum = np.zeros(len(feat_names), dtype=np.float64)
    nz_ratios = []

    for ch in range(3):
        rows, best = fit_channel_grid(x, y[:, ch], THRESHOLDS)
        coef = best["coef"]
        active = np.abs(coef) >= 0.01
        votes += active.astype(np.int64)
        coef_abs_sum += np.abs(coef)
        nz_nonbias = int(np.sum(np.abs(coef[1:]) > 1e-8))
        denom_nonbias = int(max(1, len(feat_names) - 1))
        nz_ratio = float(nz_nonbias / denom_nonbias)
        nz_ratios.append(nz_ratio)

        channel_models.append(
            {
                "channel": ch,
                "r2": float(best["r2"]),
                "mae": float(best["mae"]),
                "best_sparse": {
                    "threshold": float(best["threshold"]),
                    "nz": int(best["nz"]),
                    "nz_nonbias": nz_nonbias,
                    "nz_ratio": nz_ratio,
                },
                "top_terms": top_terms(feat_names, coef, topk=10),
                "grid": [
                    {
                        "threshold": float(z["threshold"]),
                        "nz": int(z["nz"]),
                        "r2": float(z["r2"]),
                        "mae": float(z["mae"]),
                    }
                    for z in rows
                ],
            }
        )

    score_idx = sorted(
        range(len(feat_names)),
        key=lambda k: (int(votes[k]), float(coef_abs_sum[k])),
        reverse=True,
    )
    global_top = [feat_names[k] for k in score_idx if feat_names[k] != "bias1"]
    over_complex_channels = [int(cm["channel"]) for cm in channel_models if cm["best_sparse"]["nz_ratio"] >= NZ_RATIO_THRESHOLD]
    use_sindy_prior = len(over_complex_channels) == 0

    if use_sindy_prior:
        prior_reason = (
            f"all channels nz_ratio < {NZ_RATIO_THRESHOLD:.2f}; sparse structure is sufficiently compressible"
        )
    else:
        prior_reason = (
            f"channels {over_complex_channels} have nz_ratio >= {NZ_RATIO_THRESHOLD:.2f}; sparse structure is too dense"
        )

    result = {
        "data_path": DATA_PATH,
        "dataset_meta": meta,
        "warmup_steps": WARMUP_STEPS,
        "max_lag": MAX_LAG,
        "nz_ratio_threshold": NZ_RATIO_THRESHOLD,
        "episodes_total": len(episodes_all),
        "episodes_used_source": "ppo_eval" if len([ep for ep in episodes_all if str(ep.get("source", "")) == "ppo_eval"]) >= 10 else "all",
        "episodes_used_count": len(ppo_eps),
        "episode_metrics_top10": metrics_sorted[:10],
        "fit_data_stats": data_stats,
        "threshold_grid": THRESHOLDS,
        "feature_names": feat_names,
        "channel_models": channel_models,
        "global_term_votes": {feat_names[k]: int(votes[k]) for k in range(len(feat_names))},
        "global_term_abs_coef_sum": {feat_names[k]: float(coef_abs_sum[k]) for k in range(len(feat_names))},
        "global_top_terms": global_top[:12],
        "complexity_decision": {
            "nz_ratio_threshold": float(NZ_RATIO_THRESHOLD),
            "channel_nz_ratio": {f"ch{i}": float(z) for i, z in enumerate(nz_ratios)},
            "over_complex_channels": over_complex_channels,
            "use_sindy_prior": bool(use_sindy_prior),
            "reason": prior_reason,
        },
    }

    with open(OUT_JSON, "w", encoding="utf-8") as f:
        json.dump(result, f, indent=2, ensure_ascii=False)

    print(f"saved: {OUT_JSON}")
    print("episodes used:", len(ppo_eps), "/", len(episodes_all))
    print("samples:", data_stats["samples_used"])
    print("global top terms:", result["global_top_terms"][:8])
    for cm in channel_models:
        print(
            f"ch{cm['channel']} r2={cm['r2']:.4f} mae={cm['mae']:.5f} "
            f"nz={cm['best_sparse']['nz']} nz_ratio={cm['best_sparse']['nz_ratio']:.3f}"
        )
    print("use_sindy_prior:", result["complexity_decision"]["use_sindy_prior"])
    print("reason:", result["complexity_decision"]["reason"])


if __name__ == "__main__":
    main()
