"""
Qiskit circuit diagrams for each QAOA module and each construction stage.

Stages
------
"init"  : Hadamard initialization only.
"cost"  : Init + Cost Layer (Rz / Rzz gates).
"full"  : Init + Cost Layer + Mixer Layer + measurement barrier.

Modes
-----
"standard"   : 3-qubit unconstrained (sparse connectivity).
"equality"   : 3-qubit equality penalty (fully connected K3).
"inequality" : 4-qubit slack variable (fully connected K4, q3 labelled s).

All functions return a Matplotlib Figure.
"""

import matplotlib
matplotlib.use("Agg")   # non-interactive backend for Streamlit
import matplotlib.pyplot as plt
from qiskit import QuantumCircuit
from qiskit.circuit import Parameter


_GAMMA = Parameter("γ")
_BETA  = Parameter("β")


def _make_circuit(stage: str, mode: str) -> QuantumCircuit:
    """Build the symbolic QuantumCircuit for the requested stage and mode."""
    n = 4 if mode == "inequality" else 3

    # Named qubit registers so the slack qubit shows as 's'
    if mode == "inequality":
        from qiskit.circuit import QuantumRegister
        x_reg = QuantumRegister(3, "x")
        s_reg = QuantumRegister(1, "s")
        qc = QuantumCircuit(x_reg, s_reg)
        all_qubits = list(range(4))
    else:
        qc = QuantumCircuit(n)
        all_qubits = list(range(n))

    # --- Initialization ---
    qc.h(all_qubits)

    if stage == "init":
        return qc

    # --- Cost Layer ---
    qc.barrier(label="Cost  U_C(γ)")

    if mode == "standard":
        # Linear terms: Z1, Z2
        qc.rz(2 * _GAMMA, 0)
        qc.rz(2 * _GAMMA, 1)
        # Quadratic term: Z2 Z3
        qc.rzz(2 * _GAMMA, 1, 2)

    elif mode == "equality":
        # Linear terms: Z1, Z2, Z3 (from both obj and penalty expansion)
        qc.rz(2 * _GAMMA, [0, 1, 2])
        # Quadratic terms: all pairs (K3 fully connected)
        qc.rzz(2 * _GAMMA, 0, 1)
        qc.rzz(2 * _GAMMA, 1, 2)
        qc.rzz(2 * _GAMMA, 0, 2)

    elif mode == "inequality":
        # Linear terms: x1, x2, x3, s
        qc.rz(2 * _GAMMA, all_qubits)
        # Quadratic terms: all pairs (K4 fully connected)
        for i in range(4):
            for j in range(i + 1, 4):
                qc.rzz(2 * _GAMMA, i, j)

    if stage == "cost":
        return qc

    # --- Mixer Layer ---
    qc.barrier(label="Mixer  U_B(β)")
    qc.rx(2 * _BETA, all_qubits)

    qc.barrier(label="Measure")
    qc.measure_all()

    return qc


def draw_circuit(stage: str, mode: str, scale: float = 0.99) -> plt.Figure:
    """
    Return a Matplotlib Figure of the QAOA circuit at the given stage/mode.

    Parameters
    ----------
    stage : "init" | "cost" | "full"
    mode  : "standard" | "equality" | "inequality"
    scale : float
        Passed to qc.draw(..., scale=scale) to control figure size.
    """
    qc = _make_circuit(stage, mode)
    fig = qc.draw("mpl", style="iqp", scale=scale, fold=-1)
    plt.close("all")   # prevent memory leak; fig is already captured
    return fig