import json
import os
import pickle
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

from dante_pinball.env.gym_env_dante_total_force import CustomEnv


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)

V6_RUN = "d1a3o12_250421_forces02_dante_v6_3"
V7_RUN = "d1a3o12_250421_forces02_dante_v7_1"
SEEDS = [11, 29, 47]
N_ACT = 3
EVAL_STEPS = 300


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))


@dataclass
class RunData:
    name: str
    x: np.ndarray
    y: np.ndarray
    num_initial: 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)
    with open(cfg_path, "r", encoding="utf-8") as f:
        cfg = json.load(f)

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


def ppo_npz(seed: int) -> str:
    return os.path.join(OUT_DIR, f"raw_oldenv_seed_{seed}.npz")


def ppo_point_for_v6(seed: int, basis_terms: List[str]) -> np.ndarray:
    z = np.load(ppo_npz(seed))
    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])
        rows_x.append([1.0] + [float(fd[k]) for k in basis_terms])
        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:
    z = np.load(ppo_npz(seed))
    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 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.78)
    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=130,
            c=colors[i],
            marker="*",
            edgecolors="k",
            linewidths=0.9,
            label=f"PPO fitted seed {seed}",
            zorder=6,
        )

    plt.title(f"{run_name}: PCA with 3 PPO fitted points\\nPC1 {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()


class LinearBasisController:
    def __init__(self, basis_terms: List[str]):
        self.basis_terms = list(basis_terms)
        self.num_basis = len(self.basis_terms)
        self.total_params = N_ACT * (1 + self.num_basis)
        self.params = np.zeros(self.total_params, dtype=np.float64)
        self.obs_l1 = np.zeros(2, dtype=np.float64)
        self.prev_action = np.zeros(3, dtype=np.float64)

    def reset_state(self, obs0: np.ndarray) -> None:
        obs0 = np.asarray(obs0, dtype=np.float64).reshape(-1)
        if obs0.size < 2:
            obs0 = np.zeros(2, dtype=np.float64)
        self.obs_l1 = obs0[:2].copy()
        self.prev_action = np.zeros(3, dtype=np.float64)

    def set_params(self, x: np.ndarray) -> None:
        x = np.asarray(x, dtype=np.float64).reshape(-1)
        if x.size != self.total_params:
            raise ValueError(f"controller params mismatch: {x.size} != {self.total_params}")
        self.params = np.clip(x, -1.0, 1.0)

    def _feature_dict(self, obs: np.ndarray) -> Dict[str, float]:
        o = np.asarray(obs, dtype=np.float64).reshape(-1)
        if o.size < 2:
            o = np.zeros(2, dtype=np.float64)

        o0, o1 = float(o[0]), float(o[1])
        o0_l1, o1_l1 = float(self.obs_l1[0]), float(self.obs_l1[1])
        a0, a1, a2 = float(self.prev_action[0]), float(self.prev_action[1]), float(self.prev_action[2])

        return {
            "obs0": o0,
            "obs1": o1,
            "dobs0": o0 - o0_l1,
            "dobs1": o1 - o1_l1,
            "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 predict(self, obs: np.ndarray) -> np.ndarray:
        feat = self._feature_dict(obs)
        out = np.zeros(3, dtype=np.float64)
        stride = 1 + self.num_basis

        for ch in range(3):
            off = ch * stride
            y = np.tanh(1.25 * self.params[off])
            for k, term in enumerate(self.basis_terms):
                y += (2.0 * np.tanh(1.25 * self.params[off + 1 + k])) * feat.get(term, 0.0)
            out[ch] = y

        out = np.clip(out, -1.0, 1.0)
        obs2 = np.asarray(obs, dtype=np.float64).reshape(-1)
        if obs2.size < 2:
            obs2 = np.zeros(2, dtype=np.float64)
        self.obs_l1 = obs2[:2].copy()
        self.prev_action = out.copy()
        return out.astype(np.float32)


class PeriodicOpenLoopController:
    def __init__(self, control_points_per_channel: int, period_steps: float):
        self.k = int(control_points_per_channel)
        self.ctrl_points = np.zeros((N_ACT, self.k), dtype=np.float64)
        self.phase = 0.0
        self.current_period_steps = float(period_steps)
        self.total_params = int(N_ACT * self.k)

    def reset_state(self) -> None:
        self.phase = 0.0

    def _eval_channel(self, points: np.ndarray, phase: float) -> float:
        p = np.asarray(points, dtype=np.float64).reshape(-1)
        z = float(np.mod(phase, 1.0)) * self.k
        i0 = int(np.floor(z)) % self.k
        frac = float(z - np.floor(z))
        i1 = (i0 + 1) % self.k
        return float((1.0 - frac) * p[i0] + frac * p[i1])

    def set_params(self, x: np.ndarray) -> None:
        x = np.asarray(x, dtype=np.float64).reshape(-1)
        if x.size != self.total_params:
            raise ValueError(f"controller params mismatch: {x.size} != {self.total_params}")
        core = np.clip(x, -1.0, 1.0)
        self.ctrl_points = core.reshape(N_ACT, self.k)

    def predict(self, _obs: np.ndarray) -> np.ndarray:
        action = np.zeros(N_ACT, dtype=np.float64)
        for ch in range(N_ACT):
            action[ch] = self._eval_channel(self.ctrl_points[ch], self.phase)
        action = np.clip(action, -1.0, 1.0)
        step_phase = 1.0 / max(1e-6, float(self.current_period_steps))
        self.phase = float((self.phase + step_phase) % 1.0)
        return action.astype(np.float32)


def extract_ux_uy(env: CustomEnv) -> Tuple[np.ndarray, np.ndarray]:
    nx = env.flow_field.FIELD_SHAPE[0]
    ny = env.flow_field.FIELD_SHAPE[1]
    env.flow_field.get_ddf()
    ddf = env.flow_field.ddf.copy().reshape((9, ny, nx)).transpose(2, 1, 0)
    ux = ddf[:, :, 1] + ddf[:, :, 5] + ddf[:, :, 8] - ddf[:, :, 3] - ddf[:, :, 6] - ddf[:, :, 7]
    uy = ddf[:, :, 2] + ddf[:, :, 5] + ddf[:, :, 6] - ddf[:, :, 4] - ddf[:, :, 7] - ddf[:, :, 8]
    return ux.astype(np.float64), uy.astype(np.float64)


def vorticity_from_ux_uy(ux: np.ndarray, uy: np.ndarray) -> np.ndarray:
    dvy_dx = np.gradient(uy, axis=0)
    dvx_dy = np.gradient(ux, axis=1)
    return (dvy_dx - dvx_dy).astype(np.float64)


def rollout_controller_v6(params: np.ndarray, basis_terms: List[str], device_id: int = 1) -> Dict[str, np.ndarray]:
    env = CustomEnv(device_id=int(device_id))
    ctrl = LinearBasisController(basis_terms)
    ctrl.set_params(params)

    obs, _ = env.reset()
    obs = np.asarray(obs, dtype=np.float32)
    ctrl.reset_state(obs)

    obs_hist = []
    act_hist = []
    rew_hist = []

    try:
        for _ in range(EVAL_STEPS):
            act = ctrl.predict(obs)
            next_obs, reward, done, trunc, _ = env.step(act)
            obs_hist.append(np.asarray(obs, dtype=np.float64).copy())
            act_hist.append(np.asarray(act, dtype=np.float64).copy())
            rew_hist.append(float(reward))
            obs = np.asarray(next_obs, dtype=np.float32)
            if done or trunc:
                obs, _ = env.reset()
                obs = np.asarray(obs, dtype=np.float32)
                ctrl.reset_state(obs)

        ux, uy = extract_ux_uy(env)
    finally:
        env.close()

    return {
        "obs": np.asarray(obs_hist, dtype=np.float64),
        "act": np.asarray(act_hist, dtype=np.float64),
        "rew": np.asarray(rew_hist, dtype=np.float64),
        "ux": ux,
        "uy": uy,
    }


def rollout_controller_v7(params: np.ndarray, k: int, period_steps: float, device_id: int = 0) -> Dict[str, np.ndarray]:
    env = CustomEnv(device_id=int(device_id))
    ctrl = PeriodicOpenLoopController(control_points_per_channel=k, period_steps=float(period_steps))
    ctrl.set_params(params)

    obs, _ = env.reset()
    obs = np.asarray(obs, dtype=np.float32)
    ctrl.reset_state()

    obs_hist = []
    act_hist = []
    rew_hist = []

    try:
        for _ in range(EVAL_STEPS):
            act = ctrl.predict(obs)
            next_obs, reward, done, trunc, _ = env.step(act)
            obs_hist.append(np.asarray(obs, dtype=np.float64).copy())
            act_hist.append(np.asarray(act, dtype=np.float64).copy())
            rew_hist.append(float(reward))
            obs = np.asarray(next_obs, dtype=np.float32)
            if done or trunc:
                obs, _ = env.reset()
                obs = np.asarray(obs, dtype=np.float32)
                ctrl.reset_state()

        ux, uy = extract_ux_uy(env)
    finally:
        env.close()

    return {
        "obs": np.asarray(obs_hist, dtype=np.float64),
        "act": np.asarray(act_hist, dtype=np.float64),
        "rew": np.asarray(rew_hist, dtype=np.float64),
        "ux": ux,
        "uy": uy,
    }


def best_ppo_seed() -> Tuple[int, float]:
    best_seed = None
    best_tail = -1e18
    for s in SEEDS:
        z = np.load(ppo_npz(s))
        r = np.asarray(z["ppo_rewards"], dtype=np.float64)
        tail = float(np.mean(r[-100:]))
        if tail > best_tail:
            best_tail = tail
            best_seed = s
    return int(best_seed), float(best_tail)


def load_ppo_series(seed: int) -> Dict[str, np.ndarray]:
    z = np.load(ppo_npz(seed))
    return {
        "obs": np.asarray(z["ppo_obs"], dtype=np.float64),
        "act": np.asarray(z["ppo_actions"], dtype=np.float64),
        "rew": np.asarray(z["ppo_rewards"], dtype=np.float64),
    }


def plot_obs_act_time(ppo: Dict[str, np.ndarray], v6: Dict[str, np.ndarray], v7: Dict[str, np.ndarray], out_png: str) -> None:
    t = np.arange(EVAL_STEPS)
    fig, axes = plt.subplots(5, 1, figsize=(12, 12), sharex=True, constrained_layout=True)

    series = [
        (0, "obs0", ppo["obs"][:, 0], v6["obs"][:, 0], v7["obs"][:, 0]),
        (1, "obs1", ppo["obs"][:, 1], v6["obs"][:, 1], v7["obs"][:, 1]),
        (2, "act0", ppo["act"][:, 0], v6["act"][:, 0], v7["act"][:, 0]),
        (3, "act1", ppo["act"][:, 1], v6["act"][:, 1], v7["act"][:, 1]),
        (4, "act2", ppo["act"][:, 2], v6["act"][:, 2], v7["act"][:, 2]),
    ]

    for idx, name, y0, y1, y2 in series:
        ax = axes[idx]
        ax.plot(t, y0, lw=1.2, label="PPO")
        ax.plot(t, y1, lw=1.2, label="v6")
        ax.plot(t, y2, lw=1.2, label="v7")
        ax.set_ylabel(name)
        ax.grid(alpha=0.25)
        if idx == 0:
            ax.legend(loc="best", ncol=3)
    axes[-1].set_xlabel("time step")
    fig.suptitle("Obs-Act time series comparison (300 steps)")
    fig.savefig(out_png, dpi=170)
    plt.close(fig)


def plot_vorticity_maps(omega_ppo: np.ndarray, omega_v6: np.ndarray, omega_v7: np.ndarray, out_png: str) -> float:
    gmax = float(max(np.max(np.abs(omega_ppo)), np.max(np.abs(omega_v6)), np.max(np.abs(omega_v7))))
    vmax = max(1e-12, 0.1 * gmax)

    fig, axes = plt.subplots(1, 3, figsize=(15, 4.8), constrained_layout=True)
    data = [(omega_ppo, "PPO"), (omega_v6, "v6"), (omega_v7, "v7")]
    for ax, (om, title) in zip(axes, data):
        im = ax.imshow(om.T, origin="lower", cmap="RdBu_r", vmin=-vmax, vmax=vmax)
        ax.set_title(title)
        ax.set_xlabel("x")
        ax.set_ylabel("y")
    cbar = fig.colorbar(im, ax=axes.ravel().tolist(), shrink=0.95)
    cbar.set_label("vorticity")
    fig.suptitle(f"Final vorticity maps, unified range +/-{vmax:.4f} (10% global max)")
    fig.savefig(out_png, dpi=170)
    plt.close(fig)
    return float(vmax)


def build_design(obs: np.ndarray, act: np.ndarray, terms: List[str]) -> Tuple[np.ndarray, np.ndarray]:
    X_rows = []
    Y_rows = []
    for t in range(1, len(obs)):
        fd = feature_dict(obs[t], obs[t - 1], act[t - 1])
        X_rows.append([1.0] + [float(fd[k]) for k in terms])
        Y_rows.append(act[t])
    return np.asarray(X_rows, dtype=np.float64), np.asarray(Y_rows, dtype=np.float64)


def fit_multi_seed_linear(terms: List[str]) -> Dict[str, object]:
    Ws = []
    r2s = []
    for s in SEEDS:
        z = np.load(ppo_npz(s))
        obs = np.asarray(z["ppo_obs"], dtype=np.float64)
        act = np.asarray(z["ppo_actions"], dtype=np.float64)
        X, Y = build_design(obs, act, terms)
        W = []
        r2_ch = []
        for ch in range(3):
            c, *_ = np.linalg.lstsq(X, Y[:, ch], rcond=None)
            pred = X @ c
            ssr = float(np.sum((Y[:, ch] - pred) ** 2))
            sst = float(np.sum((Y[:, ch] - np.mean(Y[:, ch])) ** 2) + 1e-12)
            r2 = float(1.0 - ssr / sst)
            W.append(c)
            r2_ch.append(r2)
        Ws.append(np.asarray(W, dtype=np.float64))
        r2s.append(np.asarray(r2_ch, dtype=np.float64))
    Wm = np.mean(np.asarray(Ws), axis=0)
    r2m = np.mean(np.asarray(r2s), axis=0)
    return {
        "terms": ["bias1"] + list(terms),
        "W_mean": Wm,
        "r2_by_action_mean": r2m,
        "r2_mean": float(np.mean(r2m)),
    }


def matrix_to_latex(W: np.ndarray, precision: int = 4) -> str:
    rows = []
    for i in range(W.shape[0]):
        rows.append(" & ".join([f"{float(v):.{precision}f}" for v in W[i]]))
    body = " \\\\ ".join(rows)
    return "\\begin{bmatrix}" + body + "\\end{bmatrix}"


def main() -> None:
    v6 = load_run(V6_RUN)
    v7 = load_run(V7_RUN)

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

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

    v6_x2, v6_p2, v6_evr = fit_pca_and_project(v6.x, v6_ppo_pts)
    v7_x2, v7_p2, v7_evr = fit_pca_and_project(v7.x, v7_ppo_pts)

    fig_v6_pca = os.path.join(OUT_DIR, "brief_v6_3_pca_with_ppo3.png")
    fig_v7_pca = os.path.join(OUT_DIR, "brief_v7_1_pca_with_ppo3.png")
    plot_pca(v6.name, v6_x2, v6.y, v6_p2, v6_evr, fig_v6_pca)
    plot_pca(v7.name, v7_x2, v7.y, v7_p2, v7_evr, fig_v7_pca)

    with open(os.path.join(ROOT, "output", f"{V6_RUN}_best_meta.pkl"), "rb") as f:
        v6_meta = pickle.load(f)
    with open(os.path.join(ROOT, "output", f"{V7_RUN}_best_meta.pkl"), "rb") as f:
        v7_meta = pickle.load(f)

    v6_best = np.asarray(v6_meta["best_params"], dtype=np.float64)
    v7_best = np.asarray(v7_meta["best_params"], dtype=np.float64)

    v6_roll = rollout_controller_v6(v6_best, basis_terms=v6_basis, device_id=1)
    v7_roll = rollout_controller_v7(v7_best, k=k, period_steps=period_steps, device_id=0)

    seed_star, tail_star = best_ppo_seed()
    ppo_series = load_ppo_series(seed_star)

    fig_ts = os.path.join(OUT_DIR, "brief_v6_v7_ppo_obs_act_time.png")
    plot_obs_act_time(ppo_series, v6_roll, v7_roll, fig_ts)

    ppo_flow_npz = os.path.join(OUT_DIR, "fig_oldenv_flow_best.npz")
    zf = np.load(ppo_flow_npz)
    ppo_ux = np.asarray(zf["ux"], dtype=np.float64)
    ppo_uy = np.asarray(zf["uy"], dtype=np.float64)

    omega_ppo = vorticity_from_ux_uy(ppo_ux, ppo_uy)
    omega_v6 = vorticity_from_ux_uy(v6_roll["ux"], v6_roll["uy"])
    omega_v7 = vorticity_from_ux_uy(v7_roll["ux"], v7_roll["uy"])

    fig_omega = os.path.join(OUT_DIR, "brief_v6_v7_ppo_final_vorticity.png")
    omega_vmax = plot_vorticity_maps(omega_ppo, omega_v6, omega_v7, fig_omega)

    with open(os.path.join(OUT_DIR, "sindy_group_sparsity_scan.json"), "r", encoding="utf-8") as f:
        group_scan = json.load(f)
    with open(os.path.join(OUT_DIR, "sindy_constrained_profile_search.json"), "r", encoding="utf-8") as f:
        constrained = json.load(f)
    with open(os.path.join(OUT_DIR, "dante_ackley_highdim_sweep.json"), "r", encoding="utf-8") as f:
        highdim = json.load(f)
    with open(os.path.join(OUT_DIR, "v6_v7_reaudit_frozen_facts_20260323.json"), "r", encoding="utf-8") as f:
        frozen = json.load(f)
    with open(os.path.join(OUT_DIR, "v6_v7_acq_diagnostics_summary_20260323.json"), "r", encoding="utf-8") as f:
        acq = json.load(f)

    full_terms = [
        "obs0", "obs1", "dobs0", "dobs1", "sin_obs0", "sin_obs1", "cos_obs0", "cos_obs1",
        "tanh_obs0", "tanh_obs1", "act0_l1", "act1_l1", "act2_l1",
    ]
    reduced_terms = list(constrained["best_chrono"]["basis_terms"])

    full_fit = fit_multi_seed_linear(full_terms)
    reduced_fit = fit_multi_seed_linear(reduced_terms)

    highdim_results = list(highdim.get("results", []))
    highdim_neg_r2 = int(sum(1 for r in highdim_results if float(r.get("holdout_r2", 0.0)) < 0.0))
    highdim_neg_gain = int(sum(1 for r in highdim_results if float(r.get("acq_gain_scaled", 0.0)) < 0.0))

    summary = {
        "runs": {"v6": V6_RUN, "v7": V7_RUN, "ppo_seed_used": int(seed_star), "ppo_tail100": float(tail_star)},
        "figures": {
            "v6_pca": fig_v6_pca,
            "v7_pca": fig_v7_pca,
            "obs_act_time": fig_ts,
            "vorticity": fig_omega,
        },
        "vorticity_unified_abs_limit": float(omega_vmax),
        "reduction": {
            "full_r2_mean_over_seeds": float(group_scan["full_model"]["r2_mean_over_seeds"]),
            "best_chrono_basis_terms": reduced_terms,
            "best_chrono_r2_mean_over_seeds": float(constrained["best_chrono"]["chrono_r2_mean_over_seeds"]),
            "best_chrono_r2_min_over_seeds": float(constrained["best_chrono"]["chrono_r2_min_over_seeds"]),
            "group_scan_chosen": group_scan.get("chosen", {}),
            "full_fit": {
                "terms": full_fit["terms"],
                "W_mean": np.asarray(full_fit["W_mean"]).tolist(),
                "r2_by_action_mean": np.asarray(full_fit["r2_by_action_mean"]).tolist(),
                "r2_mean": float(full_fit["r2_mean"]),
            },
            "reduced_fit": {
                "terms": reduced_fit["terms"],
                "W_mean": np.asarray(reduced_fit["W_mean"]).tolist(),
                "r2_by_action_mean": np.asarray(reduced_fit["r2_by_action_mean"]).tolist(),
                "r2_mean": float(reduced_fit["r2_mean"]),
            },
            "latex": {
                "full_terms": "\\phi_{full}=[1,obs_0,obs_1,\\Delta obs_0,\\Delta obs_1,\\sin(\\pi obs_0),\\sin(\\pi obs_1),\\cos(\\pi obs_0),\\cos(\\pi obs_1),\\tanh(obs_0),\\tanh(obs_1),a_{0,t-1},a_{1,t-1},a_{2,t-1}]^\\top",
                "reduced_terms": "\\phi_{red}=[1,obs_1,\\sin(\\pi obs_0),\\cos(\\pi obs_0),a_{1,t-1}]^\\top",
                "W_full_mean": matrix_to_latex(np.asarray(full_fit["W_mean"], dtype=np.float64)),
                "W_reduced_mean": matrix_to_latex(np.asarray(reduced_fit["W_mean"], dtype=np.float64)),
            },
        },
        "highdim": {
            "n_cases": int(len(highdim_results)),
            "n_neg_holdout_r2": int(highdim_neg_r2),
            "n_neg_acq_gain": int(highdim_neg_gain),
        },
        "dual_surrogate": frozen.get("surrogate_log_stats", {}),
        "acq_summary": acq,
    }

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

    md_path = os.path.join(OUT_DIR, "v6_v7_ppo_briefing_20260324.md")
    with open(md_path, "w", encoding="utf-8") as f:
        f.write("# v6/v7/PPO 综合简报\n\n")
        f.write("## 1) PCA（含3个PPO拟合点）\n")
        f.write(f"- v6图: {fig_v6_pca}\n")
        f.write(f"- v7图: {fig_v7_pca}\n\n")

        f.write("## 2) obs-act时序 + 最终涡量\n")
        f.write(f"- 时序图: {fig_ts}\n")
        f.write(f"- 最终涡量图: {fig_omega}\n")
        f.write(f"- 涡量统一色条范围: ±{omega_vmax:.6f}（按三者全局|omega|max的10%）\n")
        f.write(f"- PPO时序选用seed={seed_star}（tail100={tail_star:.5f}）\n\n")

        f.write("## 3) 函数约简、关键函数、最终组合\n")
        f.write(f"- 全特征模型平均R2: {group_scan['full_model']['r2_mean_over_seeds']:.6f}\n")
        f.write(f"- 约束搜索best_chrono基函数: {reduced_terms}\n")
        f.write(f"- best_chrono平均R2: {constrained['best_chrono']['chrono_r2_mean_over_seeds']:.6f}\n")
        f.write(f"- best_chrono最差seed R2: {constrained['best_chrono']['chrono_r2_min_over_seeds']:.6f}\n")
        f.write("- 关键函数解释: obs1给出主状态幅值，sin/cos(pi*obs0)提供周期相位，act1_l1提供单步记忆。\n\n")

        f.write("### 学到方程与拟合方程（LaTeX）\n")
        f.write("- 全特征学到方程（3动作联合线性写法）:\n")
        f.write("$$\\mathbf{a}_t = W_{full}\\,\\phi_{full,t}$$\n")
        f.write("$$" + summary['reduction']['latex']['full_terms'] + "$$\n")
        f.write("$$W_{full}=" + summary['reduction']['latex']['W_full_mean'] + "$$\n")
        f.write(f"- 全特征拟合R2(本次重算, action均值): {full_fit['r2_mean']:.6f}\n\n")

        f.write("- 约简拟合方程（best_chrono）:\n")
        f.write("$$\\mathbf{a}_t = W_{red}\\,\\phi_{red,t}$$\n")
        f.write("$$" + summary['reduction']['latex']['reduced_terms'] + "$$\n")
        f.write("$$W_{red}=" + summary['reduction']['latex']['W_reduced_mean'] + "$$\n")
        f.write(f"- 约简拟合R2(本次重算, action均值): {reduced_fit['r2_mean']:.6f}\n\n")

        f.write("## 4) 高维能力与现实约束反思\n")
        f.write(f"- 高维扫描样本数: {len(highdim_results)}\n")
        f.write(f"- holdout_r2<0 的case数: {highdim_neg_r2}/{len(highdim_results)}\n")
        f.write(f"- 首轮acq_gain<0 的case数: {highdim_neg_gain}/{len(highdim_results)}\n")
        f.write("- 解释: 高维下代理可辨识度不足时，Tree探索会被误导，出现边界漂移与收益停滞。\n")
        f.write("- 与论文对齐: DANTE强调DNN surrogate与自适应探索协同；在你的流体控制里，受噪声、时序漂移和高维参数化耦合影响，样本效率会显著下降。\n")
        f.write("- 双代理是否有帮助: 当前日志显示ensemble路径提供了可运行冗余（CNN失败时MLP兜底），但在v6_3其val_r2均值仍偏低，收益主要体现在稳定性而非显著提升最优值。\n")

    print("saved:", summary_path)
    print("saved:", md_path)
    print("saved figs:", fig_v6_pca, fig_v7_pca, fig_ts, fig_omega)


if __name__ == "__main__":
    main()
