"""Quantum Suite — a real statevector quantum simulator in ONE cpu-basic Space.

Consolidates the 11-member `QUANTUM_*` family, none of which can run on a free
CPU Space: every one of them imports `qiskit`, `torch` and `transformers` and pulls
real models (GPT-2, facebook/bart-large-mnli, the Olivetti faces set) at import time.

This suite replaces them with the physics that actually matters and nothing else:
a hand-written complex-amplitude simulator (pure Python + numpy), so the whole
family fits in one slot.

Tabs
----
  Circuit Simulator    a text circuit language, amplitudes and probabilities out
  Superposition Explorer  the uniform superposition and its 1/sqrt(N) statistics
  Grover Search        real amplitude amplification, iteration by iteration
  Entanglement Lab     Bell/GHZ states, entropy, and a CHSH bound
  Family Index         what this suite consolidates, and why

Statevector convention: qubit `q` is bit `n-1-q` of the basis index, so qubit 0 is
the most significant bit and two-qubit operators are the usual `kron(A, B)` form.
"""
import math

import numpy as np
import gradio as gr

# ---------------------------------------------------------------- gates
H = np.array([[1, 1], [1, -1]], dtype=complex) / math.sqrt(2)
X = np.array([[0, 1], [1, 0]], dtype=complex)
Y = np.array([[0, -1j], [1j, 0]], dtype=complex)
Z = np.array([[1, 0], [0, -1]], dtype=complex)
S = np.array([[1, 0], [0, 1j]], dtype=complex)
T = np.array([[1, 0], [0, np.exp(1j * math.pi / 4)]], dtype=complex)
ONE_QUBIT = {"H": H, "X": X, "Y": Y, "Z": Z, "S": S, "T": T}
ROTATIONS = ("RX", "RY", "RZ")

# two-qubit unitaries in the (control/target) = (a, b) index order, row = 2*a + b
CX = np.array([[1, 0, 0, 0], [0, 1, 0, 0], [0, 0, 0, 1], [0, 0, 1, 0]], dtype=complex)
CZ = np.diag([1, 1, 1, -1]).astype(complex)
SWAP = np.array([[1, 0, 0, 0], [0, 0, 1, 0], [0, 1, 0, 0], [0, 0, 0, 1]], dtype=complex)
TWO_QUBIT = {"CX": CX, "CNOT": CX, "CZ": CZ, "SWAP": SWAP}


def rotation(name, theta):
    c, s = math.cos(theta / 2), math.sin(theta / 2)
    if name == "RX":
        return np.array([[c, -1j * s], [-1j * s, c]], dtype=complex)
    if name == "RY":
        return np.array([[c, -s], [s, c]], dtype=complex)
    return np.diag([np.exp(-1j * theta / 2), np.exp(1j * theta / 2)]).astype(complex)


# ---------------------------------------------------------------- simulator
def apply_one(state, U, q, n):
    """Apply a 1-qubit unitary to qubit `q` by moving its axis to the front."""
    a = np.moveaxis(state.reshape([2] * n), n - 1 - q, 0)
    a = np.tensordot(U, a, axes=([1], [0]))
    return np.moveaxis(a, 0, n - 1 - q).reshape(-1)


def apply_two(state, U, qa, qb, n):
    """Apply a 2-qubit unitary; qa/qb are reordered so the matrix index is 2*a + b."""
    axa, axb = n - 1 - qa, n - 1 - qb
    a = np.moveaxis(state.reshape([2] * n), [axa, axb], [0, 1])
    a = (U @ a.reshape(4, -1)).reshape(a.shape)
    return np.moveaxis(a, [0, 1], [axa, axb]).reshape(-1)


def parse_circuit(text, n=None):
    """Parse `H 0` / `RZ 1 0.785` / `CX 0 1` lines. `*` means every qubit."""
    lines = []
    for raw in (text or "").splitlines():
        line = raw.split("#", 1)[0].strip()
        if line:
            lines.append(line.split())
    if not lines:
        raise ValueError("empty circuit — give at least one gate")
    qubits = []
    for parts in lines:
        gate = parts[0].upper()
        if gate in ONE_QUBIT or gate in ROTATIONS:
            if len(parts) < 2:
                raise ValueError(f"'{parts[0]}' needs a qubit, e.g. '{gate} 0'")
            if parts[1] != "*":
                qubits.append(int(parts[1]))
        elif gate in TWO_QUBIT:
            if len(parts) < 3:
                raise ValueError(f"'{parts[0]}' needs two qubits, e.g. '{gate} 0 1'")
            qubits += [int(parts[1]), int(parts[2])]
        else:
            raise ValueError(f"unknown gate '{parts[0]}'")
    if not qubits:
        # only possible when every gate targets `*`, which cannot imply a width
        if not n:
            raise ValueError("a circuit that only uses '*' needs the qubit count filled in")
        n = int(n)
    else:
        n = int(n) if n else (max(qubits) + 1)
    if n < 1 or n > 12:
        raise ValueError("qubit count must be between 1 and 12")
    if qubits and max(qubits) >= n:
        raise ValueError(f"circuit uses qubit {max(qubits)} but n={n}")
    return lines, n


def simulate(text, n=None):
    """Run a circuit; return (final state, parsed lines, n)."""
    lines, n = parse_circuit(text, n)
    state = np.zeros(2 ** n, dtype=complex)
    state[0] = 1.0
    for parts in lines:
        gate = parts[0].upper()
        if gate in ONE_QUBIT:
            targets = range(n) if parts[1] == "*" else [int(parts[1])]
            for q in targets:
                state = apply_one(state, ONE_QUBIT[gate], q, n)
        elif gate in ROTATIONS:
            if len(parts) < 3:
                raise ValueError(f"'{gate}' needs a qubit and an angle, e.g. '{gate} 0 0.785'")
            theta = float(parts[2])
            targets = range(n) if parts[1] == "*" else [int(parts[1])]
            U = rotation(gate, theta)
            for q in targets:
                state = apply_one(state, U, q, n)
        else:
            state = apply_two(state, TWO_QUBIT[gate], int(parts[1]), int(parts[2]), n)
        state = state / np.linalg.norm(state)  # guards against accumulated round-off
    return state, lines, n


def basis_label(i, n):
    return "|" + format(int(i), f"0{n}b") + ">"


def fmt_complex(z):
    return f"{z.real:+.4f}{z.imag:+.4f}i"


def z_expectations(state, n):
    probs = np.abs(state) ** 2
    out = []
    for q in range(n):
        sign = np.array([1.0 if not ((i >> (n - 1 - q)) & 1) else -1.0 for i in range(2 ** n)])
        out.append(float(np.sum(sign * probs)))
    return out


# ---------------------------------------------------------------- tab 1
DEFAULT_CIRCUIT = """H 0
CX 0 1
RZ 1 0.785
H 1
CX 0 2"""


def circuit_sim(text, n_qubits):
    try:
        state, lines, n = simulate(text, n_qubits or None)
    except (ValueError, TypeError) as e:
        return f"Circuit error: {e}", "", ""
    probs = np.abs(state) ** 2
    head = [
        f"{len(lines)} gates on {n} qubits   |state| = {np.linalg.norm(state):.12f}",
        f"non-zero amplitudes: {int(np.sum(probs > 1e-12))} of {2 ** n}",
        "",
        "probabilities",
    ]
    head += [f"  {basis_label(i, n):>12}  {probs[i]:.6f}" for i in np.argsort(-probs) if probs[i] > 1e-9]
    amps = ["amplitudes (top 16 by probability)"]
    for i in np.argsort(-probs)[:16]:
        if probs[i] > 1e-12:
            amps.append(f"  {basis_label(i, n):>12}  {fmt_complex(state[i])}   p = {probs[i]:.6f}")
    zs = z_expectations(state, n)
    zs_out = ["<Z_q> per qubit  (-1 / +1 is a definite bit, 0 is full superposition)"]
    zs_out += [f"  q{q}: {v:+.6f}" for q, v in enumerate(zs)]
    return "\n".join(head), "\n".join(amps), "\n".join(zs_out)


INTRO_SIM = ("A text circuit language (`H 0`, `CX 0 1`, `RZ 1 0.785`, `H *` for all qubits) "
             "run on a hand-written statevector simulator. CPU only, no qiskit.")


# ---------------------------------------------------------------- tab 2
def superposition(n_qubits, shots):
    n = int(n_qubits)
    state, _, n = simulate("\n".join(f"H {q}" for q in range(n)), n)
    probs = np.abs(state) ** 2
    rng = np.random.default_rng(0)
    counts = np.bincount(rng.choice(2 ** n, size=int(shots), p=probs), minlength=2 ** n)
    N = 2 ** n
    rows = [
        f"n = {n} qubits -> N = {N} basis states, each with amplitude 1/sqrt(N) = {1 / math.sqrt(N):.6f}",
        f"max |p - 1/N| over all states = {np.max(np.abs(probs - 1 / N)):.3e}",
        "",
        f"{'state':>12}  {'amplitude':>22}  {'p':>10}  {'count/'+str(int(shots)):>12}  freq",
    ]
    for i in range(min(N, 16)):
        rows.append(f"{basis_label(i, n):>12}  {fmt_complex(state[i]):>22}  {probs[i]:>10.6f}  "
                    f"{counts[i]:>12}  {counts[i] / shots:.4f}")
    if N > 16:
        rows.append(f"  ... {N - 16} more states, all with the same probability")
    rows += [
        "",
        "why the sqrt(N) matters",
        f"  a single measurement identifies one of {N} outcomes: P(correct) = 1/N = {1 / N:.6f}",
        f"  the standard error of the estimate falls as 1/sqrt(N) = {1 / math.sqrt(N):.6f} per shot",
        f"  to reach a fixed confidence you need O(N) shots classically — see Grover Search "
        f"for the quantum route at O(sqrt(N))",
    ]
    return "\n".join(rows)


INTRO_SUPER = ("N = 2^n basis states carried at once, with the uniform amplitude 1/sqrt(N) "
               "verified state by state against seeded sampling.")


# ---------------------------------------------------------------- tab 3
def grover(n_qubits, marked):
    n = int(n_qubits)
    N = 2 ** n
    marked = int(marked)
    if not 0 <= marked < N:
        return f"The marked item must be between 0 and {N - 1} for n = {n}."
    state = np.ones(N, dtype=complex) / math.sqrt(N)
    optimal = max(1, int(round(math.pi / 4 * math.sqrt(N))))
    rows = [
        f"search space N = {N}   marked item = {basis_label(marked, n)}",
        f"classical search: {N // 2} trials on average. Grover: about pi/4 * sqrt(N) = "
        f"{math.pi / 4 * math.sqrt(N):.2f} iterations.",
        "",
        f"{'iteration':>9}  {'P(marked)':>12}  {'amplitude(marked)':>20}  {'max other p':>12}",
    ]
    p = float(abs(state[marked]) ** 2)
    rows.append(f"{0:>9}  {p:>12.6f}  {fmt_complex(state[marked]):>20}  "
                f"{float(np.max(np.delete(np.abs(state) ** 2, marked))):>12.6f}")
    for it in range(1, optimal + 2):
        state = state.copy()
        state[marked] *= -1.0                                    # oracle: phase flip the target
        state = 2 * np.mean(state) * np.ones(N, dtype=complex) - state   # diffusion: 2|s><s| - I
        p = float(abs(state[marked]) ** 2)
        others = float(np.max(np.delete(np.abs(state) ** 2, marked)))
        rows.append(f"{it:>9}  {p:>12.6f}  {fmt_complex(state[marked]):>20}  {others:>12.6f}")
        if p > 0.999:
            break
    rows += [
        "",
        "amplitude amplification: the oracle flips the sign of the marked amplitude, the "
        "diffusion mirrors every amplitude about the mean. Both steps together rotate the "
        "state toward the marked item by a fixed angle, so the success probability rises "
        "as sin^2((2k+1) theta) with sin(theta) = 1/sqrt(N).",
        "past the optimal iteration count the same rotation carries the state away again — "
        "the table shows the rise, and the peak is the shortest route to the answer.",
    ]
    return "\n".join(rows)


INTRO_GROVER = ("Grover's algorithm run for real on the simulator: the oracle and the "
                "diffusion operator applied iteration by iteration, with the success "
                "probability after each one.")


# ---------------------------------------------------------------- tab 4
def reduced_density(state, q, n):
    a = np.moveaxis(state.reshape([2] * n), n - 1 - q, 0).reshape(2, -1)
    return a @ a.conj().T


def entropy_of(state, q, n):
    ev = np.linalg.eigvalsh(reduced_density(state, q, n))
    ev = ev[ev > 1e-12]
    return float(-np.sum(ev * np.log2(ev)))


def bell_states():
    out = ["Bell states, each produced by gates on |00> (not written down by hand):", ""]
    setup = {"|Phi+>": "H 0\nCX 0 1", "|Phi->": "H 0\nCX 0 1\nZ 0",
             "|Psi+>": "H 0\nCX 0 1\nX 0", "|Psi->": "H 0\nCX 0 1\nZ 0\nX 0"}
    for name, circ in setup.items():
        state, _, n = simulate(circ, 2)
        amps = "  ".join(f"{basis_label(i, n)}={fmt_complex(state[i])}"
                         for i in range(4) if abs(state[i]) > 1e-9)
        gates = ", ".join(circ.splitlines())
        out.append(f"{name:7} -> {amps}")
        out.append(f"  gates: {gates}")
        out.append(f"  S(qubit0) = {entropy_of(state, 0, n):.4f} bits   <- maximal entanglement")
        out.append("")
    return "\n".join(out)


def ghz_states():
    out = ["Greenberger-Horne-Zeilinger states: H on qubit 0, then CNOT down the chain.", ""]
    for n in (2, 3, 4):
        circ = "H 0\n" + "\n".join(f"CX {q} {q + 1}" for q in range(n - 1))
        state, _, n = simulate(circ, n)
        probs = np.abs(state) ** 2
        live = [(basis_label(i, n), probs[i]) for i in range(2 ** n) if probs[i] > 1e-9]
        out.append(f"n = {n}: " + ", ".join(f"{lab} p={p:.4f}" for lab, p in live)
                   + f"   S(qubit0) = {entropy_of(state, 0, n):.4f} bit")
    out += ["", "Only all-zeros and all-ones ever appear, however many qubits are added: "
                "the whole register collapses as one object."]
    return "\n".join(out)


def chsh():
    """E(a,b) = <phi| A(a) x B(b) |phi> with A(t) = cos(t) Z + sin(t) X."""
    state, _, n = simulate("H 0\nCX 0 1", 2)
    pairs = [("a=0", 0.0, "b=+pi/4", math.pi / 4), ("a=0", 0.0, "b'=-pi/4", -math.pi / 4),
             ("a'=pi/2", math.pi / 2, "b=+pi/4", math.pi / 4),
             ("a'=pi/2", math.pi / 2, "b'=-pi/4", -math.pi / 4)]
    rows = ["CHSH correlations from the simulator, |Phi+> on |00>:",
            f"{'direction pair':30} {'E(a,b)':>10}", ""]
    es = []
    for la, ta, lb, tb in pairs:
        A = math.cos(ta) * Z + math.sin(ta) * X
        B = math.cos(tb) * Z + math.sin(tb) * X
        M = np.kron(A, B)
        e = float(np.real(state.conj() @ M @ state))
        es.append(e)
        rows.append(f"{la + '  vs  ' + lb:30} {e:>10.6f}")
    S = es[0] + es[1] + es[2] - es[3]
    rows += [
        "",
        f"S = E(a,b) + E(a,b') + E(a',b) - E(a',b') = {S:.6f}",
        f"2 * sqrt(2) = {2 * math.sqrt(2):.6f}  <- the quantum maximum (Tsirelson bound)",
        "classical bound: |S| <= 2. A value above 2 cannot be produced by any local "
        "hidden-variable model, which is the measured statement of Bell's theorem.",
    ]
    return "\n".join(rows)


def entanglement_lab():
    return bell_states() + "\n\n" + ghz_states() + "\n\n" + chsh()


INTRO_ENT = ("Bell and GHZ states built from gates, their entanglement entropy, and a CHSH "
             "correlation table computed from the statevector — the classical bound is "
             "exceeded numerically.")


# ---------------------------------------------------------------- tab 5
FAMILY = [
    ("QUANTUM_AI", "60 lines: qiskit + qiskit_aer + torch + transformers + sklearn"),
    ("QUANTUM_COGNITION", "89 lines: qiskit simulation behind a chat shell"),
    ("QUANTUM_COGNITION_CHAT_AI", "97 lines: qiskit + a Hugging Face pipeline at import"),
    ("QUANTUM_COGNITION_CHATBOT", "100 lines: same stack, model loaded at import time"),
    ("QUANTUM_FUZZY_CONSTRAINTS", "106 lines: qiskit behind a fuzzy-constraint demo"),
    ("QUANTUM_GEN", "80 lines: qiskit + generated text"),
    ("QUANTUM_GENERATIVE_MODEL", "79 lines: model download at import time"),
    ("QUANTUM_GENIE", "61 lines: qiskit + a text pipeline"),
    ("QUANTUM_GENIUS", "118 lines: qiskit + transformers, no CPU path"),
    ("QUANTUM_IMG_GEN", "55 lines: qiskit + torch + facebook/bart-large-mnli + Olivetti faces"),
    ("QUANTUM_SUPERPOSITION_GEN", "60 lines: qiskit + GPT-2 weights loaded at import"),
]

REASONS = ("Every member states its dependencies as `qiskit`, `qiskit_aer`, `torch`, "
           "`transformers` and often `sklearn` and `matplotlib`, and several download model "
           "weights (`gpt2`, `facebook/bart-large-mnli`) or a dataset (`fetch_olivetti_faces`) "
           "at import time. On a free cpu-basic Space that means a multi-gigabyte image and a "
           "cold start that pulls weights before the app can serve, which is why all 11 sat "
           "asleep rather than running.")


def family_index():
    rows = ["| consolidated Space | what blocks it on a free CPU Space |", "|---|---|"]
    rows += [f"| `{name}` | {why} |" for name, why in FAMILY]
    head = ["# Family Index", "",
            f"**{len(FAMILY)} Spaces in the `QUANTUM_*` family**, measured on 2026-10-12. "
            "They are consolidated into the four tabs above, which run in **one** cpu-basic "
            "slot instead of one Space per project.", "", REASONS, "",
            "The replacement keeps the quantum mechanics (superposition, interference, "
            "amplitude amplification, entanglement, Bell violation) and drops the model "
            "weights that made the originals unhostable. Nothing here is simulated with "
            "another library: the amplitudes are computed directly.", ""]
    return "\n".join(head + rows)


INTRO_FAMILY = ("The 11 `QUANTUM_*` Spaces this suite consolidates, and the measured reason "
                "each one cannot run on a free cpu-basic Space.")


# ---------------------------------------------------------------- app
INTRO = {
    "Circuit Simulator": INTRO_SIM,
    "Superposition Explorer": INTRO_SUPER,
    "Grover Search": INTRO_GROVER,
    "Entanglement Lab": INTRO_ENT,
    "Family Index": INTRO_FAMILY,
}


def page():
    with gr.Tabs():
        with gr.Tab("Circuit Simulator"):
            gr.Markdown(INTRO_SIM)
            with gr.Row():
                code = gr.Code(label="circuit", language="python", value=DEFAULT_CIRCUIT)
                nq = gr.Number(label="qubits (blank = infer from the circuit)", value=3, precision=0)
            o1 = gr.Textbox(lines=22, label="probabilities")
            with gr.Row():
                o2 = gr.Textbox(lines=18, label="amplitudes")
                o3 = gr.Textbox(lines=18, label="expectation of Z")
            gr.Button("Run circuit").click(circuit_sim, [code, nq], [o1, o2, o3])

        with gr.Tab("Superposition Explorer"):
            gr.Markdown(INTRO_SUPER)
            with gr.Row():
                n2 = gr.Slider(1, 6, value=4, step=1, label="qubits")
                shots = gr.Slider(256, 8192, value=4096, step=256, label="shots")
            o4 = gr.Textbox(lines=26, label="uniform superposition")
            gr.Button("Explore").click(superposition, [n2, shots], o4)

        with gr.Tab("Grover Search"):
            gr.Markdown(INTRO_GROVER)
            with gr.Row():
                n3 = gr.Slider(3, 6, value=4, step=1, label="qubits")
                marked = gr.Number(label="marked item (integer)", value=11, precision=0)
            o5 = gr.Textbox(lines=26, label="iterations")
            gr.Button("Search").click(grover, [n3, marked], o5)

        with gr.Tab("Entanglement Lab"):
            gr.Markdown(INTRO_ENT)
            o6 = gr.Textbox(lines=34, label="states, entropy and CHSH")
            gr.Button("Compute").click(entanglement_lab, None, o6)

        with gr.Tab("Family Index"):
            gr.Markdown(family_index())


if __name__ == "__main__":
    with gr.Blocks(title="Quantum Suite", theme=gr.themes.Soft()) as demo:
        gr.Markdown("# Quantum Suite\n"
                    "A statevector quantum simulator over a whole project family, in a "
                    "single cpu-basic Space. **No qiskit, no torch, no transformers, no GPU.**\n\n"
                    "Companion to the [GGUF Toolkit](https://huggingface.co/spaces/AEUPH/gguf-toolkit).")
        page()
    demo.launch()