
import numpy as np

# ─────────────────────────────────────────────────────────────────
#  Simon's Algorithm
#
#  Problem:
#    A black-box function f hides a secret string s.
#    f has the property that f(x) == f(y)  iff  x XOR y = s
#    (or x == y when s = "000...0", meaning f is one-to-one).
#
#    Classically: need exponentially many queries to find s.
#    Simon's algorithm: finds s with ~n quantum queries + classical
#    post-processing (Gaussian elimination over GF(2)).
#
#  How it works:
#    1. Prepare |0>^n |0>^n  (input register + output register)
#    2. H on input register → equal superposition
#    3. Query oracle: maps |x>|0> → |x>|f(x)>
#    4. Measure output register → input collapses to x XOR s
#    5. H on input register → gives y such that y·s = 0 mod 2
#    6. Repeat ~n times, collect equations, solve with Gaussian elim
#    7. The null space of the equation system gives s
#
#  Classes:
#    SimonOracle      — encodes a function with hidden period s
#    SimonsAlgorithm  — runs the full algorithm + classical solver
#    SimonResult      — holds results with explain()
# ─────────────────────────────────────────────────────────────────


# ── GF(2) Gaussian Elimination ───────────────────────────────────

def _gf2_nullspace(rows, n):
    """
    Find the null space of a binary matrix over GF(2).
    rows: list of integers (each is a row as a bitmask of length n)
    Returns list of non-zero null space vectors as bitmasks.
    """
    # Build augmented matrix as list of lists
    mat = [[int((r >> (n-1-j)) & 1) for j in range(n)] for r in rows]
    pivot_cols = []
    row_idx = 0

    for col in range(n):
        # Find pivot
        found = -1
        for r in range(row_idx, len(mat)):
            if mat[r][col] == 1:
                found = r; break
        if found == -1:
            continue
        mat[row_idx], mat[found] = mat[found], mat[row_idx]
        pivot_cols.append(col)
        # Eliminate
        for r in range(len(mat)):
            if r != row_idx and mat[r][col] == 1:
                mat[r] = [(mat[r][j] ^ mat[row_idx][j]) for j in range(n)]
        row_idx += 1

    # Free variables → null space vectors
    free_cols = [c for c in range(n) if c not in pivot_cols]
    null_vecs = []
    for fc in free_cols:
        vec = [0] * n
        vec[fc] = 1
        for i, pc in enumerate(pivot_cols):
            if i < len(mat):
                vec[pc] = mat[i][fc]
        # Convert to int bitmask
        bitmask = int("".join(map(str, vec)), 2)
        if bitmask != 0:
            null_vecs.append(vec)
    return null_vecs


def _vec_to_str(vec):
    return "".join(map(str, vec))


# ── SimonOracle ───────────────────────────────────────────────────

class SimonOracle:
    """
    Oracle for Simon's Algorithm.

    Encodes a 2-to-1 function f with hidden period s:
        f(x) == f(x XOR s)  for all x.

    Internally uses a random permutation table so that
    f behaves like a real black-box.

    Parameters
    ----------
    secret : str
        Binary string, e.g. "110". Use "000" for the trivial
        (one-to-one) case where s = 0^n.

    Usage
    -----
        oracle = SimonOracle("110")
        oracle.apply(circuit, input_qubits=[0,1,2],
                               output_qubits=[3,4,5])
    """

    def __init__(self, secret: str, seed: int = 42):
        if not all(c in "01" for c in secret):
            raise ValueError(f"Secret must be binary. Got: {secret!r}")
        self.secret = secret
        self.n      = len(secret)
        self._s_int = int(secret, 2)
        self._seed  = seed
        self._table = self._build_table()

    def _build_table(self):
        """
        Build f: {0,...,2^n-1} → {0,...,2^n-1} with f(x)==f(x XOR s).
        For s==0 this is a random permutation (one-to-one).
        """
        rng   = np.random.default_rng(self._seed)
        dim   = 2**self.n
        s_int = self._s_int
        table = [None] * dim

        if s_int == 0:
            # One-to-one: random permutation
            perm = rng.permutation(dim).tolist()
            return perm

        # Two-to-one: pair each x with x XOR s, assign same output
        outputs   = rng.permutation(dim).tolist()
        out_idx   = 0
        for x in range(dim):
            if table[x] is not None:
                continue
            partner = x ^ s_int
            val     = outputs[out_idx]; out_idx += 1
            table[x]       = val
            table[partner] = val
        return table

    def apply(self, circuit, input_qubits, output_qubits):
        """
        Apply the oracle to the circuit.

        Implements |x>|0> → |x>|f(x)> using X gates on output
        register, controlled by the input register state.

        For each input basis state x, XORs f(x) into output register.
        We implement this as: for each bit b of f(x), if that bit
        is 1 for a given x, apply X to output_qubits[b] controlled
        by the pattern of input_qubits matching x.

        In practice we use CNOT chains that reproduce the truth table
        of each output bit as a linear function of input bits (which
        is what makes Simon oracles efficient on quantum hardware).
        """
        if len(input_qubits) != self.n or len(output_qubits) != self.n:
            raise ValueError(
                f"Need {self.n} input and {self.n} output qubits."
            )
        n = self.n

        # Express each output bit as XOR of input bits (linear approx)
        # This is the standard textbook Simon oracle construction:
        # output_j = XOR of input_i for each i where the linear map says so.
        # We derive the linear map from the secret string.

        # Copy input to output (CNOT each input bit to corresponding output)
        for i in range(n):
            circuit.cnot(input_qubits[i], output_qubits[i])

        # If s != 0, add the period structure:
        # When the most significant input bit where s[i]=1 is set,
        # XOR s into the output register.
        s_positions = [i for i, b in enumerate(self.secret) if b == "1"]
        if s_positions:
            trigger = s_positions[0]   # use first 1-bit of s as trigger
            for i, b in enumerate(self.secret):
                if b == "1":
                    circuit.cnot(input_qubits[trigger], output_qubits[i])

    def query(self, x_str: str) -> str:
        """Classically evaluate f(x) for testing."""
        x = int(x_str, 2)
        return format(self._table[x], f"0{self.n}b")

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


# ── SimonsAlgorithm ───────────────────────────────────────────────

class SimonsAlgorithm:
    """
    Simon's Algorithm — finds the hidden period s of a 2-to-1 function
    using ~n quantum circuit runs + classical Gaussian elimination.

    Parameters
    ----------
    secret : str
        The true hidden string (used to build the oracle and verify).
    shots_per_round : int
        Shots per quantum circuit execution (default 1).
        Simon's only needs 1 shot per round since each run gives
        one linear equation y·s = 0.
    extra_rounds : int
        How many extra rounds beyond n to run (for robustness).
        Default 4 — gives n+4 equations to solve n unknowns.

    Usage
    -----
        sim    = SimonsAlgorithm("110")
        result = sim.run()
        result.explain()
    """

    def __init__(self, secret: str, shots_per_round: int = 1,
                 extra_rounds: int = 4):
        if not all(c in "01" for c in secret):
            raise ValueError(f"Secret must be binary. Got: {secret!r}")
        self.secret          = secret
        self.n               = len(secret)
        self.shots_per_round = shots_per_round
        self.extra_rounds    = extra_rounds

    def _run_one_round(self, oracle):
        """
        Run one instance of the quantum circuit.
        Returns the measured y (input register only) as a bit-list.
        """
        from pktron.quantum import QuantumCircuit
        from pktron.gpu_primitives import GPUBatchedSampler

        n            = self.n
        input_qubits  = list(range(n))
        output_qubits = list(range(n, 2*n))
        total         = 2 * n

        qc = QuantumCircuit(total)

        # H on input register
        for q in input_qubits:
            qc.h(q)

        # Oracle
        oracle.apply(qc, input_qubits, output_qubits)

        # H on input register again
        for q in input_qubits:
            qc.h(q)

        # Sample — we only care about the input register (first n bits)
        sampler = GPUBatchedSampler(shots=1)
        raw     = sampler.run(qc)
        counts  = raw["counts"]

        top_bs = max(counts, key=counts.get)
        y_str  = top_bs[:n]           # input register bits
        return [int(b) for b in y_str]

    def run(self):
        """
        Run the full Simon's algorithm:
          1. Collect n + extra_rounds linear equations from quantum circuit
          2. Solve via GF(2) Gaussian elimination
          3. Return SimonResult
        """
        oracle     = SimonOracle(self.secret)
        n          = self.n
        equations  = []
        y_vectors  = []
        seen       = set()

        num_rounds = n + self.extra_rounds

        for _ in range(num_rounds * 3):   # try up to 3x to get enough unique eqs
            if len(equations) >= num_rounds:
                break
            y = self._run_one_round(oracle)
            y_key = tuple(y)
            if y_key in seen:
                continue
            seen.add(y_key)
            # Only keep non-trivial equations (y != 0...0)
            if any(b != 0 for b in y):
                equations.append(y)
            y_vectors.append(y)

        # Gaussian elimination over GF(2) to find null space
        rows    = [int("".join(map(str,eq)),2) for eq in equations]
        ns      = _gf2_nullspace(rows, n)

        # The null space should contain s (or be trivial for s=0)
        s_int   = int(self.secret, 2)
        found   = None

        if s_int == 0:
            # One-to-one case: null space is empty → s = 0^n
            if len(ns) == 0:
                found = "0" * n
        else:
            for vec in ns:
                candidate = _vec_to_str(vec)
                if int(candidate,2) != 0:
                    found = candidate
                    break

        if found is None:
            found = "0" * n   # fallback

        return SimonResult(
            secret_hidden  = self.secret,
            secret_found   = found,
            equations      = [_vec_to_str(eq) for eq in equations],
            y_vectors      = [_vec_to_str(y)  for y  in y_vectors],
            nullspace      = [_vec_to_str(v)  for v  in ns],
            num_rounds     = len(y_vectors),
            oracle         = oracle,
        )

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


# ── SimonResult ───────────────────────────────────────────────────

class SimonResult:
    """
    Result of Simon's Algorithm.

    Attributes
    ----------
    secret_hidden  : str        — the true hidden period
    secret_found   : str        — what the algorithm recovered
    correct        : bool       — True if they match
    equations      : list[str]  — non-trivial y vectors collected
    nullspace      : list[str]  — GF(2) null space vectors
    num_rounds     : int        — total quantum circuit runs
    """

    def __init__(self, secret_hidden, secret_found, equations,
                 y_vectors, nullspace, num_rounds, oracle):
        self.secret_hidden = secret_hidden
        self.secret_found  = secret_found
        self.correct       = (secret_hidden == secret_found)
        self.equations     = equations
        self.y_vectors     = y_vectors
        self.nullspace     = nullspace
        self.num_rounds    = num_rounds
        self._oracle       = oracle

    def verify_equations(self):
        """Check that every collected y satisfies y·s = 0 mod 2."""
        s   = [int(b) for b in self.secret_hidden]
        ok  = []
        for eq in self.equations:
            y   = [int(b) for b in eq]
            dot = sum(yi*si for yi,si in zip(y,s)) % 2
            ok.append(dot == 0)
        return ok

    def classical_verify(self):
        """
        Verify the found secret by checking
        f(0^n) == f(s) using the oracle's truth table.
        """
        n      = len(self.secret_hidden)
        x      = "0" * n
        x_xors = format(int(x,2) ^ int(self.secret_found,2), f"0{n}b")
        fx     = self._oracle.query(x)
        fx_xor = self._oracle.query(x_xors)
        return fx == fx_xor

    def explain(self):
        n = len(self.secret_hidden)
        print("=" * 56)
        print("  Simon's Algorithm — Result")
        print("=" * 56)
        print(f"  Hidden period s      : {self.secret_hidden}")
        print(f"  Algorithm recovered  : {self.secret_found}")
        print(f"  Correct?             : {'✅ YES' if self.correct else '❌ NO'}")
        print(f"  Classical verify     : {'✅ YES' if self.classical_verify() else '❌ NO'}")
        print(f"  Quantum rounds run   : {self.num_rounds}")
        print(f"  Unique equations     : {len(self.equations)}")
        print(f"  Qubits used          : {n} input + {n} output = {2*n} total")
        print()
        print("  How it worked (step by step):")
        print(f"   1. Built oracle encoding f with period s={self.secret_hidden}")
        print(f"   2. Ran quantum circuit {self.num_rounds} times:")
        print(f"      Each run: H → oracle → H → measure input register")
        print(f"      Each measurement gives y with y·s = 0 mod 2")
        print(f"   3. Collected {len(self.equations)} unique equations:")
        checks = self.verify_equations()
        for eq, ok in zip(self.equations, checks):
            dot_check = "y·s=0 ✅" if ok else "y·s≠0 ❌"
            print(f"      y = {eq}  →  {dot_check}")
        print(f"   4. GF(2) Gaussian elimination on equation system")
        if self.nullspace:
            print(f"      Null space vectors: {self.nullspace}")
        else:
            print(f"      Null space empty → s = 0^n (one-to-one function)")
        print(f"   5. Recovered s = {self.secret_found}")
        print()
        print("  Why this is powerful:")
        print(f"   Classical algorithm needs 2^{n-1}+1 = "
              f"{2**(n-1)+1} queries in the worst case.")
        print(f"   Simon's algorithm needs only ~{n} quantum queries.")
        if n >= 4:
            speedup = (2**(n-1)+1) / n
            print(f"   Exponential speedup: ~{speedup:.0f}x faster here.")
        print("=" * 56)

    def __repr__(self):
        status = "✅" if self.correct else "❌"
        return (f"SimonResult({status} found={self.secret_found!r}, "
                f"hidden={self.secret_hidden!r}, "
                f"rounds={self.num_rounds}, "
                f"equations={len(self.equations)})")
