"""Plot per-category WCXB dev F1 for deformat's baseline and filter pipeline.

Run:
    uv run --locked gen_category_f1.py

Output:
    category_f1.png
"""

# /// script
# requires-python = "==3.12.*"
# dependencies = [
#   "matplotlib>=3.9",
# ]
# ///

from __future__ import annotations

import json
import sys
from pathlib import Path

import matplotlib.pyplot as plt

sys.path.insert(0, str(Path(__file__).resolve().parents[2] / "posts"))
from figure_quality import configure_matplotlib, figure_output_path, validate_figure

configure_matplotlib()

HERE = Path(__file__).resolve().parent
COLORS = {"strip": "#6c757d", "triple": "#b55239"}


def load_results() -> dict[str, object]:
    return json.loads((HERE / "wcxb-dev-results.json").read_text())


def plot() -> None:
    packet = load_results()
    results = packet["results"]
    categories = [
        "article",
        "collection",
        "documentation",
        "forum",
        "listing",
        "product",
        "service",
    ]
    labels = [
        f"{category.title()}  (n={results['strip'][category]['n']})"
        for category in categories
    ]
    baseline = [results["strip"][category]["f1"] for category in categories]
    filtered = [results["triple"][category]["f1"] for category in categories]
    y = list(range(len(categories)))

    figure, axis = plt.subplots(figsize=(9.2, 5.3))
    for row, before, after in zip(y, baseline, filtered, strict=True):
        color = "#8f3b2d" if after < before else "#b8aa94"
        axis.plot([before, after], [row, row], color=color, linewidth=2.2, zorder=1)
    axis.scatter(
        baseline,
        y,
        color=COLORS["strip"],
        s=64,
        label="baseline strip",
        zorder=2,
    )
    axis.scatter(
        filtered,
        y,
        color=COLORS["triple"],
        marker="D",
        s=58,
        label="three-filter pipeline",
        zorder=3,
    )

    docs = categories.index("documentation")
    axis.annotate(
        "documentation: -0.007",
        xy=(filtered[docs], docs),
        xytext=(0.72, docs + 0.72),
        arrowprops={"arrowstyle": "->", "color": "#5f3128"},
        color="#5f3128",
    )
    axis.set_yticks(y, labels)
    axis.invert_yaxis()
    axis.set_xlim(0.35, 0.95)
    axis.set_xlabel("mean per-page word F1 within category")
    axis.set_title("One aggregate improvement, one category reversal")
    axis.grid(axis="x", alpha=0.2)
    axis.spines[["top", "right", "left"]].set_visible(False)
    axis.legend(
        loc="upper center",
        bbox_to_anchor=(0.5, -0.13),
        frameon=False,
        ncol=2,
    )
    figure.tight_layout()
    validate_figure(figure)
    output = figure_output_path(__file__, "category_f1.png")
    figure.savefig(output, dpi=300, bbox_inches="tight")


if __name__ == "__main__":
    plot()