"""Generate REAL baseline data for the HCM:21 PyTorch poster.

Runs two non-LLM policies through the standalone environment (no server, no API key):
  - Random agent  : picks valid actions per phase at random, advances when allowed
  - Heuristic agent: the built-in data-driven strategy from demo.py

Across 5 fixed seeds, records final episode score (0-1) and the per-quarter
reward trajectory. Saves results to results.json and renders poster_figure.png.
"""

from __future__ import annotations

import json
import random
from statistics import mean, pstdev

import numpy as np

from demo import run_demo_json
from hr_env.models import HRAction
from hr_env.server.environment import HRProductivityEnvironment
from hr_env.server.phases import PHASE_ACTIONS, PHASE_MIN_ACTIONS, PHASES

SEEDS = [42, 99, 123, 456, 789]
SIZE = 300
DEPARTMENTS = ["Engineering", "Sales", "Operations", "HR", "Finance"]


def _random_action(phase: str, rng: random.Random) -> HRAction:
    """Build a syntactically valid random action for the given phase."""
    atype = rng.choice(PHASE_ACTIONS[phase])
    dept = rng.choice(DEPARTMENTS)
    if atype == "query_department":
        return HRAction(action_type=atype, department=dept)
    if atype == "query_employees":
        return HRAction(action_type=atype, parameters={"min_performance": 0.0})
    if atype == "calculate_metric":
        return HRAction(action_type=atype, metric_name="all")
    if atype == "review_financials":
        return HRAction(action_type=atype)
    if atype == "set_hiring_target":
        return HRAction(action_type=atype, department=dept, count=rng.randint(0, 10))
    if atype == "set_training_budget":
        return HRAction(action_type=atype, department=dept, amount=rng.randint(0, 50_000))
    if atype == "set_compensation_policy":
        return HRAction(action_type=atype, department=dept, amount=rng.uniform(0, 5))
    if atype == "set_retention_program":
        return HRAction(action_type=atype, department=dept, amount=rng.randint(0, 30_000))
    if atype == "execute_hiring":
        return HRAction(action_type=atype, department=dept, count=rng.randint(0, 8))
    if atype == "execute_training":
        return HRAction(action_type=atype, department=dept, amount=rng.randint(0, 20))
    if atype in ("execute_promotion", "execute_transfer", "execute_termination"):
        return HRAction(action_type=atype, employee_ids=[])
    if atype == "submit_report":
        return HRAction(action_type=atype)
    return HRAction(action_type="review_financials")


def run_null_episode(seed: int, size: int = SIZE) -> dict:
    """Do-nothing baseline: satisfy phase minimums with no-op actions only."""
    env = HRProductivityEnvironment()
    env.reset(seed=seed, size=size)
    final_score = 0.0
    for _q in range(1, 7):
        for phase in PHASES:
            for _ in range(PHASE_MIN_ACTIONS[phase]):
                if phase == "scanning":
                    env.step(HRAction(action_type="review_financials"))
                elif phase == "planning":
                    env.step(HRAction(action_type="set_hiring_target", department="HR", count=0))
                elif phase == "producing":
                    env.step(HRAction(action_type="execute_hiring", department="HR", count=0))
                else:
                    env.step(HRAction(action_type="submit_report"))
            if phase != "controlling":
                env.step(HRAction(action_type="advance_phase"))
        obs = env.step(HRAction(action_type="advance_quarter"))
        if obs.done:
            final_score = (obs.data or {}).get("final_score", 0.0)
    return {"seed": seed, "final_score": round(final_score, 4)}


def run_random_episode(seed: int, size: int = SIZE) -> dict:
    """Run a full 6-quarter episode taking random valid actions each phase."""
    rng = random.Random(seed)
    env = HRProductivityEnvironment()
    env.reset(seed=seed, size=size)

    quarterly_rewards: list[float] = []
    final_score = 0.0
    total_steps = 0

    for _q in range(1, 7):
        for phase in PHASES:
            # take a couple of random valid actions to satisfy phase minimums
            n = PHASE_MIN_ACTIONS[phase] + rng.randint(0, 1)
            for _ in range(n):
                env.step(_random_action(phase, rng))
                total_steps += 1
            if phase != "controlling":
                env.step(HRAction(action_type="advance_phase"))
                total_steps += 1
        obs = env.step(HRAction(action_type="advance_quarter"))
        total_steps += 1
        data = obs.data or {}
        if obs.reward is not None and not obs.done:
            quarterly_rewards.append(round(obs.reward, 4))
        if obs.done:
            final_score = data.get("final_score", obs.reward or 0.0)
            break
    return {"seed": seed, "final_score": round(final_score, 4),
            "quarterly_rewards": quarterly_rewards, "total_steps": total_steps}


GRPO_RESULTS_FILE = "grpo_results.json"


def load_trained() -> dict | None:
    """Load Colab GRPO output if present (produced in RUNBOOK.md step 2.4).

    Expected schema (any extra keys ignored):
        {
          "trained":      [s1, s2, ...],   # per-seed final scores of the GRPO agent
          "reward_curve": [r1, r2, ...]    # optional: per-logging-step mean reward
        }
    Returns a normalized dict {"scores": [...], "mean", "std", "reward_curve"} or
    None if the file is missing/empty so the experiment runs unchanged.
    """
    import os
    if not os.path.exists(GRPO_RESULTS_FILE):
        return None
    try:
        with open(GRPO_RESULTS_FILE) as f:
            raw = json.load(f)
    except (json.JSONDecodeError, OSError):
        return None
    scores = raw.get("trained") or raw.get("trained_scores") or []
    scores = [float(s) for s in scores]
    if not scores:
        return None
    return {
        "scores": scores,
        "mean": round(mean(scores), 4),
        "std": round(pstdev(scores), 4),
        "reward_curve": raw.get("reward_curve") or [],
    }


def main() -> None:
    random_runs, heuristic_runs, null_runs = [], [], []

    print("Running NULL (do-nothing) agent...")
    for s in SEEDS:
        r = run_null_episode(s)
        null_runs.append(r)
        print(f"  seed {s:>3}: score {r['final_score']:.4f}")

    print("Running RANDOM agent...")
    for s in SEEDS:
        r = run_random_episode(s)
        random_runs.append(r)
        print(f"  seed {s:>3}: score {r['final_score']:.4f}  steps {r['total_steps']}")

    print("Running HEURISTIC agent...")
    for s in SEEDS:
        r = run_demo_json(seed=s, size=SIZE)
        heuristic_runs.append({"seed": s,
                               "final_score": round(r["final"]["score"], 4),
                               "quarterly_rewards": r["quarterly_rewards"],
                               "total_steps": r["final"]["total_steps"]})
        print(f"  seed {s:>3}: score {r['final']['score']:.4f}  steps {r['final']['total_steps']}")

    rnd = [r["final_score"] for r in random_runs]
    heu = [r["final_score"] for r in heuristic_runs]
    nul = [r["final_score"] for r in null_runs]
    summary = {
        "seeds": SEEDS, "size": SIZE,
        "null": {"runs": null_runs, "mean": round(mean(nul), 4), "std": round(pstdev(nul), 4)},
        "random": {"runs": random_runs, "mean": round(mean(rnd), 4), "std": round(pstdev(rnd), 4)},
        "heuristic": {"runs": heuristic_runs, "mean": round(mean(heu), 4), "std": round(pstdev(heu), 4)},
    }
    # Optional: fold in the GRPO-trained agent from the Colab run (RUNBOOK step 2).
    trained = load_trained()
    if trained:
        summary["trained"] = {
            "runs": [{"final_score": s} for s in trained["scores"]],
            "mean": trained["mean"], "std": trained["std"],
        }

    with open("results.json", "w") as f:
        json.dump(summary, f, indent=2)

    print(f"\nNULL      mean {summary['null']['mean']:.3f} +/- {summary['null']['std']:.3f}")
    print(f"RANDOM    mean {summary['random']['mean']:.3f} +/- {summary['random']['std']:.3f}")
    print(f"HEURISTIC mean {summary['heuristic']['mean']:.3f} +/- {summary['heuristic']['std']:.3f}")
    if trained:
        print(f"TRAINED   mean {trained['mean']:.3f} +/- {trained['std']:.3f}  "
              f"(from {GRPO_RESULTS_FILE})")

    make_figure(summary)
    out = "results.json and poster_figure.png"
    if trained and trained["reward_curve"]:
        make_grpo_curve(trained["reward_curve"])
        out += " and fig5_grpo_reward_curve.png"
    print("Wrote " + out)


def make_figure(summary: dict) -> None:
    import matplotlib
    matplotlib.use("Agg")
    import matplotlib.pyplot as plt

    # v1 (sigmoid-centred) measured numbers — the reward-hacked collapse, recorded
    # before the recalibration. v2 = current health-gated scoring (this run).
    v1 = {"null": 0.672, "random": 0.652, "heuristic": 0.649}
    labels = ["Null\n(do-nothing)", "Random", "Heuristic"]
    v2_means = [summary["null"]["mean"], summary["random"]["mean"], summary["heuristic"]["mean"]]
    v2_stds = [summary["null"]["std"], summary["random"]["std"], summary["heuristic"]["std"]]
    v1_means = [v1["null"], v1["random"], v1["heuristic"]]

    # Trained (GRPO) agent — only present after the Colab run. It has no v1
    # counterpart (it is trained against v2), so it gets a single green bar.
    trained = summary.get("trained")
    if trained:
        labels = labels + ["Trained\n(GRPO)"]

    fig, axes = plt.subplots(1, 2, figsize=(13, 5))

    # --- Left: before vs after recalibration (+ trained agent if available) ---
    x = np.arange(len(labels))
    w = 0.38
    # v1 bars only for the three non-trained baselines
    b1 = axes[0].bar(x[:3] - w / 2, v1_means, w, label="v1: sigmoid-centred (hacked)",
                     color="#bdc3c7", edgecolor="black", alpha=0.9)
    b2 = axes[0].bar(x[:3] + w / 2, v2_means, w, yerr=v2_stds, capsize=6,
                     label="v2: health-gated (this work)", color="#2980b9",
                     edgecolor="black", alpha=0.9)
    bars_to_label = [b1, b2]
    if trained:
        b3 = axes[0].bar(x[3] + w / 2, [trained["mean"]], w, yerr=[trained["std"]],
                         capsize=6, label="GRPO-trained agent (v2)", color="#27ae60",
                         edgecolor="black", alpha=0.95)
        bars_to_label.append(b3)
        # reference line at the heuristic v2 mean — the bar an agent must beat
        axes[0].axhline(summary["heuristic"]["mean"], color="#e67e22",
                        linestyle="--", linewidth=1.2)
        axes[0].text(x[3] + w / 2, summary["heuristic"]["mean"] + 0.01,
                     "heuristic bar", fontsize=7, color="#e67e22", ha="center")
    axes[0].axhspan(0.62, 0.69, color="red", alpha=0.08)
    axes[0].text(0.02, 0.70, "v1 collapse band (no separation)", fontsize=8, color="#c0392b")
    axes[0].set_xticks(x)
    axes[0].set_xticklabels(labels)
    axes[0].set_ylabel("Final episode score (0–1)")
    axes[0].set_title("Reward recalibration restores policy separation")
    axes[0].set_ylim(0, 1.0)
    axes[0].grid(True, axis="y", alpha=0.3)
    axes[0].legend(fontsize=8, loc="upper right")
    for bars in bars_to_label:
        for bar in bars:
            axes[0].text(bar.get_x() + bar.get_width() / 2, bar.get_height() + 0.015,
                         f"{bar.get_height():.2f}", ha="center", va="bottom", fontsize=8)

    # --- Right: per-quarter reward trajectory (heuristic) ---
    for r in summary["heuristic"]["runs"]:
        qr = r["quarterly_rewards"]
        axes[1].plot(range(1, len(qr) + 1), qr, marker="o", alpha=0.5,
                     label=f"seed {r['seed']}")
    # mean trajectory
    maxq = max(len(r["quarterly_rewards"]) for r in summary["heuristic"]["runs"])
    meanq = []
    for i in range(maxq):
        vals = [r["quarterly_rewards"][i] for r in summary["heuristic"]["runs"]
                if i < len(r["quarterly_rewards"])]
        meanq.append(mean(vals))
    axes[1].plot(range(1, maxq + 1), meanq, color="black", linewidth=2.5,
                 marker="s", label="mean")
    axes[1].axhline(0, color="gray", linestyle="--", linewidth=1)
    axes[1].set_xlabel("Quarter")
    axes[1].set_ylabel("Sparse quarterly reward")
    axes[1].set_title("Heuristic quarterly reward (delayed, sparse) — v2")
    axes[1].grid(True, alpha=0.3)
    axes[1].legend(fontsize=8, ncol=2)

    fig.suptitle("HCM:21 OpenEnv — long-horizon HR planning benchmark",
                 fontsize=13, fontweight="bold")
    fig.tight_layout(rect=[0, 0, 1, 0.96])
    fig.savefig("poster_figure.png", dpi=150)


def make_grpo_curve(reward_curve: list) -> None:
    """Figure 5: GRPO training reward over logging steps (from the Colab run)."""
    import matplotlib
    matplotlib.use("Agg")
    import matplotlib.pyplot as plt

    y = [float(r) for r in reward_curve]
    x = list(range(1, len(y) + 1))
    fig, ax = plt.subplots(figsize=(7, 4.5))
    ax.plot(x, y, color="#27ae60", linewidth=1.2, alpha=0.5, label="reward")
    # simple trailing moving average to show the trend
    win = max(1, len(y) // 20)
    if win > 1:
        ma = [mean(y[max(0, i - win + 1):i + 1]) for i in range(len(y))]
        ax.plot(x, ma, color="#145a32", linewidth=2.5, label=f"moving avg ({win})")
    ax.axhline(0, color="gray", linestyle="--", linewidth=1)
    ax.set_xlabel("Training step")
    ax.set_ylabel("Mean group reward")
    ax.set_title("HCM:21 — GRPO training reward (Qwen3-0.6B, health-gated)")
    ax.grid(True, alpha=0.3)
    ax.legend(fontsize=9)
    fig.tight_layout()
    fig.savefig("fig5_grpo_reward_curve.png", dpi=150)
    plt.close(fig)


if __name__ == "__main__":
    main()