src/pages/equality.py

"""
Equality-Constrained QAOA Module — 3 qubits.

Problem:  Minimize  C(x) = x1 + x2 - x2*x3
          Subject to:  x1 + x2 + x3 = 1
Penalty:  P(x) = A * (x1 + x2 + x3 - 1)^2,  A = 3.0
Total:    C_total(x) = C(x) + P(x)
Target:   |001⟩  (only valid state with lowest objective cost)
"""

import numpy as np
import pandas as pd
import plotly.graph_objects as go
import streamlit as st

from core.costs import cost_equality, cost_standard
from core.simulator import simulate_qaoa, get_statevector_after_init, get_statevector_after_cost
from core.optimizer import run_cobyla
from viz.disc_viz import show_disc
from viz.probabilities import plot_probabilities
from viz.circuits import draw_circuit

_N = 3
_MODE = "equality"
_A = 3.0


def _truth_table() -> pd.DataFrame:
    rows = []
    for i in range(2 ** _N):
        bits = [int(b) for b in format(i, f"0{_N}b")]
        c_obj = cost_standard(*bits)
        constraint_sum = sum(bits)
        penalty = _A * (constraint_sum - 1) ** 2
        c_total = c_obj + penalty
        feasible = "✅" if constraint_sum == 1 else "❌"
        rows.append({
            "State": f"|{format(i, f'0{_N}b')}⟩",
            "x₁+x₂+x₃": constraint_sum,
            "Feasible?": feasible,
            "C_obj": c_obj,
            "Penalty": penalty,
            "C_total": c_total,
        })
    return pd.DataFrame(rows)


def _hamiltonian_df():
    costs = [cost_equality(*[int(b) for b in format(i, f"0{_N}b")]) for i in range(2 ** _N)]
    df = pd.DataFrame(
        np.diag(costs),
        columns=[f"|{format(i, f'0{_N}b')}⟩" for i in range(2 ** _N)],
        index=[f"⟨{format(i, f'0{_N}b')}|" for i in range(2 ** _N)],
    )
    return df, costs


def render():
    # ── Back button ─────────────────────────────────────────────────────────
    if st.button("← Back to Home"):
        st.session_state["page"] = "home"
        st.rerun()

    st.markdown('<h1 style="color:#7C3AED;">Constrained QAOA — Equality Penalty</h1>',
                unsafe_allow_html=True)
    st.caption("3-qubit · penalty method · fully-connected K₃ cost layer · sv_disc visualization")

    # ════════════════════════════════════════════════════════════════════════
    # STEP 1 — Problem Formulation
    # ════════════════════════════════════════════════════════════════════════
    st.markdown("---")
    st.markdown("## Step 1 · Problem Formulation")
    col1, col2 = st.columns(2)
    with col1:
        st.info(r"**Minimize** $C(x) = x_1 + x_2 - x_2 x_3$")
        st.error(r"**Subject to** $x_1 + x_2 + x_3 = 1$")
        st.markdown(
            "Only one qubit may be $|1\\rangle$ at a time. "
            "Among the feasible states, $|001\\rangle$ achieves the lowest objective value."
        )
    with col2:
        st.markdown("**Feasibility & Cost Table**")
        st.dataframe(_truth_table(), hide_index=True)

    # ════════════════════════════════════════════════════════════════════════
    # STEP 2 — Penalty Method & QUBO
    # ════════════════════════════════════════════════════════════════════════
    st.markdown("---")
    st.markdown("## Step 2 · Penalty Method & QUBO")
    st.write(
        "We handle the equality constraint by adding a **quadratic penalty** "
        "to the objective. Infeasible states receive a large cost that the optimizer avoids."
    )
    st.latex(r"""
C_{\text{total}}(x) = C(x) + A\,(x_1 + x_2 + x_3 - 1)^2, \quad A = 3.0
""")

    with st.expander("How do we choose A = 3.0?"):
        st.markdown(
            "The penalty weight $A$ must be **large enough** that the penalized cost of any "
            "infeasible state exceeds the cost of the best feasible state. "
            "Here, the best feasible cost is $0$ and the smallest infeasible penalty is "
            "$3.0 \\times 1 = 3.0$, which dominates. "
            "Too small an $A$ and the optimizer ignores the constraint; "
            "too large and the landscape becomes steep, slowing convergence."
        )

    st.markdown("#### Expanded QUBO")
    st.write(
        "Expanding $(x_1 + x_2 + x_3 - 1)^2$ introduces **cross terms** between all pairs of "
        "qubits — these become the $R_{zz}$ gates in the cost circuit:"
    )
    st.latex(r"""
(x_1+x_2+x_3-1)^2 = x_1^2 + x_2^2 + x_3^2 + 2x_1x_2 + 2x_1x_3 + 2x_2x_3 - 2x_1 - 2x_2 - 2x_3 + 1
""")
    st.latex(r"""
\xrightarrow{x_i^2 = x_i} \;
x_1 + x_2 + x_3 + 2x_1x_2 + 2x_1x_3 + 2x_2x_3 - 2x_1 - 2x_2 - 2x_3 + 1
""")
    st.info("Every pair $(x_i, x_j)$ now appears — the QUBO matrix is **fully dense** (K₃ topology).")

    # ════════════════════════════════════════════════════════════════════════
    # STEP 3 — Hamiltonian
    # ════════════════════════════════════════════════════════════════════════
    st.markdown("---")
    st.markdown("## Step 3 · Penalized Hamiltonian Matrix")
    st.write("The diagonal of $H_C$ now reflects the total penalized cost:")
    diag_df, costs = _hamiltonian_df()
    min_c = min(costs)
    styled = diag_df.style.format("{:.1f}").map(
        lambda v: "background-color:#ede9fe;color:#4c1d95;font-weight:bold;" if np.isclose(v, min_c)
        else ("background-color:#fee2e2;color:#b91c1c;font-weight:bold;" if v > 5
              else ("color:black;" if v != 0 else "color:#94a3b8;"))
    )
    st.dataframe(styled, height=310)

    with st.expander("Worked example: why does |111⟩ cost 13?"):
        st.latex(r"""
C_{\text{total}}(1,1,1) = \underbrace{1+1-1}_{C_{\text{obj}}=1}
+ 3.0\,(1+1+1-1)^2 = 1 + 3.0 \times 4 = 13.0
""")

    # ════════════════════════════════════════════════════════════════════════
    # STEP 4 — Circuit Construction + Phase-Probability Disconnect
    # ════════════════════════════════════════════════════════════════════════
    st.markdown("---")
    st.markdown("## Step 4 · Circuit Construction & Phase-Probability Disconnect")

    st.markdown("#### Connectivity vs. Standard QAOA")
    st.markdown(
        "| Module | $R_z$ gates | $R_{zz}$ gates | Topology |\n"
        "|--------|------------|----------------|----------|\n"
        "| Standard | 2 | 1 | Sparse (1 edge) |\n"
        "| **Equality** | **3** | **3** | **K₃ fully connected** |"
    )
    st.write(
        "Each $R_{zz}$ corresponds to a cross term in the penalty expansion. "
        "More entangling gates = higher hardware cost = pedagogical contrast."
    )

    st.markdown("#### Stage A — Initialization")
    _, c1, c2 = st.columns([1,1, 2])
    with c1:
        fig_init = draw_circuit("init", _MODE)
        st.pyplot(fig_init, width=300)
    with c2:
        show_disc(get_statevector_after_init(_N),
                  caption="Equal superposition: all radii identical, all phases zero.",
                  phase_colors=True)

    st.markdown("#### Stage B — Cost Layer (Phases Rotate, Probabilities Flat)")
    st.warning(
        "The penalty Hamiltonian has **larger eigenvalue spread** than the standard one, "
        "causing faster phase rotation for high-cost states like $|111\\rangle$. "
        "Yet the probability bar chart remains **perfectly flat** — pure phase encoding."
    )
    c3, c4 = st.columns([1, 1])
    with c3:
        st.pyplot(draw_circuit("cost", _MODE), width="stretch")
        g_cost = st.slider("γ (Cost Layer demo)", 0.0, 6.28, 0.6, 0.05, key="eq_g_cost")
    with c4:
        sv_cost = get_statevector_after_cost(g_cost, cost_equality, _N)
        show_disc(sv_cost,
                  caption="Large penalty → fast phase rotation for high-cost states.",
                  phase_colors=True)
        probs_cost = np.array([np.abs(a) ** 2 for a in sv_cost.data])
        _, _, cv_cost, _ = simulate_qaoa(g_cost, 0.0, cost_equality, _N, shots=1, seed=0)
        st.plotly_chart(
            plot_probabilities(probs_cost, cv_cost, _N,
                               title="Probabilities (flat — phase only)", shots=1024),
            width="stretch"
        )

    st.markdown("#### Stage C — Mixer Layer (Selective Amplitude Amplification)")
    st.success(
        "With high-penalty phases on infeasible states, the mixer's interference **suppresses** "
        "their amplitudes and **amplifies** feasible states. "
        "QAOA effectively learns to 'steer' probability mass toward the constraint-satisfying region."
    )
    c5, c6 = st.columns([1, 1])
    with c5:
        st.pyplot(draw_circuit("full", _MODE), width="stretch")
        g_mix = st.slider("γ", 0.0, 6.28, 0.6, 0.05, key="eq_g_mix")
        b_mix = st.slider("β", 0.0, 3.14, 0.5, 0.05, key="eq_b_mix")
    with c6:
        probs_mix, _, cv_mix, sv_mix = simulate_qaoa(g_mix, b_mix, cost_equality, _N,
                                                      shots=1024, seed=200)
        show_disc(sv_mix, caption="After mixer: feasible states gain amplitude.", phase_colors=True)
        st.plotly_chart(
            plot_probabilities(probs_mix, cv_mix, _N,
                               title="Probabilities (feasible states amplified)", shots=1024),
            width="stretch"
        )

    # ════════════════════════════════════════════════════════════════════════
    # STEP 5 — Interactive Simulation
    # ════════════════════════════════════════════════════════════════════════
    st.markdown("---")
    st.markdown("## Step 5 · Interactive Simulation")
    col_ctrl, col_vis = st.columns([1, 2])
    with col_ctrl:
        st.subheader("Parameters")
        gamma5 = st.slider("γ", 0.0, 6.28, 0.6, 0.05, key="eq_gamma5")
        beta5  = st.slider("β", 0.0, 3.14, 0.5, 0.05, key="eq_beta5")
        shots5 = st.select_slider("Shots", [128, 512, 1024, 4096], value=1024, key="eq_shots5")
        probs5, energy5, cv5, sv5 = simulate_qaoa(gamma5, beta5, cost_equality,
                                                   _N, shots=shots5, seed=42)
        st.metric("Expected Total Cost ⟨C_total⟩", f"{energy5:.4f}")
    with col_vis:
        st.plotly_chart(plot_probabilities(probs5, cv5, _N, shots=shots5),
                        width="stretch")
        show_disc(sv5, caption="Current statevector.", phase_colors=True)

    # ════════════════════════════════════════════════════════════════════════
    # STEP 6 — COBYLA Optimization + Replay
    # ════════════════════════════════════════════════════════════════════════
    st.markdown("---")
    st.markdown("## Step 6 · COBYLA Optimization & Replay")

    if st.button("▶ Run COBYLA Optimizer", key="eq_run_opt"):
        with st.spinner("Optimizing… (up to 60 iterations)"):
            res, hist = run_cobyla(cost_equality, _N, mode="equality")
            st.session_state["eq_history"] = hist
            st.session_state["eq_opt_done"] = True
        st.success(f"Done — {len(hist['energy'])} calls · final ⟨C_total⟩ = {hist['energy'][-1]:.4f}")

    if st.session_state.get("eq_opt_done"):
        hist = st.session_state["eq_history"]

        fig_conv = go.Figure()
        fig_conv.add_trace(go.Scatter(
            x=list(range(len(hist["energy"]))), y=hist["energy"],
            mode="lines+markers", line=dict(color="#7C3AED", width=2)
        ))
        fig_conv.update_layout(
            title="COBYLA Convergence", xaxis_title="Iteration",
            yaxis_title="⟨C_total⟩",
            paper_bgcolor="rgba(0,0,0,0)", plot_bgcolor="rgba(245,243,255,0.5)",
            height=300, margin=dict(l=20, r=20, t=40, b=20)
        )
        st.plotly_chart(fig_conv, width="stretch")

        g_fin = hist["gamma"][-1]; b_fin = hist["beta"][-1]
        pf, _, _, _ = simulate_qaoa(g_fin, b_fin, cost_equality, _N, shots=8192, seed=999)
        winner = np.argmax(pf)
        feasible_check = "✅ Feasible" if sum(int(b) for b in format(winner, f"0{_N}b")) == 1 else "⚠ Infeasible"
        st.markdown(
            f'<div style="background:#f5f3ff;padding:18px;border:2px solid #7c3aed;'
            f'border-radius:10px;text-align:center;">'
            f'<div style="font-size:1.1rem;color:#4c1d95;">Most Probable State</div>'
            f'<div style="font-size:2.2rem;font-weight:800;font-family:monospace;color:#6d28d9;">'
            f'|{format(winner, f"0{_N}b")}⟩</div>'
            f'<div style="color:#7c3aed;">{feasible_check} · Probability: {pf[winner]:.2%} · '
            f'γ* = {g_fin:.3f} · β* = {b_fin:.3f}</div></div>',
            unsafe_allow_html=True
        )

        st.markdown("### Replay Optimization")
        step = st.slider("Iteration", 0, len(hist["energy"]) - 1, 0, key="eq_replay")
        g_s = hist["gamma"][step]; b_s = hist["beta"][step]
        p_s, _, cv_s, sv_s = simulate_qaoa(g_s, b_s, cost_equality, _N, shots=2048, seed=step)
        col_r1, col_r2 = st.columns(2)
        with col_r1:
            st.plotly_chart(
                plot_probabilities(p_s, cv_s, _N,
                                   title=f"Iteration {step}", shots=2048),
                width="stretch"
            )
        with col_r2:
            show_disc(sv_s,
                      caption=f"Statevector at iteration {step}  (γ={g_s:.3f}, β={b_s:.3f})",
                      phase_colors=True)