
import numpy as np
from .quantum import QuantumCircuit

def grover(num_qubits, target, iterations=None):
    """
    Grover's search algorithm — finds target state with high probability.

    Usage:
        result = grover(num_qubits=3, target=5)
        print(result["outcome"])
        print(result["success"])
    """
    n = num_qubits
    qc = QuantumCircuit(n)

    # Optimal number of iterations
    if iterations is None:
        iterations = max(1, int(np.round(np.pi/4 * np.sqrt(2**n))))

    # Step 1: Hadamard on all qubits (uniform superposition)
    for i in range(n):
        qc.h(i)

    target_bits = format(target, f"0{n}b")

    for _ in range(iterations):
        # Step 2: Oracle — flips phase of target state
        # Mark target by flipping X on qubits where target bit is 0
        for i, bit in enumerate(target_bits):
            if bit == "0":
                qc.x(i)
        # Multi-controlled Z (via H + CCX decomposition)
        qc.h(n-1)
        if n == 2:
            qc.cnot(0, 1)
        elif n >= 3:
            qc.toffoli(0, 1, n-1)
        qc.h(n-1)
        for i, bit in enumerate(target_bits):
            if bit == "0":
                qc.x(i)

        # Step 3: Diffusion operator (inversion about average)
        for i in range(n):
            qc.h(i)
            qc.x(i)
        qc.h(n-1)
        if n == 2:
            qc.cnot(0, 1)
        elif n >= 3:
            qc.toffoli(0, 1, n-1)
        qc.h(n-1)
        for i in range(n):
            qc.x(i)
            qc.h(i)

    outcome, probs = qc.measure()
    return {
        "circuit":  qc,
        "outcome":  outcome,
        "probs":    probs,
        "target":   target,
        "success":  (outcome == target),
        "target_prob": float(probs[target]),
    }


def quantum_phase_estimation(unitary_func, num_estimation_qubits, num_target_qubits):
    """
    Quantum Phase Estimation (QPE).
    Estimates the phase phi in U|psi> = e^(2*pi*i*phi)|psi>.

    unitary_func: a function that takes (circuit, control_qubit, target_qubits)
                  and applies the controlled-U gate.

    Usage:
        def my_U(qc, ctrl, targets):
            qc.cz(ctrl, targets[0])

        result = quantum_phase_estimation(my_U, num_estimation_qubits=3, num_target_qubits=1)
        print(result["phase"])
    """
    n_est = num_estimation_qubits
    n_tgt = num_target_qubits
    n_total = n_est + n_tgt
    qc = QuantumCircuit(n_total)

    est_qubits = list(range(n_est))
    tgt_qubits = list(range(n_est, n_total))

    # Superposition on estimation register
    for q in est_qubits:
        qc.h(q)

    # Prepare target in |1> (eigenstate of many common unitaries)
    qc.x(tgt_qubits[0])

    # Controlled-U^(2^k) for each estimation qubit
    for k, ctrl in enumerate(est_qubits):
        for _ in range(2**k):
            unitary_func(qc, ctrl, tgt_qubits)

    # Inverse QFT on estimation register
    for i in range(n_est-1, -1, -1):
        qc.h(est_qubits[i])
        for j in range(i-1, -1, -1):
            angle = -np.pi / (2**(i-j))
            qc.rz(est_qubits[i], angle/2)
            qc.cnot(est_qubits[i], est_qubits[j])
            qc.rz(est_qubits[j], -angle/2)
            qc.cnot(est_qubits[i], est_qubits[j])
            qc.rz(est_qubits[j], angle/2)

    outcome, probs = qc.measure()
    # Read phase from estimation register bits only
    est_outcome = outcome >> n_tgt
    phase = est_outcome / (2**n_est)

    return {
        "circuit":  qc,
        "phase":    phase,
        "outcome":  est_outcome,
        "probs":    probs,
    }


def amplitude_amplification(oracle_func, diffusion_func, num_qubits, iterations=None):
    """
    General Amplitude Amplification (generalized Grover).

    oracle_func:    function(circuit) that marks good states
    diffusion_func: function(circuit) that reflects about initial state

    Usage:
        def oracle(qc):
            qc.z(0)   # marks |1> on qubit 0

        def diffuse(qc):
            for i in range(qc.num_qubits): qc.h(i); qc.x(i)
            qc.h(qc.num_qubits-1)
            qc.cnot(0, qc.num_qubits-1)
            qc.h(qc.num_qubits-1)
            for i in range(qc.num_qubits): qc.x(i); qc.h(i)

        result = amplitude_amplification(oracle, diffuse, num_qubits=2, iterations=1)
    """
    n = num_qubits
    qc = QuantumCircuit(n)

    if iterations is None:
        iterations = max(1, int(np.pi/4 * np.sqrt(2**n)))

    for i in range(n):
        qc.h(i)

    for _ in range(iterations):
        oracle_func(qc)
        diffusion_func(qc)

    outcome, probs = qc.measure()
    return {
        "circuit": qc,
        "outcome": outcome,
        "probs":   probs,
    }


import math
import random

class ShorResult:
    """Returned by shor(). Contains factors + full step log."""

    def __init__(self, N, factors, steps, circuit=None):
        self.N       = N
        self.factors = factors
        self.steps   = steps
        self.circuit = circuit
        self.success = factors is not None

    def __repr__(self):
        if self.success:
            return (f"ShorResult(N={self.N}, "
                    f"factors={self.factors[0]} x {self.factors[1]})")
        return f"ShorResult(N={self.N}, factors=None [try again])"

    def summary(self):
        sep = "=" * 52
        lines = [sep, f"  Shor Algorithm — N = {self.N}", sep]
        for s in self.steps:
            lines.append("  " + s)
        if self.success:
            lines.append("")
            lines.append(f"  PASS  {self.N} = {self.factors[0]} x {self.factors[1]}")
        else:
            lines.append("")
            lines.append("  WARN  Could not factor in this attempt — rerun.")
        lines.append(sep)
        return "\n".join(lines)


def _build_qpe_circuit_safe(N):
    """
    Build a SMALL symbolic QPE circuit for inspection.
    Uses at most 8 counting qubits + 4 target qubits = 12 qubits MAX.
    This is purely for display/visualization — the actual order-finding
    is done classically below (as it would be simulated anyway).
    """
    from .quantum import QuantumCircuit

    # Hard cap: max 6 counting + 4 target = 10 qubits
    # 2^10 x 16 bytes = 16 KB — completely safe
    n_count  = min(6, N.bit_length())
    n_target = min(4, max(2, N.bit_length() // 2))
    total    = n_count + n_target

    qc = QuantumCircuit(total)

    # Counting register — Hadamard
    for q in range(n_count):
        qc.h(q)

    # Target register — init to |1>
    qc.x(n_count)

    # Symbolic controlled-U gates (logged, not actually applied)
    for k in range(n_count):
        power = 2 ** k
        qc.gate_history.append(
            f"Controlled-U^{power} ctrl={k} target={n_count}..{total-1}"
        )

    # Inverse QFT on counting register (symbolic)
    for i in range(n_count // 2):
        qc.swap(i, n_count - 1 - i)
    for j in range(n_count):
        qc.h(j)
        for k in range(j + 1, n_count):
            angle = -math.pi / (2 ** (k - j))
            qc.rz(angle, j)

    return qc, n_count, n_target


def _find_order(a, N):
    """
    Classically find the multiplicative order of a mod N.
    This is what QPE computes on real hardware.
    Returns r such that a^r == 1 (mod N), or None.
    """
    if math.gcd(a, N) != 1:
        return None
    r, val = 1, a % N
    while val != 1:
        val = (val * a) % N
        r  += 1
        if r > 10 * N:
            return None
    return r


def shor(N, max_attempts=20, verbose=True):
    """
    Factor N using Shor\'s Algorithm.

    Parameters
    ----------
    N            : int  — composite number to factor (must be > 3, odd)
    max_attempts : int  — random base retries (default 20)
    verbose      : bool — print step log

    Returns
    -------
    ShorResult with .factors, .steps, .circuit, .success
    """
    steps = []

    def log(msg):
        steps.append(msg)
        if verbose:
            print("  " + msg)

    # ── 0. Sanity checks ────────────────────────────────────
    if N < 4:
        log(f"N={N} is too small.")
        return ShorResult(N, None, steps)

    if N % 2 == 0:
        log(f"N={N} is even — trivial factor: 2 x {N//2}")
        return ShorResult(N, (2, N // 2), steps)

    # Perfect power check
    for k in range(2, N.bit_length() + 1):
        p = round(N ** (1 / k))
        for c in [p - 1, p, p + 1]:
            if c >= 2 and c ** k == N:
                log(f"N={N} is a perfect power: {c}^{k}")
                return ShorResult(N, (c, N // c), steps)

    log(f"Input N = {N}  ({N.bit_length()} bits)")

    # ── 1. Build small symbolic QPE circuit ─────────────────
    qpe_circuit, n_count, n_target = _build_qpe_circuit_safe(N)
    log(f"QPE circuit (symbolic): {n_count} counting + {n_target} target = "
        f"{qpe_circuit.num_qubits} qubits, {len(qpe_circuit.gate_history)} gates")

    # ── 2. Shor main loop ────────────────────────────────────
    for attempt in range(1, max_attempts + 1):
        log(f"Attempt {attempt}/{max_attempts}")

        a = random.randint(2, N - 1)
        log(f"  base a = {a}")

        g = math.gcd(a, N)
        if g != 1:
            log(f"  gcd({a},{N}) = {g} — lucky classical shortcut!")
            return ShorResult(N, (g, N // g), steps, qpe_circuit)

        r = _find_order(a, N)
        if r is None:
            log(f"  order not found — skip")
            continue

        log(f"  period r = {r}")

        if r % 2 != 0:
            log(f"  r is odd — retry")
            continue

        half = pow(a, r // 2, N)
        if half == N - 1:
            log(f"  a^(r/2) mod N = -1 — trivial, retry")
            continue

        f1 = math.gcd(half + 1, N)
        f2 = math.gcd(half - 1, N)
        log(f"  gcd(a^(r/2)+1, N) = {f1}")
        log(f"  gcd(a^(r/2)-1, N) = {f2}")

        for factor in [f1, f2]:
            if 1 < factor < N:
                other = N // factor
                log(f"  Non-trivial factor: {factor} x {other} = {N}")
                return ShorResult(N, (factor, other), steps, qpe_circuit)

        log(f"  Trivial factors — retry")

    log(f"Could not factor {N} in {max_attempts} attempts.")
    return ShorResult(N, None, steps, qpe_circuit)
