from __future__ import annotations

from pathlib import Path

import matplotlib

matplotlib.use("Agg")

import matplotlib.pyplot as plt

from engine.controller import Action, ControllerConfig, WhenToSpeakController
from engine.traces import SignalTrace, surprising_claim_trace


ROOT = Path(__file__).resolve().parents[1]
DEFAULT_OUTPUT = ROOT / "eval" / "controller_timeline.png"
ACTION_COLORS = {
    Action.BACKCHANNEL: "tab:blue",
    Action.TAKE_FLOOR: "tab:green",
    Action.INTERRUPT: "tab:red",
}


def render_timeline(
    trace: SignalTrace | None = None,
    output_path: str | Path = DEFAULT_OUTPUT,
    config: ControllerConfig | None = None,
) -> Path:
    trace = trace or surprising_claim_trace()
    controller = WhenToSpeakController(trace.agent_ids, config=config)
    ticks = [controller.tick(step.signals_by_agent, floor_holder=step.floor_holder) for step in trace.steps]

    steps = [tick.step for tick in ticks]
    primary = trace.agent_ids[0]
    surprise = [step.signals_by_agent[primary].surprise for step in trace.steps]
    readiness = [step.signals_by_agent[primary].readiness for step in trace.steps]
    p_end = [step.signals_by_agent[primary].p_end for step in trace.steps]
    change = [tick.decisions[primary].change_score for tick in ticks]

    fig, axes = plt.subplots(4, 1, figsize=(10, 8), sharex=True)
    fig.suptitle(f"WhenToSpeak controller timeline: {trace.name}", fontsize=13)

    axes[0].plot(steps, surprise, marker="o", color="tab:orange", label="surprise/NLL")
    axes[0].set_ylabel("surprise")
    axes[0].legend(loc="upper right")

    axes[1].plot(steps, change, marker="o", color="tab:purple", label="change score")
    axes[1].set_ylabel("change")
    axes[1].legend(loc="upper right")

    axes[2].plot(steps, readiness, marker="o", color="tab:green", label="readiness")
    axes[2].plot(steps, p_end, marker=".", linestyle="--", color="tab:gray", label="p_end")
    axes[2].set_ylabel("probability")
    axes[2].set_ylim(-0.05, 1.05)
    axes[2].legend(loc="upper right")

    for agent_id in trace.agent_ids:
        urges = [tick.decisions[agent_id].urge for tick in ticks]
        axes[3].plot(steps, urges, marker="o", label=f"{agent_id} urge")
    axes[3].axhline(controller.config.tau, color="black", linestyle="--", linewidth=1, label="tau")
    axes[3].set_ylabel("urge")
    axes[3].set_xlabel("step")
    axes[3].legend(loc="upper right")

    for tick in ticks:
        for agent_id, decision in tick.decisions.items():
            if decision.action == Action.SILENT:
                continue
            color = ACTION_COLORS[decision.action]
            axes[3].scatter([tick.step], [decision.urge], color=color, s=90, zorder=5)
            axes[3].annotate(
                decision.action.value,
                (tick.step, decision.urge),
                textcoords="offset points",
                xytext=(0, 8),
                ha="center",
                fontsize=8,
                color=color,
            )

    for axis in axes:
        axis.grid(True, alpha=0.25)

    output = Path(output_path)
    output.parent.mkdir(parents=True, exist_ok=True)
    fig.tight_layout()
    fig.savefig(output, dpi=160)
    plt.close(fig)
    return output


def main() -> None:
    output = render_timeline()
    print(f"Wrote {output}")


if __name__ == "__main__":
    main()