import json
import os
from dataclasses import dataclass
from typing import Dict, List, Tuple

import matplotlib.pyplot as plt
import numpy as np
from sklearn.decomposition import PCA
from sklearn.preprocessing import StandardScaler

ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), os.pardir))
OUT_DIR = os.path.join(ROOT, "output", "report_dante_v2_v5_v6")
os.makedirs(OUT_DIR, exist_ok=True)

SEEDS = [11, 29, 47]

V6_NAME = "d1a3o12_250421_forces02_dante_v6_2"
V7_NAME = "d1a3o12_250421_forces02_dante_v7"


@dataclass
class RunData:
    name: str
    x: np.ndarray
    y: np.ndarray
    num_initial: int
    dims: int
    cfg: Dict


def load_run(name: str) -> RunData:
    db_path = os.path.join(ROOT, "output", f"{name}_database_live.npz")
    cfg_path = os.path.join(ROOT, "output", f"{name}_structure_decision.json")

    z = np.load(db_path, allow_pickle=True)
    x = np.asarray(z["input_x"], dtype=np.float64)
    y = np.asarray(z["input_y"], dtype=np.float64)
    cfg = json.load(open(cfg_path, "r", encoding="utf-8"))

    return RunData(
        name=name,
        x=x,
        y=y,
        num_initial=int(cfg["num_initial"]),
        dims=int(cfg["controller_dims"]),
        cfg=cfg,
    )


def feature_dict(obs_t: np.ndarray, obs_prev: np.ndarray, act_prev: np.ndarray) -> Dict[str, float]:
    o0, o1 = float(obs_t[0]), float(obs_t[1])
    p0, p1 = float(obs_prev[0]), float(obs_prev[1])
    a0, a1, a2 = float(act_prev[0]), float(act_prev[1]), float(act_prev[2])
    return {
        "obs0": o0,
        "obs1": o1,
        "dobs0": o0 - p0,
        "dobs1": o1 - p1,
        "sin_obs0": float(np.sin(np.pi * o0)),
        "sin_obs1": float(np.sin(np.pi * o1)),
        "cos_obs0": float(np.cos(np.pi * o0)),
        "cos_obs1": float(np.cos(np.pi * o1)),
        "tanh_obs0": float(np.tanh(o0)),
        "tanh_obs1": float(np.tanh(o1)),
        "act0_l1": a0,
        "act1_l1": a1,
        "act2_l1": a2,
    }


def inv_tanh_map(q: float) -> float:
    qq = float(np.clip(q, -0.999, 0.999))
    x = np.arctanh(qq) / 1.25
    return float(np.clip(x, -1.0, 1.0))


def ppo_point_for_v6(seed: int, basis_terms: List[str]) -> np.ndarray:
    p = os.path.join(ROOT, "output", "report_dante_v2_v5_v6", f"raw_oldenv_seed_{seed}.npz")
    z = np.load(p)
    obs = np.asarray(z["ppo_obs"], dtype=np.float64)
    act = np.asarray(z["ppo_actions"], dtype=np.float64)

    rows_x = []
    rows_y = []
    for t in range(1, len(obs)):
        fd = feature_dict(obs[t], obs[t - 1], act[t - 1])
        row = [1.0] + [float(fd[k]) for k in basis_terms]
        rows_x.append(row)
        rows_y.append(act[t])

    X = np.asarray(rows_x, dtype=np.float64)
    Y = np.asarray(rows_y, dtype=np.float64)

    params = []
    for ch in range(3):
        coef, *_ = np.linalg.lstsq(X, Y[:, ch], rcond=None)
        bias = float(coef[0])
        params.append(inv_tanh_map(bias / 1.0))
        for c in coef[1:]:
            params.append(inv_tanh_map(float(c) / 2.0))

    return np.asarray(params, dtype=np.float64)


def phase_weights(phase: float, k: int) -> np.ndarray:
    z = float(np.mod(phase, 1.0)) * k
    i0 = int(np.floor(z)) % k
    frac = float(z - np.floor(z))
    i1 = (i0 + 1) % k
    w = np.zeros(k, dtype=np.float64)
    w[i0] += (1.0 - frac)
    w[i1] += frac
    return w


def ppo_point_for_v7(seed: int, k: int, period_steps: float) -> np.ndarray:
    p = os.path.join(ROOT, "output", "report_dante_v2_v5_v6", f"raw_oldenv_seed_{seed}.npz")
    z = np.load(p)
    act = np.asarray(z["ppo_actions"], dtype=np.float64)

    n = int(act.shape[0])
    phase = 0.0
    step_phase = 1.0 / max(1e-6, float(period_steps))

    W = np.zeros((n, k), dtype=np.float64)
    for t in range(n):
        W[t] = phase_weights(phase, k)
        phase = (phase + step_phase) % 1.0

    params = []
    for ch in range(3):
        coef, *_ = np.linalg.lstsq(W, act[:, ch], rcond=None)
        coef = np.clip(coef, -1.0, 1.0)
        params.extend(coef.tolist())

    return np.asarray(params, dtype=np.float64)


def fit_pca_and_project(x: np.ndarray, extra: np.ndarray) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
    scaler = StandardScaler()
    xs = scaler.fit_transform(x)
    extra_s = scaler.transform(extra)

    pca = PCA(n_components=2, random_state=0)
    x2 = pca.fit_transform(xs)
    extra2 = pca.transform(extra_s)
    return x2, extra2, pca.explained_variance_ratio_


def distance_metrics(x: np.ndarray, y: np.ndarray, ppo_points: np.ndarray, n0: int) -> Dict:
    centroid = np.mean(ppo_points, axis=0)
    d = np.linalg.norm(x - centroid.reshape(1, -1), axis=1)

    n = len(d)
    xx = np.arange(n, dtype=np.float64)
    slope = float(np.polyfit(xx, d, deg=1)[0]) if n >= 2 else 0.0

    init_mean = float(np.mean(d[:n0])) if n0 > 0 else float(np.mean(d))
    late_mean = float(np.mean(d[n0:])) if n > n0 else init_mean

    thr = np.quantile(y, 0.9)
    idx_hi = np.where(y >= thr)[0]
    hi_mean = float(np.mean(d[idx_hi])) if len(idx_hi) > 0 else float("nan")

    return {
        "init_mean": init_mean,
        "late_mean": late_mean,
        "late_minus_init": float(late_mean - init_mean),
        "slope": slope,
        "high_reward_dist_mean": hi_mean,
        "overall_dist_mean": float(np.mean(d)),
        "distance": d.tolist(),
        "ppo_centroid": centroid.tolist(),
    }


def plot_pca(
    run_name: str,
    x2: np.ndarray,
    y: np.ndarray,
    ppo2: np.ndarray,
    evr: np.ndarray,
    out_png: str,
) -> None:
    plt.figure(figsize=(8, 6))
    sc = plt.scatter(x2[:, 0], x2[:, 1], c=y, cmap="viridis", s=12, alpha=0.75)
    plt.colorbar(sc, label="reward_scaled")

    colors = ["#e41a1c", "#377eb8", "#ff7f00"]
    for i, seed in enumerate(SEEDS):
        plt.scatter(
            [ppo2[i, 0]],
            [ppo2[i, 1]],
            s=120,
            c=colors[i],
            marker="*",
            edgecolors="k",
            linewidths=0.9,
            label=f"PPO seed {seed}",
            zorder=6,
        )

    plt.title(
        f"{run_name} PCA with PPO points\\n"
        f"PC1 {evr[0]*100:.1f}% | PC2 {evr[1]*100:.1f}%"
    )
    plt.xlabel("PC1")
    plt.ylabel("PC2")
    plt.legend(loc="best", fontsize=9)
    plt.grid(alpha=0.25)
    plt.tight_layout()
    plt.savefig(out_png, dpi=170)
    plt.close()


def plot_distance_curve(run_name: str, d: np.ndarray, n0: int, out_png: str) -> None:
    n = len(d)
    x = np.arange(1, n + 1)
    plt.figure(figsize=(9, 4.8))
    plt.plot(x, d, lw=1.4, color="#1f77b4")
    plt.axvline(n0, color="k", ls="--", lw=1.0, label=f"init end @ {n0}")
    plt.title(f"{run_name}: distance to PPO centroid")
    plt.xlabel("sample order")
    plt.ylabel("L2 distance")
    plt.grid(alpha=0.25)
    plt.legend(loc="best")
    plt.tight_layout()
    plt.savefig(out_png, dpi=170)
    plt.close()


def diagnose(run_name: str, m: Dict) -> Dict:
    outward = bool(m["late_minus_init"] > 0.0 and m["slope"] > 0.0)
    high_reward_near = bool(m["high_reward_dist_mean"] < m["overall_dist_mean"])

    return {
        "run": run_name,
        "outward_sampling_still_exists": outward,
        "high_reward_closer_to_ppo_than_overall": high_reward_near,
        "metrics": {
            "late_minus_init": float(m["late_minus_init"]),
            "slope": float(m["slope"]),
            "high_reward_dist_mean": float(m["high_reward_dist_mean"]),
            "overall_dist_mean": float(m["overall_dist_mean"]),
        },
    }


def main() -> None:
    v6 = load_run(V6_NAME)
    v7 = load_run(V7_NAME)

    basis_terms = list(v6.cfg["basis_terms"])
    v6_ppo = np.vstack([ppo_point_for_v6(s, basis_terms) for s in SEEDS])

    k = int(v7.cfg["control_points_per_channel"])
    period_steps = float(v7.cfg["period_steps"])
    v7_ppo = np.vstack([ppo_point_for_v7(s, k=k, period_steps=period_steps) for s in SEEDS])

    v6_x2, v6_ppo2, v6_evr = fit_pca_and_project(v6.x, v6_ppo)
    v7_x2, v7_ppo2, v7_evr = fit_pca_and_project(v7.x, v7_ppo)

    v6_m = distance_metrics(v6.x, v6.y, v6_ppo, v6.num_initial)
    v7_m = distance_metrics(v7.x, v7.y, v7_ppo, v7.num_initial)

    p_v6_pca = os.path.join(OUT_DIR, "v6_2_pca_with_ppo_points.png")
    p_v7_pca = os.path.join(OUT_DIR, "v7_pca_with_ppo_points.png")
    p_v6_dist = os.path.join(OUT_DIR, "v6_2_distance_order_vs_ppo_centroid.png")
    p_v7_dist = os.path.join(OUT_DIR, "v7_distance_order_vs_ppo_centroid.png")

    plot_pca(v6.name, v6_x2, v6.y, v6_ppo2, v6_evr, p_v6_pca)
    plot_pca(v7.name, v7_x2, v7.y, v7_ppo2, v7_evr, p_v7_pca)
    plot_distance_curve(v6.name, np.asarray(v6_m["distance"]), v6.num_initial, p_v6_dist)
    plot_distance_curve(v7.name, np.asarray(v7_m["distance"]), v7.num_initial, p_v7_dist)

    summary = {
        "v6": {
            "name": v6.name,
            "dims": int(v6.dims),
            "num_samples": int(len(v6.y)),
            "num_initial": int(v6.num_initial),
            "best_reward": float(np.max(v6.y) / 100.0),
            "distance_metrics": {k: v for k, v in v6_m.items() if k != "distance"},
            "diagnosis": diagnose(v6.name, v6_m),
        },
        "v7": {
            "name": v7.name,
            "dims": int(v7.dims),
            "num_samples": int(len(v7.y)),
            "num_initial": int(v7.num_initial),
            "best_reward": float(np.max(v7.y) / 100.0),
            "period_steps": period_steps,
            "distance_metrics": {k: v for k, v in v7_m.items() if k != "distance"},
            "diagnosis": diagnose(v7.name, v7_m),
        },
        "figures": {
            "v6_pca": p_v6_pca,
            "v6_distance": p_v6_dist,
            "v7_pca": p_v7_pca,
            "v7_distance": p_v7_dist,
        },
    }

    out_json = os.path.join(OUT_DIR, "v6_v7_pca_distance_with_ppo_summary.json")
    with open(out_json, "w", encoding="utf-8") as f:
        json.dump(summary, f, indent=2, ensure_ascii=False)

    print("saved", out_json)
    print(json.dumps({
        "v6_late_minus_init": summary["v6"]["distance_metrics"]["late_minus_init"],
        "v6_slope": summary["v6"]["distance_metrics"]["slope"],
        "v7_late_minus_init": summary["v7"]["distance_metrics"]["late_minus_init"],
        "v7_slope": summary["v7"]["distance_metrics"]["slope"],
    }, indent=2, ensure_ascii=False))


if __name__ == "__main__":
    main()
