!pip install 'classiq[qsp]' -qq

import matplotlib.pyplot as plt
import numpy as np

from classiq import *
from classiq.applications.qsp import qsp_approximate, qsvt_phases

np.random.seed(7)

N = 8  # problem size; the register has int(log2(N)) qubits

# A random Boolean target function f: [N] -> {0, 1}
F_TABLE = np.random.randint(0, 2, size=N)
print(f"f(x) for x = 0,...,{N-1}:  {F_TABLE.tolist()}")

@qfunc
def apply_basis_phase(theta: CReal, x_int: CInt, qvar: QNum) -> None:
    r"""Apply $\exp(i \cdot \theta \cdot |x_\text{int}\rangle\langle x_\text{int}|)$ on ``qvar``.

Implemented by controlling on the equality predicate ``qvar == x_int``;
    the inner ``phase($\theta$)`` becomes a relative phase on the controlled
    subspace and the identity elsewhere, which is exactly $V_t$.
    """
    control(qvar == x_int, lambda: phase(theta))


@qfunc
def quantum_oracle_sketch_boolean(
    theta: CReal, x_samples: CArray[CInt], qvar: QNum
) -> None:
    r"""Sketched Boolean phase oracle $V = \prod_t \exp(i \theta |x_t\rangle\langle x_t|)$.

    ``x_samples`` should be pre-filtered to the indices with $f(x_t) = 1$;
    a single fixed angle $\theta = \pi N / M$ is then applied per sample.
    """
    repeat(
        count=x_samples.len,
        iteration=lambda t: apply_basis_phase(theta, x_samples[t], qvar),
    )

# For circuit synthesis we use a smaller M so the visualization stays readable;
# the numpy benchmark below validates the algorithm at larger M.
M_demo = 200
demo_samples = np.random.randint(0, N, size=M_demo)
positive_samples = demo_samples[F_TABLE[demo_samples] == 1].tolist()
theta = np.pi * N / M_demo  # τ / M with τ = π N
n_qubits = int(np.log2(N))


@qfunc
def main(qvar: Output[QNum]) -> None:
    allocate(n_qubits, qvar)
    hadamard_transform(qvar)
    quantum_oracle_sketch_boolean(theta, positive_samples, qvar)


qprog = synthesize(main)
show(qprog)

from classiq.execution import sample

# 2-qubit Bell-state setup: f = [0,0,1,1] phases the MSB subspace.
N_ex, M_ex = 4, 500
F_ex = np.array([0, 0, 1, 1])
theta_ex = np.pi * N_ex / M_ex

np.random.seed(0)
demo_samples_ex = np.random.randint(0, N_ex, size=M_ex)
positive_samples_ex = demo_samples_ex[F_ex[demo_samples_ex] == 1].tolist()


@qfunc
def main(qvar: Output[QNum]) -> None:
    bell = QArray[QBit, 2]()
    allocate(2, qvar)
    # 1) Prepare |Φ+⟩.
    # 2) Apply the sketched oracle on the QNum view.
    # 3) Inverse Bell prep — rotates the Bell basis to the computational basis,
    #    in particular |Φ-⟩ → |10⟩ (qvar = 1).
    within_apply(
        within=lambda: (
            bind(qvar, bell),  # qvar (initialised) → bell
            H(bell[0]),
            CX(bell[0], bell[1]),
            bind(bell, qvar),  # bell → qvar so the oracle sees a QNum
        ),
        apply=lambda: quantum_oracle_sketch_boolean(
            theta_ex, positive_samples_ex, qvar
        ),
    )


qprog_bell = synthesize(main)

df = sample(qprog_bell, "simulator", num_shots=10_000)

fidelity = float(df.loc[df["qvar"] == 1, "probability"].sum())
print(f"Fidelity |⟨Φ-|V|Φ+⟩|² at M={M_ex}:  {fidelity:.4f}")

def ideal_oracle(f_table):
    r"""Diagonal Boolean phase oracle $O|x\rangle = (-1)^{f(x)} |x\rangle$."""
    return np.diag((-1.0) ** f_table.astype(float))


def sketched_oracle(f_table, x_samples):
    r"""Numpy reference for $V = \prod_t \exp(i \tau f(x_t) |x_t\rangle\langle x_t| / M)$, $\tau = \pi N$.

Each $V_t$ is diagonal, so the product collapses to a single diagonal:
    $V_{xx} = \exp(i \pi N m_x f(x))$ where $m_x$ is the empirical frequency of $x$.
    """
    n = f_table.size
    m_total = x_samples.size
    counts = np.bincount(x_samples, minlength=n)
    m = counts / m_total
    return np.diag(np.exp(1j * np.pi * n * m * f_table))


# Single realization at M = 6000
M = 4000
x_samples = np.random.randint(0, N, size=M)
V = sketched_oracle(F_TABLE, x_samples)
O = ideal_oracle(F_TABLE)
err = np.linalg.norm(V 

- O, ord=2)
print(f"||V 

- O||_2  at  M={M}, N={N}:  {err:.3e}")

# Sweep M to verify the single-realization vs channel-bias scalings.
M_values = np.unique(np.logspace(2, 4.5, 12, dtype=int))

# n_trials = number of independent random sample sequences λ drawn per M.
n_trials = 1500

single_err = np.empty(M_values.size)
channel_err = np.empty(M_values.size)
for k, M_k in enumerate(M_values):
    Vs = [
        sketched_oracle(F_TABLE, np.random.randint(0, N, size=int(M_k)))
        for _ in range(n_trials)
    ]
    single_err[k] = np.mean([np.linalg.norm(Vk 

- O, ord=2) for Vk in Vs])
    channel_err[k] = np.linalg.norm(np.mean(Vs, axis=0) 

- O, ord=2)

# Theoretical guides: single-realization ~ sqrt(N/M); channel bias ~ N/M.
c_single = float(np.median(single_err / np.sqrt(N / M_values)))
c_channel = float(np.median(channel_err * M_values / N))

fig, ax = plt.subplots(figsize=(6, 4.5))
ax.loglog(M_values, single_err, "o-", label=r"single realization $\|V_\lambda - O\|_2$")
ax.loglog(
    M_values, channel_err, "s-", label=r"channel bias $\|\mathbb{E}[V_\lambda] - O\|_2$"
)
ax.loglog(
    M_values,
    c_single * np.sqrt(N / M_values),
    "--",
    color="C0",
    alpha=0.5,
    label=rf"${c_single:.2f}\,\sqrt{{N/M}}$",
)
ax.loglog(
    M_values,
    c_channel * N / M_values,
    "--",
    color="C1",
    alpha=0.5,
    label=rf"${c_channel:.2f}\,N/M$",
)
ax.set_xlabel("Number of samples $M$")
ax.set_ylabel("Spectral-norm error")
ax.set_title(f"Boolean oracle sketching, $N = {N}$, {n_trials} trials")
ax.legend(fontsize=9)
ax.grid(True, which="both", alpha=0.3)
plt.tight_layout()
plt.show()

@qfunc
def quantum_oracle_sketch_known_p(
    probs: CArray[CReal],
    x_samples: CArray[CInt],
    qvar: QNum,
) -> None:
    r"""
    Sketched Boolean phase oracle for a known non-uniform distribution $p$.

    Args:
        probs: The probabilities of the samples.
        x_samples: The classical input samples.
        qvar: The quantum variable to apply the oracle to.
    """
    repeat(
        count=x_samples.len,
        iteration=lambda t: apply_basis_phase(
            np.pi / (probs[t] * probs.len), x_samples[t], qvar
        ),
    )

# Problem setup: known bounds on the distribution p, but p itself unknown to the algorithm
N_unk = 4
n_qubits_unk = int(np.log2(N_unk))
p_min, p_max = 0.10, 0.45
kappa_unk = p_max / p_min  # condition number of p restricted to its support

# Toy ground-truth p (used ONLY to sample data; the algorithm never accesses it).
p_vec = np.array([0.10, 0.15, 0.30, 0.45])
F_TABLE_unk = np.array([0, 0, 1, 1])  # Boolean target; active basis states: 2, 3 (I⊗Z)


# Compute QSVT phases for the threshold polynomial.
lambda_star = np.sin(1.0 / kappa_unk)

EPS_QSVT = 2e-2
degree_qsvt = int(np.ceil(np.log(1.0 / EPS_QSVT) / lambda_star))
if degree_qsvt % 2 == 1:
    degree_qsvt += 1  # enforce even parity

# Smooth threshold approximant: ≈ +1 near w=0, ≈ -1 for |w| ≥ λ*; tanh transition at λ*/

2.
SCALE = 0.95  # keep |P| < 1 strictly inside [-1, 1] for QSP convergence


def threshold_target(w):
    return -SCALE * np.tanh(4.0 * (np.abs(w) - 0.5 * lambda_star) / lambda_star)


poly_coeffs, fit_err = qsp_approximate(
    threshold_target,
    degree=degree_qsvt,
    parity=0,  # even
    interval=[0.0, 1.0],
    plot=True,
)
threshold_phases = qsvt_phases(poly_coeffs)
print(
    f"κ = {kappa_unk:.2f},  λ* = {lambda_star:.3f},  "
    f"degree = {degree_qsvt},  fit err = {fit_err:.2e},  "
    f"# QSVT phases = {len(threshold_phases)}"
)


# Classical sample preparation for the U = e^{iΛ} sketch (Stage 1).
np.random.seed(0)
M_unk = 20  # small M so the QSVT-chain depth stays inside the synthesis timeout
x_samples_unk = np.random.choice(np.arange(N_unk), size=M_unk, p=p_vec)
active_unk = x_samples_unk[F_TABLE_unk[x_samples_unk] == 1].tolist()

# Fixed evolution time t = 1/p_max gives a SINGLE per-sample angle θ = t/M,
# independent of x_t — we don't divide by p(x_t) because p is unknown.
# This is the uniform-θ Boolean sketcher, with the scaling set by the bound p_max
# rather than by 1/N.
theta_unk = 1.0 / (p_max * M_unk)


# Hadamard-test block encoding of sin(Λ) (Stage 2)
@qfunc
def hadamard_test_W(
    theta: CReal,
    x_samples: CArray[CInt],
    qvar: QNum,
    aux_W: QBit,
) -> None:
    r"""$W = X_a \cdot S_a \cdot X_a \cdot H_a \cdot (cU^\dagger) \cdot X_a \cdot (cU) \cdot H_a$

    Block-encodes $\sin(\Lambda)$ in the $|0\rangle_{\text{aux\_W}}$ subspace:
        $\langle 0_{\text{aux\_W}}| W |0_{\text{aux\_W}}\rangle = (U - U^\dagger)/(2i) = \sin(\Lambda)$.
    """
    H(aux_W)
    control(aux_W, lambda: quantum_oracle_sketch_boolean(theta, x_samples, qvar))
    X(aux_W)
    control(
        aux_W,
        lambda: invert(lambda: quantum_oracle_sketch_boolean(theta, x_samples, qvar)),
    )
    X(aux_W)
    H(aux_W)
    S(aux_W)  # upto a global phase, gives diag(-i, 1), which is the desired.
    X(aux_W)  # swap |0⟩_a ↔ |1⟩_a so sin(Λ) lives in the |0⟩_a block


@qfunc
def proj_W_block(aux_W: QBit, flag: QBit) -> None:
    r"""Flip `flag` when `aux_W` is in $|0\rangle$ — projector onto the $W$ block-encoded subspace."""
    control(aux_W == 0, lambda: X(flag))


# Top-level program: apply the QSVT threshold polynomial to W
@qfunc
def main(
    qvar: Output[QNum],
    aux_W: Output[QBit],
    qsvt_aux: Output[QBit],
) -> None:
    allocate(n_qubits_unk, qvar)
    allocate(1, aux_W)
    allocate(1, qsvt_aux)

    # QSVT chain applied to W; the polynomial P(sin Λ) approximates
    # the target oracle O = Σ_x (-1)^{f(x)} |x⟩⟨x| in the |0⟩_{aux_W} block.
    qsvt(
        phase_seq=threshold_phases.tolist(),
        proj_cnot_1=lambda flag: proj_W_block(aux_W, flag),
        proj_cnot_2=lambda flag: proj_W_block(aux_W, flag),
        u=lambda: hadamard_test_W(theta_unk, active_unk, qvar, aux_W),
        aux=qsvt_aux,
    )


# Increasing the synthesis timeout above the 300 s default to implement the QSVT
qprog_unk = synthesize(main, preferences=Preferences(timeout_seconds=600))
show(qprog_unk)

# 2-qubit Bell-state setup for the QSVT-corrected sketch.


@qfunc
def main(
    qvar: Output[QNum],
    aux_W: Output[QBit],
    qsvt_aux: Output[QBit],
) -> None:
    bell = QArray[QBit, 2]()
    allocate(2, qvar)
    allocate(1, aux_W)
    allocate(1, qsvt_aux)

    # W * A * W^{-1}.
    # W: cast qvar -> bell, prep |Φ+⟩, cast back to qvar so qsvt sees a QNum.
    # W^{-1} (auto-generated): cast qvar -> bell, inverse Bell prep
    # (rotates Bell basis to computational basis, |Φ-⟩ → |10⟩), cast back to qvar.
    within_apply(
        within=lambda: (
            bind(qvar, bell),
            H(bell[0]),
            CX(bell[0], bell[1]),
        ),
        apply=lambda: qsvt(
            phase_seq=threshold_phases.tolist(),
            proj_cnot_1=lambda flag: proj_W_block(aux_W, flag),
            proj_cnot_2=lambda flag: proj_W_block(aux_W, flag),
            u=lambda: hadamard_test_W(theta_unk, active_unk, bell, aux_W),
            aux=qsvt_aux,
        ),
    )


# Same timeout bump as the cell above — the QSVT chain is the same size.
qprog_unk_bell = synthesize(main, preferences=Preferences(timeout_seconds=600))

# printing the circuit metrics, showing that the QSVT chain is indeed the dominant contributor to the depth and gate count
print(
    f"depth = {qprog_unk_bell.transpiled_circuit.depth}, "
    f"width = {qprog_unk_bell.data.width}, "
)

df_unk = sample(qprog_unk_bell, "simulator", num_shots=10_000)

# Post-select on the QSVT block (aux_W = qsvt_aux = 0) and read |Φ-⟩ off qvar = 

1.
in_block = df_unk[(df_unk["aux_W"] == 0) & (df_unk["qsvt_aux"] == 0)]
p_block = float(in_block["probability"].sum())
fidelity_unk = (
    float(in_block.loc[in_block["qvar"] == 1, "probability"].sum()) / p_block
    if p_block > 0
    else float("nan")
)
print(f"P(QSVT block, aux_W = qsvt_aux = 0)              :  {p_block:.4f}")
print(f"Fidelity |⟨Φ-|P_QSVT(W)|Φ+⟩|² (post-selected)    :  {fidelity_unk:.4f}")