
import numpy as np

class BVOracle:
    """
    The black-box oracle for Bernstein-Vazirani.

    Encodes secret string s. Applies CNOT from input qubit i
    to ancilla wherever s[i] == "1", implementing f(x) = s·x mod 2.

    Usage:
        oracle = BVOracle("101")
        oracle.apply(circuit, input_qubits=[0,1,2], ancilla=3)
    """
    def __init__(self, secret: str):
        if not all(c in "01" for c in secret):
            raise ValueError(f"Secret must be a binary string. Got: {secret!r}")
        self.secret = secret
        self.n = len(secret)

    def apply(self, circuit, input_qubits, ancilla):
        if len(input_qubits) != self.n:
            raise ValueError(f"Need {self.n} input qubits, got {len(input_qubits)}.")
        for i, bit in enumerate(self.secret):
            if bit == "1":
                circuit.cnot(input_qubits[i], ancilla)

    def __repr__(self):
        return f"BVOracle(secret={self.secret!r}, n={self.n})"


class BernsteinVazirani:
    """
    Bernstein-Vazirani Algorithm.
    Finds a hidden binary string in exactly ONE oracle query.

    Usage:
        bv     = BernsteinVazirani("1011")
        result = bv.run()
        result.explain()
    """
    def __init__(self, secret: str, shots: int = 1024):
        if not all(c in "01" for c in secret):
            raise ValueError(f"Secret must be a binary string. Got: {secret!r}")
        self.secret = secret
        self.shots  = shots
        self.n      = len(secret)

    def run(self):
        from pktron.quantum import QuantumCircuit
        from pktron.gpu_primitives import GPUBatchedSampler

        n       = self.n
        ancilla = n
        total   = n + 1

        qc = QuantumCircuit(total)

        # Step 1: ancilla → |1>
        qc.x(ancilla)

        # Step 2: H on all qubits
        for q in range(total):
            qc.h(q)

        # Step 3: oracle
        BVOracle(self.secret).apply(qc, list(range(n)), ancilla)

        # Step 4: H on input qubits only
        for q in range(n):
            qc.h(q)

        # Step 5: sample using GPUBatchedSampler (reads gate_history correctly)
        sampler = GPUBatchedSampler(shots=self.shots)
        raw     = sampler.run(qc)
        counts  = raw["counts"]

        # Most frequent bitstring → first n bits are the secret
        top     = max(counts, key=counts.get)
        found   = top[:n]

        return BVResult(
            secret_hidden = self.secret,
            secret_found  = found,
            counts        = counts,
            circuit       = qc,
            shots         = self.shots,
        )

    def __repr__(self):
        return f"BernsteinVazirani(secret={self.secret!r}, shots={self.shots})"


class BVResult:
    """
    Result of a Bernstein-Vazirani run.

    Attributes:
        secret_hidden : str  — the original hidden string
        secret_found  : str  — what the algorithm recovered
        correct       : bool — True if they match
        counts        : dict — raw measurement counts
        circuit_depth : int  — number of gates applied
        shots         : int  — shots used
    """
    def __init__(self, secret_hidden, secret_found, counts, circuit, shots):
        self.secret_hidden = secret_hidden
        self.secret_found  = secret_found
        self.correct       = (secret_hidden == secret_found)
        self.counts        = counts
        self.circuit_depth = len(circuit.gate_history)
        self.shots         = shots
        self._circuit      = circuit

    def top_results(self, k=5):
        """Return the k most common measurement outcomes."""
        return sorted(self.counts.items(), key=lambda x: x[1], reverse=True)[:k]

    def success_probability(self):
        """Fraction of shots that gave the correct answer."""
        n = len(self.secret_hidden)
        correct_shots = sum(
            v for bs, v in self.counts.items()
            if bs[:n] == self.secret_hidden
        )
        return correct_shots / self.shots

    def explain(self):
        n = len(self.secret_hidden)
        print("=" * 54)
        print("  Bernstein-Vazirani Algorithm — Result")
        print("=" * 54)
        print(f"  Hidden secret string : {self.secret_hidden}")
        print(f"  Algorithm recovered  : {self.secret_found}")
        print(f"  Correct?             : {'✅ YES' if self.correct else '❌ NO'}")
        print(f"  Success probability  : {self.success_probability():.1%}")
        print(f"  Shots used           : {self.shots}")
        print(f"  Circuit gates used   : {self.circuit_depth}")
        print(f"  Qubits used          : {n} input + 1 ancilla = {n+1} total")
        print()
        print("  How it worked (step by step):")
        print(f"   1. Ancilla qubit set to |1>  (X gate)")
        print(f"   2. Hadamard on all {n+1} qubits → superposition")
        print(f"   3. Oracle queried ONCE:")
        cnot_positions = [i for i,b in enumerate(self.secret_hidden) if b=="1"]
        if cnot_positions:
            print(f"      CNOTs at qubit(s) {cnot_positions} → ancilla")
        else:
            print(f"      No CNOTs (secret is all zeros)")
        print(f"   4. Hadamard on {n} input qubits → interference")
        print(f"   5. Measure → secret appears in first {n} bits")
        print()
        print("  Top measurement outcomes:")
        for bs, cnt in self.top_results(5):
            pct    = cnt / self.shots * 100
            marker = " ← secret" if bs[:n] == self.secret_hidden else ""
            print(f"    |{bs}⟩  {cnt:>5} shots  ({pct:5.1f}%){marker}")
        print("=" * 54)

    def __repr__(self):
        status = "✅" if self.correct else "❌"
        return (f"BVResult({status} found={self.secret_found!r}, "
                f"hidden={self.secret_hidden!r}, "
                f"p_success={self.success_probability():.1%})")
