
import numpy as np

# ─────────────────────────────────────────────────────────────────
#  Quantum Counting Algorithm
#
#  Problem:
#    Given a search space of N = 2^n items, how many of them
#    satisfy a condition (are "marked")? Call this count M.
#
#    Classically: need O(N) queries to count M solutions.
#    Quantum Counting: needs only O(sqrt(N/M)) queries.
#
#  How it works:
#    Quantum Counting = Quantum Phase Estimation (QPE)
#    applied to the Grover operator G.
#
#    The Grover operator G has eigenvalues e^(±2iθ) where
#        sin²(θ) = M/N
#    So if we estimate the phase θ via QPE, we recover M:
#        M = N · sin²(θ)
#
#  Three layers in this module:
#    1. GroverOracle       — marks solutions via phase flip
#    2. QuantumCounting    — full QPE-based counting algorithm
#    3. QuantumCountResult — result with explain() and statistics
#
#  Bonus classes:
#    AmplitudeEstimator    — estimates amplitude directly
#    MonteCarloBaseline    — classical baseline for comparison
# ─────────────────────────────────────────────────────────────────


# ── Utility: QFT inverse (for QPE readout) ────────────────────────

def _inverse_qft_matrix(m):
    """
    Return the m x m inverse QFT matrix.
    Used to decode the phase register after QPE.
    """
    omega = np.exp(-2j * np.pi / m)
    idx   = np.arange(m)
    return (1/np.sqrt(m)) * (omega ** np.outer(idx, idx))


# ── Utility: statevector simulator (self-contained, no _ops) ──────

def _simulate_sv(n, gate_ops):
    """
    Minimal statevector simulator.
    gate_ops: list of (gate_name, qubits, params)
    Returns probability vector.
    """
    dim = 2**n
    sv  = np.zeros(dim, dtype=complex)
    sv[0] = 1.0

    I  = np.eye(2, dtype=complex)
    H  = np.array([[1,1],[1,-1]], dtype=complex) / np.sqrt(2)
    X  = np.array([[0,1],[1,0]], dtype=complex)
    Z  = np.array([[1,0],[0,-1]], dtype=complex)

    def apply_1q(sv, G, q, n):
        sv = sv.reshape([2]*n)
        sv = np.tensordot(G, sv, axes=[[1],[q]])
        order = list(range(1,q+1)) + [0] + list(range(q+1,n))
        return np.transpose(sv, order).reshape(dim)

    def apply_2q(sv, G4, q0, q1, n):
        sv = sv.reshape([2]*n)
        G  = G4.reshape(2,2,2,2)
        sv = np.tensordot(G, sv, axes=[[2,3],[q0,q1]])
        remaining = [i for i in range(n) if i not in (q0,q1)]
        perm = [None]*n
        perm[q0]=0; perm[q1]=1
        ri=2
        for pos in remaining:
            perm[pos]=ri; ri+=1
        return np.transpose(sv, perm).reshape(dim)

    CNOT = np.array([[1,0,0,0],[0,1,0,0],[0,0,0,1],[0,0,1,0]],dtype=complex)
    CZ   = np.array([[1,0,0,0],[0,1,0,0],[0,0,1,0],[0,0,0,-1]],dtype=complex)

    for gate, qubits, params in gate_ops:
        if gate == "H":
            sv = apply_1q(sv, H, qubits[0], n)
        elif gate == "X":
            sv = apply_1q(sv, X, qubits[0], n)
        elif gate == "Z":
            sv = apply_1q(sv, Z, qubits[0], n)
        elif gate == "PHASE":
            # Phase gate: |1> → e^(iθ)|1>
            theta = params[0]
            P = np.array([[1,0],[0,np.exp(1j*theta)]], dtype=complex)
            sv = apply_1q(sv, P, qubits[0], n)
        elif gate == "CPHASE":
            # Controlled-phase on (control, target)
            theta = params[0]
            mat = np.eye(4,dtype=complex)
            mat[3,3] = np.exp(1j*theta)
            sv = apply_2q(sv, mat, qubits[0], qubits[1], n)
        elif gate == "CNOT":
            sv = apply_2q(sv, CNOT, qubits[0], qubits[1], n)
        elif gate == "CZ":
            sv = apply_2q(sv, CZ, qubits[0], qubits[1], n)
        elif gate == "RZ":
            theta = params[0]
            RZ = np.array([[np.exp(-1j*theta/2),0],
                           [0,np.exp(1j*theta/2)]], dtype=complex)
            sv = apply_1q(sv, RZ, qubits[0], n)
        elif gate == "ORACLE_PHASE":
            # Mark solution states with a phase flip (Z on each marked state)
            marked = params  # list of marked state indices
            for idx in marked:
                # Flip phase of basis state |idx>
                sv[idx] *= -1

    probs = np.abs(sv)**2
    probs = np.clip(probs, 0, None)
    if probs.sum() > 0:
        probs /= probs.sum()
    return sv, probs


# ── GroverOracle ──────────────────────────────────────────────────

class GroverOracle:
    """
    Grover oracle — marks a set of solution states with a phase flip.

    Two ways to specify solutions:
      (a) Provide a list of marked state indices directly.
      (b) Provide a boolean function f(x) → True/False.

    Parameters
    ----------
    n        : int
        Number of search qubits (search space = 2^n).
    marked   : list[int] or None
        Indices of marked (solution) states. E.g. [3, 7].
    condition: callable or None
        A function f(x: int) → bool. Used if marked is None.
        E.g. condition=lambda x: x % 3 == 0

    Usage
    -----
        oracle = GroverOracle(n=3, marked=[5, 6])
        oracle = GroverOracle(n=4, condition=lambda x: bin(x).count("1")==2)
        ops    = oracle.get_phase_ops()   # list of gate ops
        print(oracle.true_count)          # how many are marked
    """

    def __init__(self, n, marked=None, condition=None):
        if marked is None and condition is None:
            raise ValueError("Provide either `marked` or `condition`.")
        self.n = n
        if marked is not None:
            self._marked = list(marked)
        else:
            self._marked = [x for x in range(2**n) if condition(x)]

        if not self._marked:
            raise ValueError("No states are marked. Count would be 0.")
        if any(m >= 2**n or m < 0 for m in self._marked):
            raise ValueError(f"Marked indices must be in [0, {2**n-1}].")

    @property
    def true_count(self):
        """The true number of marked solutions M."""
        return len(self._marked)

    @property
    def true_fraction(self):
        """M / N — the true fraction of marked states."""
        return self.true_count / 2**self.n

    def get_phase_ops(self):
        """Return gate ops that flip the phase of every marked state."""
        return [("ORACLE_PHASE", [], list(self._marked))]

    def get_diffusion_ops(self, qubits):
        """
        Return gate ops for the Grover diffusion operator on `qubits`.
        Diffusion = 2|s><s| - I  where |s> is uniform superposition.
        Implemented as: H^n · (2|0><0|-I) · H^n
        """
        n   = len(qubits)
        ops = []
        # H on all
        for q in qubits:
            ops.append(("H", [q], []))
        # X on all
        for q in qubits:
            ops.append(("X", [q], []))
        # Multi-controlled Z (approximated as CZ chain for n<=4,
        # or phase kickback for larger n)
        if n == 1:
            ops.append(("Z", [qubits[0]], []))
        elif n == 2:
            ops.append(("CZ", [qubits[0], qubits[1]], []))
        else:
            # Toffoli-chain approximation: CZ on last two, rest are controls
            # For simulation we apply a direct phase flip on |11...1> state
            ops.append(("MCPHASE", qubits, [np.pi]))
        # X on all
        for q in qubits:
            ops.append(("X", [q], []))
        # H on all
        for q in qubits:
            ops.append(("H", [q], []))
        return ops

    def __repr__(self):
        return (f"GroverOracle(n={self.n}, marked={self._marked}, "
                f"true_count={self.true_count})")


# ── QuantumCounting (QPE-based) ───────────────────────────────────

class QuantumCounting:
    """
    Quantum Counting via Quantum Phase Estimation on the Grover operator.

    The Grover operator G = Diffusion · Oracle has eigenvalues e^(±2iθ)
    where sin²(θ) = M/N. QPE estimates θ, then we compute M = N·sin²(θ).

    Parameters
    ----------
    oracle       : GroverOracle
        The oracle defining which states are marked.
    phase_qubits : int
        Number of qubits in the phase estimation register.
        More qubits → higher precision. Default 4.
        Precision in M: ±N·sin(π/2^phase_qubits)² roughly.
    shots        : int
        Measurement shots per run. Default 2048.

    Usage
    -----
        oracle = GroverOracle(n=4, marked=[3,7,11])
        qc_obj = QuantumCounting(oracle, phase_qubits=5)
        result = qc_obj.run()
        result.explain()
    """

    def __init__(self, oracle: GroverOracle, phase_qubits: int = 4,
                 shots: int = 2048):
        self.oracle        = oracle
        self.phase_qubits  = phase_qubits
        self.shots         = shots
        self.n             = oracle.n       # search register size
        self.N             = 2**oracle.n   # search space size

    def _build_grover_unitary(self):
        """
        Build the full 2^n × 2^n Grover operator matrix G = D · O.
        D = diffusion, O = oracle (phase flip on marked states).
        This is used for the controlled-G gates in QPE.
        """
        N = self.N
        n = self.n

        # Oracle matrix: flip phase of marked states
        O = np.eye(N, dtype=complex)
        for m in self.oracle._marked:
            O[m, m] = -1

        # Diffusion matrix: D = 2|s><s| - I
        s    = np.ones(N, dtype=complex) / np.sqrt(N)
        D    = 2 * np.outer(s, s.conj()) - np.eye(N, dtype=complex)

        # Grover = D · O
        return D @ O

    def _qpe_simulation(self, G):
        """
        Simulate QPE using the Grover matrix G directly.
        Returns estimated phase θ in [0, 2π).
        """
        m   = self.phase_qubits
        M_p = 2**m              # phase register dimension

        # Eigendecompose G to get phases
        eigenvalues, eigenvectors = np.linalg.eig(G)
        phases = np.angle(eigenvalues) % (2*np.pi)  # in [0, 2π)

        # The Grover eigenvalues are e^(±2iθ)
        # |s> (uniform state) has overlap with both eigenvectors
        s        = np.ones(self.N, dtype=complex) / np.sqrt(self.N)
        overlaps = np.abs(eigenvectors.conj().T @ s)**2

        # Simulate the QPE measurement distribution
        # Each eigenvalue contributes with probability |<eig|s>|²
        # and maps to phase bin k = round(phi/(2π) * M_p)
        probs = np.zeros(M_p)
        for phase, weight in zip(phases, overlaps):
            k = int(round(phase / (2*np.pi) * M_p)) % M_p
            probs[k] += weight

        # Normalize
        if probs.sum() > 0:
            probs /= probs.sum()

        # Sample from phase register
        outcomes = np.random.choice(M_p, self.shots, p=probs)
        counts   = {}
        for o in outcomes:
            counts[o] = counts.get(o, 0) + 1

        # Most likely phase bin
        top_k = max(counts, key=counts.get)
        theta_est = top_k / M_p * 2 * np.pi

        return theta_est, counts, phases, overlaps

    def run(self):
        """
        Run Quantum Counting and return a QuantumCountResult.
        """
        G                         = self._build_grover_unitary()
        theta_est, phase_counts, true_phases, overlaps = self._qpe_simulation(G)

        # From the QPE phase θ, compute M estimate
        # Eigenvalues of G are e^(±2iθ) where sin²(θ) = M/N
        # The smaller phase (closer to 0) gives θ directly
        # We need to handle the two conjugate eigenvalues

        # Try both θ and π - θ (the two Grover phases)
        # and pick the one that gives a count in [1, N]
        N   = self.N
        candidates = []
        for phase in [theta_est, 2*np.pi - theta_est]:
            theta = phase / 2       # because eigenvalue = e^(2iθ)
            M_est = N * (np.sin(theta)**2)
            if 0 <= M_est <= N:
                candidates.append((abs(M_est - round(M_est)), round(M_est), theta, M_est))

        candidates.sort()
        if candidates:
            _, M_int, theta_used, M_float = candidates[0]
        else:
            theta_used = theta_est / 2
            M_float    = N * np.sin(theta_used)**2
            M_int      = int(round(M_float))

        # True eigenvalues for reference
        true_theta = np.arcsin(np.sqrt(self.oracle.true_count / N))

        return QuantumCountResult(
            true_count    = self.oracle.true_count,
            estimated_count = M_int,
            estimated_float = M_float,
            theta_estimated = theta_used,
            theta_true      = true_theta,
            N               = N,
            n               = self.n,
            phase_qubits    = self.phase_qubits,
            shots           = self.shots,
            phase_counts    = phase_counts,
            oracle          = self.oracle,
        )

    def __repr__(self):
        return (f"QuantumCounting(n={self.n}, "
                f"phase_qubits={self.phase_qubits}, "
                f"N={self.N})")


# ── QuantumCountResult ────────────────────────────────────────────

class QuantumCountResult:
    """
    Result of the Quantum Counting algorithm.

    Attributes
    ----------
    true_count       : int   — actual number of solutions (M)
    estimated_count  : int   — what the algorithm estimated (M̂)
    correct          : bool  — True if estimated_count == true_count
    relative_error   : float — |M̂ - M| / M
    theta_estimated  : float — estimated Grover phase θ
    theta_true       : float — true Grover phase θ
    N                : int   — total search space size (2^n)
    """

    def __init__(self, true_count, estimated_count, estimated_float,
                 theta_estimated, theta_true, N, n, phase_qubits,
                 shots, phase_counts, oracle):
        self.true_count      = true_count
        self.estimated_count = estimated_count
        self.estimated_float = estimated_float
        self.correct         = (estimated_count == true_count)
        self.relative_error  = (abs(estimated_count - true_count) / true_count
                                if true_count > 0 else 0.0)
        self.theta_estimated = theta_estimated
        self.theta_true      = theta_true
        self.N               = N
        self.n               = n
        self.phase_qubits    = phase_qubits
        self.shots           = shots
        self.phase_counts    = phase_counts
        self._oracle         = oracle

    def precision_bits(self):
        """Effective bits of precision in the count estimate."""
        return self.phase_qubits

    def confidence(self):
        """
        Fraction of shots that fell in the correct phase bin
        (or adjacent bins within 1 step).
        """
        M_p     = 2**self.phase_qubits
        N       = self.N
        theta_t = self.theta_true
        k_true  = int(round(2 * theta_t / (2*np.pi) * M_p)) % M_p
        nearby  = {(k_true - 1) % M_p, k_true, (k_true + 1) % M_p}
        hits    = sum(self.phase_counts.get(k, 0) for k in nearby)
        return hits / self.shots

    def speedup_over_classical(self):
        """
        Quantum counting queries: O(sqrt(N/M) * phase_qubits)
        Classical queries:        O(N)
        Returns the approximate speedup factor.
        """
        M = max(self.true_count, 1)
        quantum_queries  = int(np.sqrt(self.N / M) * self.phase_qubits)
        classical_queries = self.N
        return classical_queries / max(quantum_queries, 1)

    def explain(self):
        print("=" * 58)
        print("  Quantum Counting Algorithm — Result")
        print("=" * 58)
        print(f"  Search space N      : {self.N}  ({self.n} qubits)")
        print(f"  True solution count : {self.true_count}  (M)")
        print(f"  Estimated count     : {self.estimated_count}  (M̂)")
        print(f"  Correct?            : {'✅ YES' if self.correct else '❌ NO'}")
        print(f"  Relative error      : {self.relative_error:.1%}")
        print(f"  Confidence          : {self.confidence():.1%}")
        print()
        print(f"  Phase estimation:")
        print(f"   Phase qubits used  : {self.phase_qubits}")
        print(f"   True θ             : {self.theta_true:.6f} rad")
        print(f"   Estimated θ        : {self.theta_estimated:.6f} rad")
        print(f"   sin²(θ_true)       : {np.sin(self.theta_true)**2:.6f}")
        print(f"   sin²(θ_est)        : {np.sin(self.theta_estimated)**2:.6f}")
        print(f"   M = N·sin²(θ_est)  : {self.estimated_float:.4f} → {self.estimated_count}")
        print()
        print(f"  Performance:")
        print(f"   Classical queries  : O({self.N})")
        quantum_q = int(np.sqrt(self.N/max(self.true_count,1)) * self.phase_qubits)
        print(f"   Quantum queries    : O({quantum_q})  "
              f"[√(N/M) × phase_qubits]")
        print(f"   Speedup            : ~{self.speedup_over_classical():.1f}x")
        print()
        print(f"  Top phase register outcomes:")
        M_p   = 2**self.phase_qubits
        top5  = sorted(self.phase_counts.items(),
                       key=lambda x: x[1], reverse=True)[:5]
        for k, cnt in top5:
            theta_k = k / M_p * np.pi    # θ = phase/2
            M_k     = self.N * np.sin(theta_k)**2
            pct     = cnt/self.shots*100
            print(f"   bin={k:>3}  θ≈{theta_k:.4f}  M≈{M_k:6.2f}  "
                  f"{cnt:>5} shots ({pct:.1f}%)")
        print("=" * 58)

    def __repr__(self):
        status = "✅" if self.correct else "❌"
        return (f"QuantumCountResult({status} "
                f"estimated={self.estimated_count}, "
                f"true={self.true_count}, "
                f"error={self.relative_error:.1%}, "
                f"N={self.N})")


# ── AmplitudeEstimator ────────────────────────────────────────────

class AmplitudeEstimator:
    """
    Estimates the amplitude a = sqrt(M/N) directly via QPE,
    without going through Grover counting.

    Useful when you want the raw amplitude rather than a count.

    Usage:
        ae = AmplitudeEstimator(oracle, phase_qubits=5)
        r  = ae.run()
        print(r)  # {"amplitude": 0.433, "estimated_M": 3, "true_M": 3}
    """
    def __init__(self, oracle: GroverOracle, phase_qubits: int = 5):
        self.oracle       = oracle
        self.phase_qubits = phase_qubits
        self.N            = 2**oracle.n

    def run(self):
        qc  = QuantumCounting(self.oracle, self.phase_qubits, shots=4096)
        res = qc.run()
        a   = np.sin(res.theta_true)
        a_est = np.sin(res.theta_estimated)
        return {
            "true_amplitude"     : float(a),
            "estimated_amplitude": float(a_est),
            "true_M"             : self.oracle.true_count,
            "estimated_M"        : res.estimated_count,
            "relative_error"     : float(res.relative_error),
            "phase_qubits"       : self.phase_qubits,
        }

    def __repr__(self):
        return (f"AmplitudeEstimator(n={self.oracle.n}, "
                f"phase_qubits={self.phase_qubits})")


# ── MonteCarloBaseline ────────────────────────────────────────────

class MonteCarloBaseline:
    """
    Classical Monte Carlo estimate of M for comparison.
    Randomly samples from the search space and counts hits.

    Usage:
        baseline = MonteCarloBaseline(oracle, samples=500)
        result   = baseline.run()
        print(result)
    """
    def __init__(self, oracle: GroverOracle, samples: int = 500):
        self.oracle  = oracle
        self.samples = samples
        self.N       = 2**oracle.n

    def run(self):
        marked_set = set(self.oracle._marked)
        hits = sum(
            1 for _ in range(self.samples)
            if np.random.randint(0, self.N) in marked_set
        )
        M_est = int(round(hits / self.samples * self.N))
        return {
            "true_M"           : self.oracle.true_count,
            "estimated_M"      : M_est,
            "samples_used"     : self.samples,
            "hit_rate"         : hits / self.samples,
            "relative_error"   : abs(M_est - self.oracle.true_count) /
                                  max(self.oracle.true_count, 1),
        }

    def __repr__(self):
        return (f"MonteCarloBaseline(N={self.N}, "
                f"samples={self.samples})")
