"""Generate Fig. 7 from NPA data."""
from __future__ import annotations

import json
import sys
from pathlib import Path

import matplotlib.pyplot as plt
import numpy as np

sys.path.insert(0, str(Path(__file__).parent.parent / "src"))
from npa_solver import WORDS, build_moment_matrix, reverse_word, word_to_canonical


def load_data(json_path: Path):
    with open(json_path, "r") as f:
        return json.load(f)


def plot_fig7(data_dir: Path, out_dir: Path):
    json_path = data_dir / "tilted_chsh_npa_level2.json"
    manifest = load_data(json_path)

    fig, ax = plt.subplots(figsize=(4.8, 3.6))
    colors = ["#1f77b4", "#2ca02c", "#ff7f0e", "#9467bd"]

    for idx, (alpha_str, records) in enumerate(manifest["data"].items()):
        alpha = float(alpha_str)
        s12 = np.array([r["authorized_score"] for r in records])
        s13 = np.array([r["collusive_score"] if r["collusive_score"] is not None else np.nan for r in records])
        certified = np.array([bool(r.get("certified", False)) for r in records])

        ax.plot(s12, s13, color=colors[idx], linewidth=1.2, alpha=0.6)
        cert_mask = certified
        if np.any(cert_mask):
            ax.plot(s12[cert_mask], s13[cert_mask], "o", color=colors[idx], markersize=3.5, label=rf"$\alpha={alpha}$")
        noncert_mask = ~cert_mask
        if np.any(noncert_mask):
            ax.plot(s12[noncert_mask], s13[noncert_mask], "o", markerfacecolor="none", markeredgecolor=colors[idx], markersize=3.5, alpha=0.5)

        if alpha == 0.0:
            s_ref = np.linspace(0.0, np.sqrt(8.0), 400)
            s13_ref = np.sqrt(np.maximum(0.0, 8.0 - s_ref**2))
            ax.plot(s_ref, s13_ref, color=colors[idx], linestyle=":", linewidth=1.8, label=r"analytic CHSH ($\alpha=0$)")

    s_diag = np.linspace(0, 4.0, 200)
    ax.plot(s_diag, s_diag, color="#d62728", linestyle="--", linewidth=1.5, label="copied seed")
    ax.axhline(2.0, color="gray", linestyle=":", linewidth=1)
    ax.set_xlim(0, 4.5)
    ax.set_ylim(0, 4.5)
    ax.set_xlabel(r"authorized tilted-CHSH score $I_\alpha^{12}$")
    ax.set_ylabel(r"collusive tilted-CHSH score $I_\alpha^{13}$")
    ax.set_title("NPA upper envelopes beyond CHSH")
    ax.legend(loc="upper right", frameon=False, fontsize=8)
    ax.grid(alpha=0.2)

    out_dir.mkdir(parents=True, exist_ok=True)
    fig.savefig(out_dir / "fig7_npa_upper_envelopes.pdf")
    fig.savefig(out_dir / "fig7_npa_upper_envelopes.png", dpi=300)
    plt.close(fig)
    print(f"Saved figures to {out_dir}")


if __name__ == "__main__":
    data_dir = Path(__file__).parent.parent / "data" / "raw"
    out_dir = Path(__file__).parent.parent / "figures"
    plot_fig7(data_dir, out_dir)