import numpy as np

from classiq import *

np.random.seed(552)
NUM_QUBITS = 2

BASIS_ANGLES = np.array(
    [
        [np.pi / 2, np.pi],
        [np.pi / 2, 3 * np.pi / 2],
        [0.0, 0.0],
    ],
    dtype=float,
)

@qfunc
def main(
    thetas: CArray[CReal, NUM_QUBITS],
    phis: CArray[CReal, NUM_QUBITS],
    qarr: Output[QArray],
) -> None:
    """
    Prepares the |Phi^-> Bell state and applies the parametric per-qubit
    basis-change rotation.

The rotation angles are bound at execution time,
    one parameter set per snapshot, so a single synthesized program is
    reused across the entire shadow.

The circuit realizes a measurement in the X, Y, or Z basis (up to global phase).

Because the angles are bound at execution time, the circuit is independent
    of the random basis choice and only needs to be synthesized once.
    """
    prepare_bell_state(1, qarr)
    # performing a basis change by applying RY and RZ rotations with the angles specified in thetas and phis
    for i in range(NUM_QUBITS):
        RZ(phis[i], qarr[i])
        RY(thetas[i], qarr[i])

qprog = synthesize(main)
show(qprog)

def calculate_shadow_pauli(num_snapshots, num_qubits):
    """
    Computes a Pauli-shadow of size `num_snapshots` from a single synthesized
    program.

The per-qubit basis-change angles are sampled in Python and
    submitted in one `sample` call.

    Args:
        num_snapshots (int): Size of the shadow, also the number of measurements.
        num_qubits (int): Number of qubits in the system.

    Returns:
        snapshots (list[list[int]]): per-snapshot measurement bits
            (`num_qubits` long each).
        ids (list[list[int]]): per-snapshot basis identifiers in {0, 1, 2}
            (X, Y, Z) for each qubit.
    """
    # 1) Sample the random measurement bases in Python (no need to bake them
    #    into separate circuits).
    basis_id_array = np.random.randint(3, size=(num_snapshots, num_qubits))
    ids = [[int(b) for b in row] for row in basis_id_array]

    # 2) Translate each random basis row into RY/RZ angles via BASIS_ANGLES.
    param_batch = [
        {
            "thetas": BASIS_ANGLES[basis_id_array[i], 0].tolist(),
            "phis": BASIS_ANGLES[basis_id_array[i], 1].tolist(),
        }
        for i in range(num_snapshots)
    ]

    # 3) Execute every snapshot in one cloud round-trip via `sample`.

With a
    #    list of parameter sets, `sample` returns one DataFrame per element;
    #    with `num_shots=1`, each frame has exactly one row whose `qarr`
    #    entry is the measured bitstring as a list of ints.
    result_dfs = sample(
        qprog,
        parameters=param_batch,
        num_shots=1,
        random_seed=int(np.random.randint(1e6)),
    )
    snapshots = [list(df["qarr"].iloc[0]) for df in result_dfs]

    return snapshots, ids

num_snapshots = 400
snapshots, ids = calculate_shadow_pauli(num_snapshots, NUM_QUBITS)
print(f"first ten snapshots: {snapshots[:10]}")
print(f"first ten ids: {ids[:10]}")

def snapshot_reconstruction_pauli(snapshot, snapshot_ids, num_qubits):
    """
    Reconstructs the state density matrix from one snapshot.

Helper for `reconstruction_pauli()`.

    Args:
        snapshot (list[int]): Element of snapshots; `num_qubits` long.
        snapshot_ids (list[int]): Element of the `ids` list returned by `calculate_shadow_pauli`.
        num_qubits (int): Number of qubits in the system.

    Returns:
        final_state (np.ndarray): Snapshot density matrix reconstruction.

    """
    # Density matrices of computational basis states
    zero_state = np.array([[1, 0], [0, 0]])
    one_state = np.array([[0, 0], [0, 1]])

    # S, H, I matrices
    S_matrix = np.array([[1, 0], [0, 1j]], dtype=complex)
    H_matrix = (1 / np.sqrt(2)) * np.array([[1, 1], [1, -1]])
    I_matrix = np.array([[1, 0], [0, 1]])

    # Unitaries applied in circuit.
    unitaries = [H_matrix, H_matrix @ S_matrix, I_matrix]

    final_state = [1]
    for i in range(num_qubits):
        state = zero_state if snapshot[i] == 0 else one_state
        U = unitaries[int(snapshot_ids[i])]

        # Applying Eq. (S44) from Huang, et al.
        local_state = 3 * (U.conjugate().transpose() @ state @ U) 

- I_matrix
        final_state = np.kron(final_state, local_state)

    return final_state

def reconstruction_pauli(in_snapshots, in_ids, num_qubits):
    """
    Performs a full reconstruction of the quantum state density matrix, where the quantum state is that
    prepared in `prep_state()`.

    Args:
        in_snapshots (list): List of snapshots returned by `calculate_shadow_pauli` (list of ints in {0, 1}).
        in_ids (list): List of unitary identifiers returned by `calculate_shadow_pauli` (list of ints in {0, 1, 2}).
        num_qubits (int): Number of qubits in the system.

    Returns:
        reconstructed_state (np.ndarray): Full state density matrix reconstruction (2d numpy complex array).
    """
    # Average over snapshots.
    reconstructed_state = np.zeros((2**num_qubits, 2**num_qubits), dtype=complex)
    for i in range(num_snapshots):
        result = snapshot_reconstruction_pauli(in_snapshots[i], in_ids[i], num_qubits)
        reconstructed_state += result

    return reconstructed_state / num_snapshots

reconstructed_state = reconstruction_pauli(snapshots, ids, NUM_QUBITS)
print(f"Full state density matrix: \n {np.round(reconstructed_state, decimals=6)}\n")

bell_state_density_matrix = np.array(
    [[0.5, 0, 0, -0.5], [0, 0, 0, 0], [0, 0, 0, 0], [-0.5, 0, 0, 0.5]]
)
print(f"Bell State Density Matrix: \n {bell_state_density_matrix}\n")

print(f"Comparing to: \n {bell_state_density_matrix}\n")
# distance = norm(ref_state_density_matrix, reconstructed_state)
difference = bell_state_density_matrix - reconstructed_state
distance = np.linalg.norm(difference, "fro")
print(f"Distance: {np.round(distance, decimals=6)}")

def estimate_observable_pauli(in_snapshots, in_ids, observable, div, num_qubits):
    """
    Estimates the expectation value of an observable using a classical shadow.

    Args:
        in_snapshots (list[int]): List of snapshots returned by `calculate_shadow_pauli`.
        in_ids (list[int]): List of unitary identifiers returned by `calculate_shadow_pauli`.
        observable (SparsePauliOp): Observable to be estimated.
        div (int): Number of divisions for median-of-means.
        num_qubits (int): Number of qubits in the system.

    Returns:
        expval (float): Esimated expectation value.
    """
    sum = 0
    in_snapshots = np.array(in_snapshots)
    in_ids = np.array(in_ids)

    # Loop takes care of observables that are sums.
    for element in observable.terms:
        means = []
        ops = []  # operators in observable

        # Pauli.I = 0, Pauli.X = 1, Pauli.Y = 2, Pauli.Z = 3
        for i in range(num_qubits):
            if i < len(element.paulis):
                ops.append(
                    element.paulis[i].pauli - 1
                )  # Subtracting 1 to stay consistent with format of ids, where X, Y, Z correspond to 0, 1, 

2.
            else:
                ops.append(-1)  # Add I operator if no operator present.

        ops = np.array(ops)

        for i in range(0, num_snapshots, num_snapshots // div):
            # Divide in_snapshots and in_ids into div chunks.
            in_snapshots_div, in_ids_div = (
                in_snapshots[i : i + num_snapshots // div],
                in_ids[i : i + num_snapshots // div],
            )

            prods = []
            for j in range(num_snapshots // div):
                # Terms that will be in the product.
                terms = np.ones_like(ops)

                # Create masks for identity, matching, and non-matching operators.
                non_id_mask = ops != -1
                match_mask = ops == in_ids_div[j]
                mismatch_mask = (~match_mask) & non_id_mask

                # Handle matches.
                terms[match_mask & (in_snapshots_div[j] == 0)] = 3.0
                terms[match_mask & (in_snapshots_div[j] == 1)] = -3.0

                # Handle mismatches.
                terms[mismatch_mask] = 0.0
                prods.append(np.prod(terms))

            means.append(np.mean(prods))
        sum += np.median(means) * element.coefficient
    return sum

observable = Pauli.Z(0) * Pauli.Z(1)
estimate = estimate_observable_pauli(
    snapshots, ids, observable, div=2, num_qubits=NUM_QUBITS
)
print(f"Estimated <ZZ> (Pauli shadow): {estimate:.4f}")
print(f"Error relative to ground truth value of 1.0: {1.0 - estimate:.4f}")

def error_bound_pauli(error, observables, delta):
    """
    Calculates the minimum shadow size that guarantees a given error and
    probability of failure for a unitary ensemble consisting of Pauli basis
    measurements.

Implements Eq. (5) with the random-Pauli shadow-norm bound for a sum of
    Pauli terms,

        || sum_a c_a P_a ||_shadow^2  <=  ( sum_a |c_a| * 3^(weight(P_a)/2) )^2,

    where `weight(P_a)` is the number of non-identity factors in the Pauli
    string `P_a`.

The 3^k locality factor is essential: without it the bound
    underestimates the required samples by 3^k for a k-local observable.

    Args:
        error (float): Desired error.
        observables (list[SparsePauliOp]): List of SparsePauliOps.
        delta (float): Probability of failure.

    Returns:
        samples (tuple(int, int)): (N*K, K) -- total snapshots required and
        the number of median-of-means divisions K.
    """
    M = len(observables)
    K = 2 * np.log(2 * M / delta)

    shadow_norms = []
    for obs in observables:
        bound = 0.0
        for term in obs.terms:
            weight = sum(1 for p in term.paulis if p.pauli != 0)
            bound += abs(term.coefficient) * np.sqrt(3**weight)
        shadow_norms.append(bound**2)

    N = 34 * max(shadow_norms) / error**2

    return int(np.ceil(N * K)), int(K)

observables = [observable]
required_snapshots = error_bound_pauli(0.2, observables, 0.01)
print(
    f"Number of samples required: {required_snapshots[0]}, Number of divisions: {required_snapshots[1]}"
)

# Single-qubit gate matrices used only for the lazy matrix materialization below.
_H1 = (1.0 / np.sqrt(2.0)) * np.array([[1, 1], [1, -1]], dtype=complex)
_S1 = np.array([[1, 0], [0, 1j]], dtype=complex)
_I1 = np.eye(2, dtype=complex)


def _embed_1q(g: np.ndarray, q: int, num_qubits: int) -> np.ndarray:
    """Embed a single-qubit gate `g` at qubit `q` in the n-qubit Hilbert space.

Tensor order matches `qarr`: qubit 0 is the MSB / leftmost factor."""
    out = np.array([[1.0]], dtype=complex)
    for i in range(num_qubits):
        out = np.kron(out, g if i == q else _I1)
    return out


def _cnot_matrix(ctrl: int, tgt: int, num_qubits: int) -> np.ndarray:
    """Dense CNOT matrix on `(ctrl, tgt)` in the full n-qubit space, MSB convention."""
    dim = 2**num_qubits
    u = np.zeros((dim, dim), dtype=complex)
    for x in range(dim):
        bits = [(x >> (num_qubits - 1 - i)) & 1 for i in range(num_qubits)]
        new_bits = bits[:]
        if bits[ctrl] == 1:
            new_bits[tgt] = 1 - new_bits[tgt]
        new_x = sum(new_bits[i] << (num_qubits - 1 - i) for i in range(num_qubits))
        u[new_x, x] = 1.0
    return u


def sample_clifford_circuit(num_qubits: int, depth: int = 20):
    """Sample a depth-`depth` random gate sequence from {H, S, CNOT}.

    Returns:
        seq (list[tuple]): each entry is `(label, q1, q2)` with
            `label in {"H", "S", "CX"}` and `q2 = -1` for single-qubit gates.

For n = 2 a depth of ~20 is past the 3-design mixing time of random
    Clifford circuits, so the resulting distribution is close enough to Haar
    on the moments we need.
    """
    seq = []
    for _ in range(depth):
        kind = np.random.randint(0, 3)
        if kind < 2:
            q = int(np.random.randint(0, num_qubits))
            label = "H" if kind == 0 else "S"
            seq.append((label, q, -1))
        else:
            c, t = np.random.choice(num_qubits, size=2, replace=False)
            seq.append(("CX", int(c), int(t)))
    return seq


def sequence_to_unitary(seq, num_qubits: int) -> np.ndarray:
    """Materialize a `(label, q1, q2)` gate sequence as the dense
    `2**num_qubits x 2**num_qubits` Clifford unitary.
    """
    U = np.eye(2**num_qubits, dtype=complex)
    for label, q1, q2 in seq:
        if label == "H":
            U = _embed_1q(_H1, q1, num_qubits) @ U
        elif label == "S":
            U = _embed_1q(_S1, q1, num_qubits) @ U
        else:
            U = _cnot_matrix(q1, q2, num_qubits) @ U
    return U

def calculate_shadow_clifford(num_snapshots: int, num_qubits: int):
    """Collect `num_snapshots` single-shot snapshots under random Cl(2^n) unitaries.

Each iteration synthesizes a fresh program (a new random Clifford gate
    sequence is sampled at trace-time), then samples one shot via `sample`.

The sampled gate sequence is captured in a closure-local `sequences`
    list rather than a global, so successive calls don't interfere.

    Args:
        num_snapshots (int): Number of single-shot snapshots to collect.
        num_qubits (int): Number of qubits in the system.

    Returns:
        snapshots (list[list[int]]): per-snapshot measured bitstrings.
        sequences (list[list[tuple]]): per-snapshot Clifford gate sequence
            (each a list of `(label, q1, q2)` tuples) — same length and
            order as `snapshots`.
    """
    sequences = []

    @qfunc
    def clifford_application(state: QArray) -> None:
        """Sample a random Cl(2^n) gate sequence and apply it directly using
        Classiq's native `H`, `S`, and `CX` qfuncs.

The sampled sequence is
        appended to the enclosing `sequences` list (no module-level state)."""
        seq = sample_clifford_circuit(num_qubits)
        sequences.append(seq)
        for label, q1, q2 in seq:
            if label == "H":
                H(state[q1])
            elif label == "S":
                S(state[q1])
            else:
                CX(state[q1], state[q2])

    @qfunc
    def main(qarr: Output[QArray]) -> None:
        """Bell state, followed by a random global Clifford applied as native gates."""
        prepare_bell_state(1, qarr)
        clifford_application(qarr)

    snapshots = []
    for i in range(num_snapshots):
        qprog = synthesize(main)
        if i == 0:
            show(qprog)  # render the very first synthesised circuit
        df = sample(
            qprog,
            num_shots=1,
            random_seed=int(np.random.randint(1e6)),
        )
        snapshots.append(list(df["qarr"].iloc[0]))
    return snapshots, sequences


# We use a smaller shadow than the Pauli case (50 vs 200) to decrease runtime.
num_snapshots_clifford = 50
snapshots_clifford, clifford_sequences = calculate_shadow_clifford(
    num_snapshots_clifford, NUM_QUBITS
)
print(f"Collected {len(snapshots_clifford)} Clifford-shadow snapshots")
print(f"first 5 outcomes:           {snapshots_clifford[:5]}")
print(f"first sampled gate sequence: {clifford_sequences[0]}")

def snapshot_reconstruction_clifford(snapshot, U, num_qubits):
    """Single-snapshot Clifford-shadow reconstruction.

Uses the fact that ``|b><b|`` is rank-1, so
    ``U^dagger |b><b| U = conj(v) outer v`` where ``v = U[b, :]`` is a single
    row of the Clifford unitary.

    Args:
        snapshot (list[int]): Measurement outcome, a list of `num_qubits` bits in
            `{0, 1}` ordered with qubit 0 as MSB (matching the convention used
            by `qarr` in the Pauli-shadow data acquisition).
        U (np.ndarray): Dense `2**num_qubits x 2**num_qubits` complex matrix of
            the random Clifford unitary that was applied to the state just
            before measurement.
        num_qubits (int): Number of qubits in the system.

    Returns:
        np.ndarray: The single-snapshot density-matrix estimator
        `(2**n + 1) * U^dagger |b><b| U 

- I`, of shape
        `(2**num_qubits, 2**num_qubits)`.
    """
    dim = 2**num_qubits
    # Convert bit list to basis index (qubit 0 = MSB by convention used in `qarr`).
    b_idx = sum(int(snapshot[i]) << (num_qubits - 1 - i) for i in range(num_qubits))
    v = U[b_idx, :]
    return (dim + 1) * np.outer(v.conjugate(), v) - np.eye(dim, dtype=complex)


def reconstruction_clifford(in_snapshots, in_sequences, num_qubits):
    """Average the snapshot reconstructions across all Clifford-shadow samples.

    `in_sequences[i]` is the gate sequence for snapshot `i`; the corresponding
    `2**n x 2**n` unitary is materialized inline via `sequence_to_unitary`.
    """
    rho_hat = np.zeros((2**num_qubits, 2**num_qubits), dtype=complex)
    for snap, seq in zip(in_snapshots, in_sequences):
        U = sequence_to_unitary(seq, num_qubits)
        rho_hat += snapshot_reconstruction_clifford(snap, U, num_qubits)
    return rho_hat / len(in_snapshots)


reconstructed_state_clifford = reconstruction_clifford(
    snapshots_clifford, clifford_sequences, NUM_QUBITS
)
print("Clifford-shadow reconstruction:")
print(np.round(reconstructed_state_clifford, decimals=4))
print("\nBell state target:")
print(bell_state_density_matrix)

distance_clifford = np.linalg.norm(
    bell_state_density_matrix - reconstructed_state_clifford, "fro"
)
print(f"\nFrobenius distance (Clifford shadow): {distance_clifford:.4f}")

def estimate_observable_clifford(
    in_snapshots, in_sequences, observable, div, num_qubits
):
    """
    Estimates <O> = tr(O rho) from a Clifford classical shadow with
    median-of-means.

Each snapshot estimator is rho_hat_i = (2^n + 1) U_i^dagger |b_i><b_i| U_i 

- I,
    so tr(O rho_hat_i) = (2^n + 1) (U_i O U_i^dagger)[b_i, b_i] - tr(O).

    Args:
        in_snapshots (list[list[int]]): Bits per snapshot, num_qubits long each.
        in_sequences (list[list[tuple]]): Sampled Clifford gate sequences,
            one per snapshot.

Each is materialized to a 2^n x 2^n unitary
            inline via `sequence_to_unitary`.
        observable (SparsePauliOp): Observable to estimate.
        div (int): Number of median-of-means divisions K.
        num_qubits (int): Number of qubits in the system.

    Returns:
        float: estimated expectation value.
    """
    dim = 2**num_qubits
    N = len(in_snapshots)
    chunk = N // div

    O = hamiltonian_to_matrix(observable)
    trace_O = np.trace(O).real

    means = []
    for k in range(div):
        snaps_k = in_snapshots[k * chunk : (k + 1) * chunk]
        seq_k = in_sequences[k * chunk : (k + 1) * chunk]

        vals = []
        for snap, seq in zip(snaps_k, seq_k):
            U = sequence_to_unitary(seq, num_qubits)
            b_idx = sum(int(snap[i]) << (num_qubits - 1 - i) for i in range(num_qubits))
            v = U[b_idx, :]
            diag = (v @ O @ v.conjugate()).real
            vals.append((dim + 1) * diag - trace_O)
        means.append(np.mean(vals))

    return float(np.median(means))


estimate_clifford = estimate_observable_clifford(
    snapshots_clifford, clifford_sequences, observable, div=10, num_qubits=NUM_QUBITS
)
print(f"Estimated <ZZ> (Clifford shadow): {estimate_clifford:.1f}")
print(f"Error relative to ground truth value of 1.0: {(1.0 - estimate_clifford)}")

def error_bound_clifford(error, observables, delta):
    """
    Calculates the minimum Clifford-shadow size guaranteeing additive error
    `error` on every estimated observable with probability `>= 1 - delta`.

Implements Eq. (5) with the Clifford-ensemble shadow-norm upper bound
        || O - tr(O)/2**n * I ||^2_shadow  <=  3 * tr( (O - tr(O)/2**n * I)^2 ).

    Args:
        error (float): Desired additive error per observable.
        observables (list[SparsePauliOp]): Observables to estimate.
        delta (float): Probability of failure.

    Returns:
        samples (tuple(int, int)): (N*K, K) — total snapshots required and the
        number of median-of-means divisions.
    """
    observable_ops = [hamiltonian_to_matrix(term) for term in observables]

    M = len(observable_ops)
    K = 2 * np.log(2 * M / delta)

    shadow_norms = []
    for op in observable_ops:
        n = int(np.log2(op.shape[0]))
        traceless = op - np.trace(op) / 2**n * np.eye(2**n)
        hs_sq = float(np.real(np.trace(traceless.conj().T @ traceless)))
        shadow_norms.append(3 * hs_sq)

    N = 34 * max(shadow_norms) / error**2

    return int(np.ceil(N * K)), int(K)


required_snapshots_clifford = error_bound_clifford(0.2, [observable], 0.01)
print(
    f"Bound predicts {required_snapshots_clifford[0]} Clifford samples are required to achieve 0.2 error\n"
    f"median-of-means divisions: {required_snapshots_clifford[1]}"
)