#!/usr/bin/env python3
"""Plot the red/blue phase diagram from a `sweep.py` result file.

Reads the JSONL produced by `tools/sweep.py` and renders:

  1. the winning side per (view_dist, radius) point, with the estimated phase
     boundary (mean RED-BLU food-share difference == 0) overlaid, and
  2. the continuous mean-difference field with a diverging colour map.

Markers are sized by the number of seeds the sequential test needed, so it is
easy to see where the search had to work hardest.

Example
-------
    python3 tools/plot_phase.py sweep_out/results.jsonl -o sweep_out/phase.png
"""

from __future__ import annotations

import argparse
import json
from pathlib import Path

import matplotlib

matplotlib.use("Agg")
import matplotlib.pyplot as plt
import numpy as np
from matplotlib.colors import TwoSlopeNorm
from matplotlib.lines import Line2D

RED = "#d62728"
BLU = "#1f77b4"
TIE = "#9e9e9e"


def load(path: Path) -> dict:
    if path.suffix == ".csv":
        import csv
        with path.open() as f:
            rows = list(csv.DictReader(f))
    else:
        rows = []
        with path.open() as f:
            for line in f:
                line = line.strip()
                if line:
                    rows.append(json.loads(line))
    cols = {
        "vd": np.array([float(r["view_dist"]) for r in rows]),
        "rad": np.array([float(r["weapon_radius"]) for r in rows]),
        "diff": np.array([float(r["mean_diff"]) for r in rows]),
        "sigma": np.array([float(r["sigma"]) for r in rows]),
        "n": np.array([float(r["n"]) for r in rows]),
        "winner": [r["winner"] for r in rows],
        "stopped": [r["stopped"] for r in rows],
    }
    return cols


def outcome(winner: str, stopped: str) -> str:
    return "tie" if stopped != "sigma" else winner


def marker_sizes(n: np.ndarray, lo: float = 18.0, hi: float = 260.0) -> np.ndarray:
    if n.max() <= n.min():
        return np.full_like(n, (lo + hi) / 2)
    scaled = (n - n.min()) / (n.max() - n.min())
    return lo + scaled * (hi - lo)


def main() -> int:
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("results", type=Path, help="results.jsonl or results.csv")
    ap.add_argument("-o", "--out", type=Path, default=Path("sweep_out/phase.png"))
    ap.add_argument("--title", default="Adversarial foraging: red (aggressive) vs blue (pacifist)")
    ap.add_argument("--dpi", type=int, default=180)
    args = ap.parse_args()

    d = load(args.results)
    vd, rad, diff, sigma, n = d["vd"], d["rad"], d["diff"], d["sigma"], d["n"]
    labels = [outcome(w, s) for w, s in zip(d["winner"], d["stopped"])]
    color_of = {"red": RED, "blu": BLU, "tie": TIE}
    colors = [color_of[l] for l in labels]
    sizes = marker_sizes(n)

    # Regular grid for a smooth background field, if the points are lattice-like.
    fig, axes = plt.subplots(1, 2, figsize=(15, 6.2), sharex=True, sharey=True)

    # ---- Panel 1: discrete winner regions -------------------------------- #
    ax = axes[0]
    ax.scatter(vd, rad, s=sizes, c=colors, marker="s", edgecolors="black",
               linewidths=0.6, alpha=0.9, zorder=3)
    # Boundary contour from the continuous field.
    try:
        ax.tricontour(vd, rad, diff, levels=[0.0], colors="black",
                      linewidths=2.0, zorder=4)
    except Exception:
        pass
    ax.set_title("Winner (marker size ∝ seeds needed)")
    ax.set_xlabel("view_dist")
    ax.set_ylabel("weapon radius")
    ax.set_axisbelow(True)
    ax.grid(alpha=0.25)
    handles = [
        Line2D([], [], marker="s", linestyle="", color=RED, markersize=11, label="red wins"),
        Line2D([], [], marker="s", linestyle="", color=BLU, markersize=11, label="blue wins"),
        Line2D([], [], marker="s", linestyle="", color=TIE, markersize=11, label="undecided (hit cap)"),
        Line2D([], [], color="black", linewidth=2, label="boundary: mean diff = 0"),
    ]
    ax.legend(handles=handles, loc="best", framealpha=0.9)

    # ---- Panel 2: continuous mean-difference field ----------------------- #
    ax = axes[1]
    vmax = max(abs(diff).max(), 1e-6)
    norm = TwoSlopeNorm(vmin=-vmax, vcenter=0.0, vmax=vmax)
    sc = ax.scatter(vd, rad, s=sizes, c=diff, cmap="coolwarm", norm=norm,
                    marker="s", edgecolors="black", linewidths=0.6, zorder=3)
    try:
        cs = ax.tricontour(vd, rad, diff, levels=[0.0], colors="black",
                           linewidths=2.0, zorder=4)
    except Exception:
        pass
    # A faint sigma=3 contour shows where confidence is marginal vs decisive.
    try:
        ax.tricontour(vd, rad, sigma, levels=[3.0], colors="black",
                      linestyles="dotted", linewidths=1.2, zorder=4)
    except Exception:
        pass
    cb = fig.colorbar(sc, ax=ax, fraction=0.046, pad=0.04)
    cb.set_label("mean RED-BLU food-share difference")
    ax.set_title("Mean difference (dotted: |diff|/SE = 3)")
    ax.set_xlabel("view_dist")
    ax.set_axisbelow(True)
    ax.grid(alpha=0.25)

    fig.suptitle(args.title, fontsize=13)
    fig.tight_layout(rect=(0, 0, 1, 0.96))
    args.out.parent.mkdir(parents=True, exist_ok=True)
    fig.savefig(args.out, dpi=args.dpi)
    print(f"wrote {args.out}")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())