src/generate_figures.py

"""
Generate high-resolution journal-quality figures for the IEEE QSEEC paper.

Run from the qaoa_pedagogical_tool/ directory:
    python generate_figures.py

Output: paper_figures/
  fig1_circuit_standard.pdf   - Standard 3-qubit QAOA circuit (full)
  fig2_circuit_equality.pdf   - Equality-constrained circuit (full)
  fig3_circuit_inequality.pdf - Inequality / slack-variable circuit (full)
  fig4_topology.pdf           - NetworkX qubit-interaction topologies (3 panels)
  fig5_phase_disconnect.pdf   - sv_disc at 3 stages (init / cost / mixer)
  fig6_convergence.pdf        - COBYLA convergence + final probability bar chart
"""

import sys, os
import numpy as np
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import matplotlib.gridspec as gridspec
import networkx as nx

# Make project root importable
_ROOT = os.path.dirname(os.path.abspath(__file__))
if _ROOT not in sys.path:
    sys.path.insert(0, _ROOT)

from core.costs import cost_standard, cost_equality, cost_inequality
from core.simulator import simulate_qaoa, get_statevector_after_init, get_statevector_after_cost
from core.optimizer import run_cobyla
from viz.circuits import draw_circuit
from sv_disc import sv_disc

_OUT = os.path.join(_ROOT, "paper_figures")
os.makedirs(_OUT, exist_ok=True)

DPI = 300
_FONT = {"family": "serif", "size": 9}
matplotlib.rc("font", **_FONT)
matplotlib.rc("text", usetex=False)   # set True if LaTeX is available on system


# ─────────────────────────────────────────────────────────────────────────────
# FIG 1-3  QAOA Circuit Diagrams
# ─────────────────────────────────────────────────────────────────────────────
def fig_circuits():
    for mode, fname in [
        ("standard",    "fig1_circuit_standard.pdf"),
        ("equality",    "fig2_circuit_equality.pdf"),
        ("inequality",  "fig3_circuit_inequality.pdf"),
    ]:
        fig = draw_circuit("full", mode, scale=0.75)
        fig.savefig(os.path.join(_OUT, fname), dpi=DPI, bbox_inches="tight")
        plt.close(fig)
        print(f"  saved {fname}")


# ─────────────────────────────────────────────────────────────────────────────
# FIG 4  Qubit Interaction Topology (NetworkX)
# ─────────────────────────────────────────────────────────────────────────────
def fig_topology():
    fig, axes = plt.subplots(1, 3, figsize=(7, 2.4))
    titles = ["Standard\n(Sparse)", "Equality\n(K₃ full)", "Inequality\n(K₄ + slack)"]
    node_labels = [
        {0: "x₁", 1: "x₂", 2: "x₃"},
        {0: "x₁", 1: "x₂", 2: "x₃"},
        {0: "x₁", 1: "x₂", 2: "x₃", 3: "s"},
    ]
    edges = [
        [(1, 2)],                                           # standard: only Z2Z3
        [(0, 1), (1, 2), (0, 2)],                          # equality: K3
        [(0,1),(0,2),(0,3),(1,2),(1,3),(2,3)],              # inequality: K4
    ]
    colors = ["#6366F1", "#7C3AED", "#0F766E"]
    layouts = [
        {0: (0, 0), 1: (1, 0), 2: (2, 0)},                # line
        {0: (0, 0), 1: (1, 1), 2: (2, 0)},                # triangle
        {0: (0, 1), 1: (1, 2), 2: (2, 1), 3: (1, 0)},     # diamond
    ]

    for ax, title, labels, edg, col, pos in zip(axes, titles, node_labels, edges, colors, layouts):
        n = len(labels)
        G = nx.Graph()
        G.add_nodes_from(range(n))
        G.add_edges_from(edg)
        nx.draw(G, pos=pos, ax=ax, with_labels=True, labels=labels,
                node_color=col, node_size=600, font_color="white",
                font_size=8, font_weight="bold",
                edge_color="#334155", width=2)
        ax.set_title(title, fontsize=9, fontweight="bold", pad=6)

    fig.suptitle("Qubit Interaction Topologies (Cost Layer)", fontsize=10, fontweight="bold", y=1.02)
    fig.tight_layout()
    fname = "fig4_topology.pdf"
    fig.savefig(os.path.join(_OUT, fname), dpi=DPI, bbox_inches="tight")
    plt.close(fig)
    print(f"  saved {fname}")


# ─────────────────────────────────────────────────────────────────────────────
# FIG 5  Phase–Probability Disconnect (sv_disc, 3 stages)
# ─────────────────────────────────────────────────────────────────────────────
def fig_phase_disconnect():
    GAMMA = 3.14
    BETA  = 0.5

    sv_init  = get_statevector_after_init(3)
    sv_cost  = get_statevector_after_cost(GAMMA, cost_standard, 3)
    _, _, _, sv_full = simulate_qaoa(GAMMA, BETA, cost_standard, 3, shots=1024, seed=42)

    stages = [
        (sv_init, "Stage A: Init\n(H gates only)"),
        (sv_cost, f"Stage B: Cost Layer\n(γ = {GAMMA:.2f})"),
        (sv_full, f"Stage C: After Mixer\n(γ = {GAMMA:.2f}, β = {BETA:.2f})"),
    ]

    fig = plt.figure(figsize=(10, 3.2))
    outer = gridspec.GridSpec(1, 3, figure=fig, wspace=0.05)

    for col_idx, (sv, title) in enumerate(stages):
        # draw sv_disc into a sub-figure
        sub_fig = sv_disc(sv, show_labels=True, phase_colors=True, num_columns=4)
        sub_axes = sub_fig.get_axes()
        inner = gridspec.GridSpecFromSubplotSpec(
            len(sub_axes) // 4, 4, subplot_spec=outer[col_idx], wspace=0, hspace=0
        )
        # Copy patches and lines from sv_disc sub-figure into main figure axes
        for i, ax_src in enumerate(sub_axes):
            ax_dst = fig.add_subplot(inner[i // 4, i % 4])
            ax_dst.set_xlim(ax_src.get_xlim())
            ax_dst.set_ylim(ax_src.get_ylim())
            ax_dst.set_aspect("equal")
            ax_dst.set_axis_off()
            for patch in ax_src.patches:
                import copy
                ax_dst.add_patch(copy.copy(patch))
            for line in ax_src.lines:
                ax_dst.add_line(copy.copy(line))
            for txt in ax_src.texts:
                ax_dst.text(*txt.get_position(), txt.get_text(),
                            ha=txt.get_ha(), va=txt.get_va(),
                            fontsize=txt.get_fontsize())
        plt.close(sub_fig)
        # Column title
        fig.text(
            (col_idx + 0.5) / 3, 1.02, title,
            ha="center", va="bottom", fontsize=9, fontweight="bold",
            transform=fig.transFigure,
        )

    fig.suptitle("Phase–Probability Disconnect in QAOA (Standard, 3 Qubits)",
                 fontsize=10, fontweight="bold", y=1.08)
    fname = "fig5_phase_disconnect.pdf"
    fig.savefig(os.path.join(_OUT, fname), dpi=DPI, bbox_inches="tight")
    plt.close(fig)
    print(f"  saved {fname}")


# ─────────────────────────────────────────────────────────────────────────────
# FIG 6  COBYLA Convergence + Final Probability Bar Chart
# ─────────────────────────────────────────────────────────────────────────────
def fig_convergence():
    print("  running COBYLA for standard module (fig6)…")
    _, hist = run_cobyla(cost_standard, 3, mode="standard")

    g_fin = hist["gamma"][-1]; b_fin = hist["beta"][-1]
    probs_fin, _, cv_fin, _ = simulate_qaoa(g_fin, b_fin, cost_standard, 3,
                                             shots=8192, seed=999)

    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(7.5, 2.8))

    # Convergence curve
    ax1.plot(hist["energy"], color="#4F46E5", linewidth=1.8, marker="o",
             markersize=3, label="⟨C⟩")
    ax1.set_xlabel("COBYLA Iteration", fontsize=9)
    ax1.set_ylabel("Expected Energy ⟨C⟩", fontsize=9)
    ax1.set_title("Variational Convergence", fontsize=9, fontweight="bold")
    ax1.grid(True, linestyle="--", alpha=0.4)
    ax1.legend(fontsize=8)

    # Probability bar chart
    n = 3
    states = [f"|{format(i, f'0{n}b')}⟩" for i in range(2 ** n)]
    min_c = float(np.min(cv_fin))
    bar_colors = ["#4F46E5" if np.isclose(c, min_c) else "#CBD5E1" for c in cv_fin]
    ax2.bar(states, probs_fin, color=bar_colors, edgecolor="white", linewidth=0.5)
    ax2.set_xlabel("Basis State", fontsize=9)
    ax2.set_ylabel("Probability", fontsize=9)
    ax2.set_title("Final Measurement Distribution", fontsize=9, fontweight="bold")
    ax2.set_ylim(0, 1)
    ax2.tick_params(axis="x", rotation=45, labelsize=7)
    ax2.grid(axis="y", linestyle="--", alpha=0.4)

    fig.tight_layout()
    fname = "fig6_convergence.pdf"
    fig.savefig(os.path.join(_OUT, fname), dpi=DPI, bbox_inches="tight")
    plt.close(fig)
    print(f"  saved {fname}")


# ─────────────────────────────────────────────────────────────────────────────
if __name__ == "__main__":
    print("Generating paper figures → paper_figures/")
    fig_circuits()
    fig_topology()
    fig_phase_disconnect()
    fig_convergence()
    print("All figures saved.")