"""Evaluation helpers for PPO on CarRacing-v3.

Functions:
  - evaluate_agent(agent, n_episodes, seed_start)  -> list[float] returns
  - record_demo_video(agent, out_path, seed)       -> save mp4 of one episode
  - load_tb_scalars(run_dir, tag)                  -> (steps, values) from TensorBoard
  - plot_eval_bar(returns, baseline, save_path)    -> bar chart of eval returns
  - plot_training_curves(run_dirs, labels, save_path) -> 6-panel curves figure
"""

from pathlib import Path

import gymnasium as gym
import imageio
import matplotlib.pyplot as plt
import numpy as np

from src.env_wrappers import FrameStack, GrayScaleResize, SkipFrame, make_env


# ---------------------------------------------------------------------------
# Numerical evaluation
# ---------------------------------------------------------------------------

def evaluate_agent(agent, n_episodes=20, seed_start=1000, deterministic=False):
    """Roll the agent for n_episodes on freshly-seeded envs.

    Args:
        agent: a PPOAgent (loaded with a checkpoint).
        n_episodes: how many evaluation episodes.
        seed_start: starting seed; each ep uses seed_start + ep.
        deterministic: if True, take argmax over policy logits instead of
            sampling (slightly higher mean, slightly lower variance).

    Returns:
        list[float] of per-episode returns.
    """
    import torch
    from torch.distributions import Categorical

    returns = []
    env = make_env(seed=seed_start)
    for ep in range(n_episodes):
        obs, _ = env.reset(seed=seed_start + ep)
        ep_return, done = 0.0, False
        while not done:
            if deterministic:
                obs_t = torch.as_tensor(obs, device=agent.device).unsqueeze(0)
                logits, _ = agent.net(obs_t)
                action = int(logits.argmax(dim=-1).item())
            else:
                action, _, _ = agent.act(obs)
            obs, r, term, trunc, _ = env.step(action)
            ep_return += r
            done = term or trunc
        returns.append(ep_return)
    env.close()
    return returns


# ---------------------------------------------------------------------------
# Video recording
# ---------------------------------------------------------------------------

class _DualEnv:
    """Step a wrapped env (for the agent) and a raw env (for nice video) in lockstep."""

    def __init__(self, seed):
        self.wrapped = make_env(seed=seed)
        self.raw = gym.make("CarRacing-v3", continuous=False, render_mode="rgb_array")
        # Don't pre-reset raw env here; reset() below will do it once cleanly.

    def reset(self, seed):
        # Re-create wrapped to ensure both envs see the EXACT same first reset
        # under the given seed. This mirrors how the eval pipeline scores seeds.
        self.wrapped = make_env(seed=seed)
        wrapped_obs, _ = self.wrapped.reset(seed=seed)
        raw_obs, _ = self.raw.reset(seed=seed)
        return wrapped_obs, raw_obs

    def step(self, action):
        wrapped_obs, _, term, trunc, _ = self.wrapped.step(action)
        # SkipFrame inside wrapped env runs 4 raw frames per call; mirror that.
        raw_frames = []
        raw_done = False
        for _ in range(4):
            if not raw_done:
                raw_obs, _, t, tr, _ = self.raw.step(action)
                raw_done = t or tr
                raw_frames.append(raw_obs.copy())
        return wrapped_obs, raw_frames, (term or trunc) or raw_done

    def close(self):
        self.wrapped.close()
        self.raw.close()


def record_demo_video(agent, out_path, seed=42, fps=30, max_steps=600):
    """Record one evaluation episode as an mp4 using the original RGB renderer."""
    out_path = Path(out_path)
    out_path.parent.mkdir(parents=True, exist_ok=True)

    denv = _DualEnv(seed=seed)
    wrapped_obs, _ = denv.reset(seed=seed)

    frames = []
    total_r = 0.0
    for _ in range(max_steps):
        action, _, _ = agent.act(wrapped_obs)
        wrapped_obs, raw_frames, done = denv.step(action)
        frames.extend(raw_frames)
        if done:
            break

    denv.close()
    imageio.mimsave(str(out_path), frames, fps=fps)
    return len(frames), out_path


# ---------------------------------------------------------------------------
# TensorBoard data extraction
# ---------------------------------------------------------------------------

def load_tb_scalars(run_dir, tag):
    """Read a single scalar tag from a TensorBoard run directory.

    Returns (steps_list, values_list); empty if tag is absent.
    """
    from tensorboard.backend.event_processing.event_accumulator import EventAccumulator

    ea = EventAccumulator(str(run_dir))
    ea.Reload()
    if tag not in ea.Tags()["scalars"]:
        return [], []
    events = ea.Scalars(tag)
    return [e.step for e in events], [e.value for e in events]


# ---------------------------------------------------------------------------
# Plotting
# ---------------------------------------------------------------------------

def plot_eval_bar(returns, baseline, save_path, title=None):
    """Bar chart of per-episode returns + baseline reference line."""
    save_path = Path(save_path)
    save_path.parent.mkdir(parents=True, exist_ok=True)

    mean_r = float(np.mean(returns))
    std_r = float(np.std(returns))

    fig, ax = plt.subplots(figsize=(10, 5))
    xs = np.arange(len(returns))
    ax.bar(xs, returns, color="steelblue", edgecolor="black", alpha=0.7)
    ax.axhline(y=mean_r, color="red", linestyle="--",
               label=f"Mean = {mean_r:.1f} ± {std_r:.1f}")
    ax.axhline(y=baseline, color="gray", linestyle=":",
               label=f"Random baseline = {baseline:.1f}")
    ax.set_xlabel("Evaluation episode")
    ax.set_ylabel("Episode return")
    ax.set_title(title or f"Evaluation returns over {len(returns)} unseen seeds")
    ax.legend()
    ax.grid(True, alpha=0.3)
    plt.tight_layout()
    plt.savefig(save_path, dpi=150)
    plt.close(fig)
    return save_path


def plot_training_curves(run_dirs, labels, save_path,
                         tags=None, smooth_window=10):
    """Multi-run multi-panel training curves (one panel per tag).

    Args:
        run_dirs: list[Path] — TensorBoard run directories
        labels: list[str] — label per run for the legend
        save_path: where to write the PNG
        tags: list[str] of TensorBoard scalar tags to plot. Defaults to a
              standard 6-panel set.
        smooth_window: rolling-mean window for visual smoothing (1 = none).
    """
    save_path = Path(save_path)
    save_path.parent.mkdir(parents=True, exist_ok=True)

    if tags is None:
        tags = [
            "episode/avg_return_100",
            "losses/value_loss",
            "losses/entropy",
            "losses/approx_kl",
            "losses/clip_frac",
            "episode/length",
        ]

    n = len(tags)
    rows = (n + 2) // 3
    fig, axes = plt.subplots(rows, 3, figsize=(15, 4 * rows))
    axes = axes.flatten()

    for ax, tag in zip(axes, tags):
        for run_dir, label in zip(run_dirs, labels):
            steps, values = load_tb_scalars(run_dir, tag)
            if not steps:
                continue
            if smooth_window > 1 and len(values) > smooth_window:
                values = np.convolve(
                    values, np.ones(smooth_window) / smooth_window, mode="valid"
                )
                steps = steps[smooth_window - 1:]
            ax.plot(steps, values, label=label, alpha=0.85)
        ax.set_title(tag)
        ax.set_xlabel("Env steps")
        ax.grid(True, alpha=0.3)
        ax.legend(fontsize=8)

    # Hide unused axes
    for ax in axes[n:]:
        ax.set_visible(False)

    plt.tight_layout()
    plt.savefig(save_path, dpi=150)
    plt.close(fig)
    return save_path
