
import numpy as np
from fractions import Fraction
from scipy.optimize import minimize
from .gates import h, rx, ry, rz, cnot, cz, s, toffoli, t, x, swap
from .quantum import QuantumCircuit
try:
    import pennylane as qml
    _HAS_PENNYLANE = True
except ImportError:
    qml = None
    _HAS_PENNYLANE = False

def quantum_flip(circuit):
    circuit.h(0)
    if circuit.num_qubits > 1:
        circuit.cnot(0, 1)
    circuit.rx(0, np.pi/4)
    circuit.gate_history.append("Quantum Flip")

def super_quantum_dance(circuit):
    circuit.h(0)
    if circuit.num_qubits > 1:
        circuit.ry(1, np.pi/2)
        circuit.cnot(0, 1)
    if circuit.num_qubits > 2:
        circuit.toffoli(0, 1, 2)
    circuit.t(0)
    circuit.gate_history.append("Super Quantum Dance")

def quantum_treasure_hunt(circuit):
    target_state = 2**circuit.num_qubits - 1
    circuit.h(0)
    if circuit.num_qubits > 1:
        circuit.cnot(0, 1)
        circuit.ry(1, np.random.uniform(0, np.pi))
    circuit.rx(0, np.random.uniform(0, np.pi))
    outcome, probs = circuit.measure()
    score = probs[target_state] * 100
    circuit.gate_history.append(f"Quantum Treasure Hunt (Score: {score:.1f})")
    return score

def quantum_magic_spin(circuit):
    target_state = 0
    circuit.rx(0, np.random.uniform(0, 2*np.pi))
    if circuit.num_qubits > 1:
        circuit.ry(1, np.random.uniform(0, np.pi))
        circuit.cz(0, 1)
    circuit.s(0)
    outcome, probs = circuit.measure()
    score = probs[target_state] * 100
    circuit.gate_history.append(f"Quantum Magic Spin (Score: {score:.1f})")
    return score

def bb84_qkd(num_qubits, eve_intercept_prob=0.5, seed=None):
    if seed is not None:
        np.random.seed(seed)
    qc = QuantumCircuit(num_qubits, "statevector")
    alice_bits = np.random.randint(0, 2, num_qubits)
    alice_bases = np.random.randint(0, 2, num_qubits)
    eve_intercepts = np.random.choice([0,1], size=num_qubits, p=[1-eve_intercept_prob, eve_intercept_prob])
    eve_bases = np.random.randint(0, 2, num_qubits)
    bob_bases = np.random.randint(0, 2, num_qubits)
    for i in range(num_qubits):
        if alice_bits[i] == 1: qc.x(i)
        if alice_bases[i] == 1: qc.h(i)
    for i in range(num_qubits):
        if eve_intercepts[i]:
            if eve_bases[i] == 1: qc.h(i)
            if np.random.randint(0, 2) == 1: qc.x(i)
            if eve_bases[i] == 1: qc.h(i)
    for i in range(num_qubits):
        if bob_bases[i] == 1: qc.h(i)
    outcome, probs = qc.measure()
    measured_bits = [int(b) for b in format(outcome, f"0{qc.num_qubits}b")]
    alice_key, bob_key = [], []
    for i in range(num_qubits):
        if alice_bases[i] == bob_bases[i]:
            alice_key.append(alice_bits[i])
            bob_key.append(measured_bits[i])
    error_count = sum(a != b for a, b in zip(alice_key, bob_key))
    error_rate = error_count / len(alice_key) if alice_key else 0.0
    eve_detected = error_rate > 0.25
    qc.gate_history.append(f"BB84 QKD (Error rate: {error_rate:.2f}, Eve detected: {eve_detected})")
    return {"alice_key": alice_key, "bob_key": bob_key, "error_rate": error_rate,
            "eve_detected": eve_detected, "circuit": qc}

def grover_search(num_qubits, target=None, oracle_func=None, iterations=None, noise_prob=0.0):
    if num_qubits < 2:
        raise ValueError("Grover requires at least 2 qubits")
    qc = QuantumCircuit(num_qubits, "statevector")
    if iterations is None:
        iterations = int(np.pi / 4 * np.sqrt(2**num_qubits))
    def default_oracle():
        target_bin = format(target, f"0{num_qubits}b")
        for i, bit in enumerate(target_bin[::-1]):
            if bit == "0": qc.x(i)
        qc.h(num_qubits-1)
        if num_qubits > 2:
            qc.toffoli(0, 1, num_qubits-1)
        else:
            qc.cnot(0, 1)
        qc.h(num_qubits-1)
        for i, bit in enumerate(target_bin[::-1]):
            if bit == "0": qc.x(i)
    def diffusion():
        for i in range(num_qubits): qc.h(i); qc.x(i)
        qc.h(num_qubits-1)
        if num_qubits > 2:
            qc.toffoli(0, 1, num_qubits-1)
        else:
            qc.cnot(0, 1)
        qc.h(num_qubits-1)
        for i in range(num_qubits): qc.x(i); qc.h(i)
    for i in range(num_qubits): qc.h(i)
    for _ in range(iterations):
        if oracle_func is not None: oracle_func(qc)
        else: default_oracle()
        if np.random.random() < noise_prob: qc.x(np.random.randint(0, num_qubits))
        diffusion()
    outcome, probs = qc.measure()
    qc.gate_history.append(f"Grover Search (Target: {target})")
    return {"circuit": qc, "outcome": outcome, "probs": probs}

def qft(circuit, start=0, end=None):
    if end is None: end = circuit.num_qubits
    for i in range(start, end):
        for j in range(i+1, end):
            lam = np.pi / (2**(j-i))
            circuit.rz(i, lam/2)
            circuit.cnot(i, j)
            circuit.rz(j, -lam/2)
            circuit.cnot(i, j)
            circuit.rz(j, lam/2)
        circuit.h(i)
    for i in range(start, (start+end)//2):
        circuit.swap(i, end-1-(i-start))
    circuit.gate_history.append("Quantum Fourier Transform")

def phase_estimation(circuit, unitary, estimation_qubits, target_qubits):
    for q in estimation_qubits: circuit.h(q)
    for i, eq in enumerate(estimation_qubits):
        for _ in range(2**i): unitary(circuit, control=eq, targets=target_qubits)
    n_est = len(estimation_qubits)
    for i in range(n_est-1, -1, -1):
        for j in range(i):
            lam = -np.pi / (2**(i-j))
            circuit.rz(estimation_qubits[i], lam/2)
            circuit.cnot(estimation_qubits[i], estimation_qubits[j])
            circuit.rz(estimation_qubits[j], -lam/2)
            circuit.cnot(estimation_qubits[i], estimation_qubits[j])
            circuit.rz(estimation_qubits[j], lam/2)
        circuit.h(estimation_qubits[i])
    circuit.gate_history.append("Phase Estimation")
    outcome, probs = circuit.measure()
    return {"circuit": circuit, "phase": outcome/(2**n_est), "probs": probs}

def shor_period_finding(N, a, num_qubits=None):
    if num_qubits is None:
        num_qubits = 2 * int(np.ceil(np.log2(N))) + 1
    if num_qubits % 2 != 0: num_qubits += 1
    qc = QuantumCircuit(num_qubits, "statevector")
    counting_qubits = num_qubits // 2
    register_qubits = num_qubits - counting_qubits
    for i in range(counting_qubits): qc.h(i)
    qc.x(counting_qubits)
    period, factors = None, []
    for _ in range(1000):
        outcome, probs = qc.measure()
        measured = outcome >> register_qubits
        try:
            phase = measured / (2**counting_qubits)
            frac = Fraction(phase).limit_denominator(N)
            r = frac.denominator
            if r != 0 and r < N and pow(a, r, N) == 1:
                cand1 = np.gcd(pow(a, r//2, N)-1, N)
                cand2 = np.gcd(pow(a, r//2, N)+1, N)
                new_factors = [f for f in [cand1, cand2] if 1 < f < N]
                if new_factors:
                    period = r
                    factors = sorted(new_factors)
                    break
        except: pass
    qc.gate_history.append(f"Shor Period Finding (N={N}, a={a}, period={period})")
    return {"circuit": qc, "period": period, "factors": factors}

def vqe(num_qubits, hamiltonian=None, ansatz_depth=3, noise_prob=0.0, max_iter=300):
    if hamiltonian is None:
        hamiltonian = np.diag([1.0 if i==0 else -1.0 for i in range(2**num_qubits)])
    def ansatz(params, circuit):
        param_idx = 0
        for d in range(ansatz_depth):
            for i in range(num_qubits):
                circuit.ry(i, params[param_idx])
                circuit.rz(i, params[param_idx+num_qubits])
                param_idx += 1
            if num_qubits > 1:
                for i in range(num_qubits-1): circuit.cnot(i, i+1)
    def objective(params):
        qc = QuantumCircuit(num_qubits, "statevector")
        ansatz(params, qc)
        state = qc.get_state()
        return float(np.real(state.conj() @ hamiltonian @ state))
    init = np.random.uniform(0, np.pi, 2*num_qubits*ansatz_depth)
    result = minimize(objective, init, method="Powell", options={"maxiter": max_iter})
    qc = QuantumCircuit(num_qubits, "statevector")
    ansatz(result.x, qc)
    energy = objective(result.x)
    qc.gate_history.append(f"VQE (Energy: {energy:.4f})")
    return {"circuit": qc, "energy": energy, "params": result.x}

def qaoa(num_qubits, graph=None, p=1, noise_prob=0.0, max_iter=100):
    if graph is None:
        graph = [(0,1),(1,2)] if num_qubits >= 3 else [(0,1)]
    def cost_h(circuit, gamma):
        for edge in graph: circuit.cz(edge[0], edge[1])
    def mixer_h(circuit, beta):
        for i in range(num_qubits): circuit.rx(i, 2*beta)
    def objective(params):
        qc = QuantumCircuit(num_qubits, "statevector")
        for i in range(num_qubits): qc.h(i)
        for layer in range(p):
            cost_h(qc, params[2*layer])
            mixer_h(qc, params[2*layer+1])
        outcome, probs = qc.measure()
        cut = sum(probs[i]*sum(1 for u,v in graph if (i>>u)&1 != (i>>v)&1) for i in range(2**num_qubits))
        return -cut
    init = np.random.uniform(0, np.pi, 2*p)
    result = minimize(objective, init, method="COBYLA", options={"maxiter": max_iter})
    qc = QuantumCircuit(num_qubits, "statevector")
    for i in range(num_qubits): qc.h(i)
    for layer in range(p):
        cost_h(qc, result.x[2*layer])
        mixer_h(qc, result.x[2*layer+1])
    outcome, probs = qc.measure()
    cut = sum(probs[i]*sum(1 for u,v in graph if (i>>u)&1 != (i>>v)&1) for i in range(2**num_qubits))
    qc.gate_history.append(f"QAOA (Cut Value: {cut:.2f})")
    return {"circuit": qc, "cut_value": cut, "params": result.x}

def hhl(num_qubits, matrix=None, vector=None):
    if matrix is None: matrix = np.array([[1,-0.5],[-0.5,1]])
    if vector is None: vector = np.array([1,0])
    if num_qubits < 3: raise ValueError("HHL needs at least 3 qubits")
    if not np.allclose(matrix, matrix.conj().T): raise ValueError("Matrix must be Hermitian")
    n_state = int(np.log2(len(vector)))
    n_eigen = num_qubits - n_state - 1
    qc = QuantumCircuit(num_qubits, "statevector")
    ancilla = 0
    eigen_qubits = list(range(1, n_eigen+1))
    state_qubits = list(range(n_eigen+1, num_qubits))
    norm = np.linalg.norm(vector)
    if norm == 0: raise ValueError("Input vector cannot be zero")
    vector = vector / norm
    for i in range(n_state):
        if abs(vector[i]) > 1e-10:
            theta = 2*np.arccos(min(1, abs(vector[i])))
            qc.ry(state_qubits[i], theta)
    for i, eq in enumerate(eigen_qubits):
        for _ in range(2**i):
            theta = np.arccos(matrix[0,0] / np.linalg.norm(matrix))
            cos, sin = np.cos(theta/2), np.sin(theta/2)
            ry_gate = np.array([[cos,-sin],[sin,cos]])
            for st in state_qubits:
                dim = 4
                fg = np.eye(dim, dtype=complex)
                fg[2:,2:] = ry_gate
                qc._apply_gate(fg, [eq, st])
                qc.gate_history.append(f"Controlled-RY control {eq} target {st}")
    for i, q in enumerate(eigen_qubits):
        angle = np.pi / (2**(i+1))
        qc.cnot(q, ancilla)
        qc.ry(ancilla, angle)
        qc.cnot(q, ancilla)
    qc.normalize()
    outcome, probs = qc.measure()
    probs = probs / probs.sum()
    solution = np.array([probs[int("1"+"0"*n_eigen+format(i, f"0{n_state}b"), 2)]
                         if int("1"+"0"*n_eigen+format(i, f"0{n_state}b"), 2) < len(probs) else 0
                         for i in range(2**n_state)])
    sn = np.linalg.norm(solution)
    if sn > 1e-10: solution = solution / sn
    qc.gate_history.append(f"HHL (Solution norm: {np.linalg.norm(solution):.4f})")
    return {"circuit": qc, "solution": solution, "probs": probs}
