"""Evaluate sumo levels: round-robin win-rate matrix + robustness under perturbations.

Final evaluation uses plain match results only (win / loss / tie) -- no reward
shaping. "Score" = (wins + 0.5 * ties) / matches, so 0.5 means evenly matched.

Usage:
  python evaluate.py                    # 400 matches per cell (~2-3 min)
  python evaluate.py --quick            # 100 matches per cell (~40 s)
  python evaluate.py --weights-dir weights_quick
Outputs: plots/winrate_matrix.png, plots/robustness.png, plots/training_curves.png,
         results/eval.json, and markdown tables on stdout.
"""

from __future__ import annotations

import argparse
import json
import os
import time

import matplotlib

matplotlib.use("Agg")
import matplotlib.pyplot as plt  # noqa: E402
import numpy as np  # noqa: E402
from matplotlib.colors import LinearSegmentedColormap  # noqa: E402

from policies import LEVELS, WEIGHTS_DIR, load_level, run_matches  # noqa: E402
from sumo_env import PERTURBATIONS, PhysParams, SumoEnv, apply_perturbation  # noqa: E402

HERE = os.path.dirname(os.path.abspath(__file__))

# Chart colours (validated reference palette: categorical slots in fixed order,
# blue sequential ramp, blue<->red diverging with a gray midpoint).
SERIES = ["#2a78d6", "#eb6834", "#1baf7a", "#eda100", "#e87ba4", "#008300"]
SURFACE, INK, INK2, GRID = "#fcfcfb", "#0b0b0b", "#52514e", "#e4e3df"
DIVERGING = LinearSegmentedColormap.from_list("rb", ["#d03b3b", "#f0efec", "#2a78d6"])
SEQUENTIAL = LinearSegmentedColormap.from_list("blue", ["#f0efec", "#cde2fb", "#6da7ec", "#256abf", "#104281"])


def style_axes(ax):
    ax.set_facecolor(SURFACE)
    for s in ("top", "right"):
        ax.spines[s].set_visible(False)
    for s in ("left", "bottom"):
        ax.spines[s].set_color(INK2)
    ax.tick_params(colors=INK2)


def play(level0, level1, n, perturbation="normal", seed=0, weights_dir=WEIGHTS_DIR):
    """n matches of level0 (robot 0) vs level1 (robot 1). Returns W/L/T rates."""
    params = apply_perturbation(PhysParams.nominal(n), perturbation)
    env = SumoEnv(n, params, seed=seed)  # same seed => same start positions for every pairing
    res = run_matches(env, load_level(level0, weights_dir), load_level(level1, weights_dir),
                      np.random.default_rng(seed))
    o = res["outcome"]
    return {"win": float((o == 1).mean()), "loss": float((o == -1).mean()), "tie": float((o == 0).mean())}


def score(r):
    return r["win"] + 0.5 * r["tie"]


def round_robin(levels, n, weights_dir):
    """Matrix of results, row level = robot 0, column level = robot 1."""
    return [[play(a, b, n, seed=1000 + 17 * i + j, weights_dir=weights_dir) for j, b in enumerate(levels)]
            for i, a in enumerate(levels)]


def robustness(levels, opponents, n, weights_dir):
    """{opponent: {level: {perturbation: result}}}. Perturbations hit robot 0 (the level)
    or the environment; the opponent itself is never handicapped."""
    return {opp: {lv: {p: play(lv, opp, n, p, seed=5000 + k, weights_dir=weights_dir)
                       for k, p in enumerate(PERTURBATIONS)} for lv in levels} for opp in opponents}


# ---------------------------------------------------------------- plots
def plot_matrix(levels, rr, path):
    S = np.array([[score(c) for c in row] for row in rr])
    fig, ax = plt.subplots(figsize=(6.4, 5.4), facecolor=SURFACE)
    im = ax.imshow(S, cmap=DIVERGING, vmin=0, vmax=1)
    for i in range(len(levels)):
        for j in range(len(levels)):
            c = rr[i][j]
            ax.text(j, i, f"{score(c):.2f}\n{c['win']:.0%}/{c['tie']:.0%}", ha="center", va="center",
                    fontsize=8, color=INK)
    ax.set_xticks(range(len(levels)), levels)
    ax.set_yticks(range(len(levels)), levels)
    ax.set_xlabel("opponent (column)", color=INK2)
    ax.set_ylabel("player (row)", color=INK2)
    ax.set_title("Round-robin score of row vs column\n(score = win + ½·tie;  cell: score, win% / tie%)",
                 color=INK, fontsize=10)
    ax.tick_params(colors=INK2)
    cb = fig.colorbar(im, ax=ax, fraction=0.046)
    cb.set_label("score (0.5 = even)", color=INK2)
    fig.tight_layout()
    fig.savefig(path, dpi=130, facecolor=SURFACE)
    plt.close(fig)


def plot_robustness(levels, rob, path):
    opps = list(rob)
    fig, axes = plt.subplots(1, len(opps), figsize=(5.2 * len(opps), 4.6), facecolor=SURFACE, squeeze=False)
    for ax, opp in zip(axes[0], opps):
        S = np.array([[score(rob[opp][lv][p]) for p in PERTURBATIONS] for lv in levels])
        ax.imshow(S, cmap=SEQUENTIAL, vmin=0, vmax=1)
        for i in range(len(levels)):
            for j in range(len(PERTURBATIONS)):
                ax.text(j, i, f"{S[i, j]:.2f}", ha="center", va="center", fontsize=9,
                        color="white" if S[i, j] > 0.6 else INK)
        ax.set_xticks(range(len(PERTURBATIONS)), PERTURBATIONS, rotation=20)
        ax.set_yticks(range(len(levels)), levels)
        ax.tick_params(colors=INK2)
        ax.set_title(f"score vs {opp} under perturbations", color=INK, fontsize=10)
    fig.tight_layout()
    fig.savefig(path, dpi=130, facecolor=SURFACE)
    plt.close(fig)


def plot_training(weights_dir, path):
    """Win rate of the training policy vs fixed opponents, across L2 -> L3 -> L4."""
    logs = []
    for st in ("L2", "L3", "L4"):
        f = os.path.join(weights_dir, f"log_{st}.json")
        if os.path.exists(f):
            with open(f) as fh:
                logs.append((st, json.load(fh)))
    if not logs:
        return False
    fig, ax = plt.subplots(figsize=(8, 4.2), facecolor=SURFACE)
    style_axes(ax)
    offset, series = 0, {}
    for st, log in logs:
        for g, ev in zip(log["eval_gen"], log["eval"]):
            for opp, (w, l, t) in ev.items():
                key = "L2 (first checkpoint of pool)" if opp == "pool0" else opp
                series.setdefault(key, ([], []))
                series[key][0].append(offset + g)
                series[key][1].append(w + 0.5 * t)
        ax.axvline(offset, color=GRID, lw=1)
        ax.text(offset + 2, 1.02, st, color=INK2, fontsize=9)
        offset += log["gen"][-1]
    for k, (name, (x, y)) in enumerate(series.items()):
        ax.plot(x, y, color=SERIES[k], lw=2, marker="o", ms=4, label=name)
    ax.axhline(0.5, color=INK2, lw=0.8, ls=":")
    ax.set_ylim(0, 1.08)
    ax.set_xlabel("ES generation (cumulative over stages)", color=INK2)
    ax.set_ylabel("score of current policy", color=INK2)
    ax.set_title("Training progress (nominal physics; L4 trains under randomization)", color=INK, fontsize=10)
    ax.legend(frameon=False, fontsize=8, loc="lower left")
    ax.grid(axis="y", color=GRID, lw=0.6)
    fig.tight_layout()
    fig.savefig(path, dpi=130, facecolor=SURFACE)
    plt.close(fig)
    return True


# ---------------------------------------------------------------- tables
def md_matrix(levels, rr):
    lines = ["| row vs col | " + " | ".join(levels) + " |", "|---" * (len(levels) + 1) + "|"]
    for i, a in enumerate(levels):
        lines.append(f"| **{a}** | " + " | ".join(f"{score(c):.2f}" for c in rr[i]) + " |")
    return "\n".join(lines)


def md_robust(levels, res):
    lines = ["| level | " + " | ".join(PERTURBATIONS) + " | mean |", "|---" * (len(PERTURBATIONS) + 2) + "|"]
    for lv in levels:
        s = [score(res[lv][p]) for p in PERTURBATIONS]
        lines.append(f"| {lv} | " + " | ".join(f"{x:.2f}" for x in s) + f" | {np.mean(s):.2f} |")
    return "\n".join(lines)


def main():
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--quick", action="store_true", help="100 matches per cell instead of 400")
    ap.add_argument("-n", type=int, default=None, help="matches per cell")
    ap.add_argument("--weights-dir", default=WEIGHTS_DIR)
    ap.add_argument("--levels", nargs="+", default=None, help=f"default: {' '.join(LEVELS)}")
    ap.add_argument("--robust-opponents", nargs="+", default=["L1", "Dodger", "L3"])
    ap.add_argument("--out", default=HERE, help="where plots/ and results/ go")
    args = ap.parse_args()
    n = args.n or (100 if args.quick else 400)

    levels = args.levels or LEVELS
    available = []
    for lv in levels:
        try:
            load_level(lv, args.weights_dir)
            available.append(lv)
        except FileNotFoundError as e:
            print(f"skipping {lv}: {e}")
    levels = available
    os.makedirs(os.path.join(args.out, "plots"), exist_ok=True)
    os.makedirs(os.path.join(args.out, "results"), exist_ok=True)

    t0 = time.time()
    rr = round_robin(levels, n, args.weights_dir)
    print(f"\n## Round-robin score (row vs column, {n} matches each)  [{time.time() - t0:.0f}s]\n")
    print(md_matrix(levels, rr))
    print("\nMean score vs all other levels (excluding self):")
    for i, a in enumerate(levels):
        print(f"  {a:7s} {np.mean([score(rr[i][j]) for j in range(len(levels)) if j != i]):.2f}")

    rob = robustness(levels, args.robust_opponents, n, args.weights_dir)
    for opp in rob:
        print(f"\n## Robustness: score vs {opp} under perturbations ({n} matches each)\n")
        print(md_robust(levels, rob[opp]))
    print(f"\n[total {time.time() - t0:.0f}s]")

    plot_matrix(levels, rr, os.path.join(args.out, "plots", "winrate_matrix.png"))
    plot_robustness(levels, rob, os.path.join(args.out, "plots", "robustness.png"))
    if plot_training(args.weights_dir, os.path.join(args.out, "plots", "training_curves.png")):
        print("saved plots/training_curves.png")
    with open(os.path.join(args.out, "results", "eval.json"), "w") as f:
        json.dump({"levels": levels, "n": n, "round_robin": rr, "robustness": rob}, f, indent=1)
    print("saved plots/winrate_matrix.png, plots/robustness.png, results/eval.json")


if __name__ == "__main__":
    main()