from typing import List, Dict
import numpy as np
from qiskit import QuantumCircuit, transpile
from qiskit.circuit.library import Diagonal
from qiskit_aer import AerSimulator
import requests
import math
# Se eliminó la importación de Sampler (no se usaba)

class QuantumOracle:
    def __init__(self, api_key: str, license_key: str, tier: str = "enterprise",
                 base_url: str = "https://dsf-quantum-oracle-m2p0xf347-api-dsfuptech.vercel.app"):
        self.api_key = api_key
        self.license_key = license_key
        self.tier = tier
        self.base_url = base_url

    def _call_api(self, similitudes: List[float], weights: List[float],
                  criticalities: List[float], threshold: float) -> Dict:
        
        payload = {
            "similitudes": similitudes,
            "weights": weights,
            "criticalities": criticalities,
            "threshold": threshold,
            "license_key": self.license_key,
            "tier": self.tier 
        }
        
        response = requests.post(
            self.base_url,
            json=payload,
            headers={
                "X-Api-Key": self.api_key,
                "Content-Type": "application/json" 
            },
            timeout=30
        )
        response.raise_for_status()
        return response.json()

    def _build_grover_diffuser(self, selector_qubits: List[int]) -> QuantumCircuit:
        n = len(selector_qubits)
        qc = QuantumCircuit(n)
        qc.h(selector_qubits)
        qc.x(selector_qubits)
        qc.h(selector_qubits[-1])
        if n > 1:
            qc.mcx(selector_qubits[:-1], selector_qubits[-1], mode='noancilla') 
        else:
            qc.z(selector_qubits[0])
        qc.h(selector_qubits[-1])
        qc.x(selector_qubits)
        qc.h(selector_qubits)
        return qc

    def _benchmark_ancillas(self, n_controls: int) -> Dict:
        backend = AerSimulator()
        
        n_ancillas = max(0, n_controls - 2)
        qc_trad = QuantumCircuit(n_controls + 1 + n_ancillas)
        qc_trad.mcx(list(range(n_controls)), n_controls,
                    list(range(n_controls + 1, n_controls + 1 + n_ancillas)))
        qc_trad_trans = transpile(qc_trad.decompose(reps=5), backend, optimization_level=0)
        
        qc_opt = QuantumCircuit(n_controls + 1)
        qc_opt.mcx(list(range(n_controls)), n_controls, mode='noancilla')
        qc_opt_trans = transpile(qc_opt.decompose(reps=5), backend, optimization_level=0)
        
        return {
            "traditional_qubits": qc_trad_trans.num_qubits,
            "optimized_qubits": qc_opt_trans.num_qubits,
            "ancillas_saved": qc_trad_trans.num_qubits - qc_opt_trans.num_qubits,
            "traditional_depth": qc_trad_trans.depth(),
            "optimized_depth": qc_opt_trans.depth()
        }

    def search(self, similitudes: List[float], weights: List[float],
               criticalities: List[float], threshold: float = 0.8,
               include_ancilla_benchmark: bool = True) -> Dict:
        
        # Llamar API para lógica clásica
        api_response = self._call_api(similitudes, weights, criticalities, threshold)
        
        prepared = api_response['prepared_vectors']
        config = api_response['grover_config']
        
        marked_indices = prepared['marked_indices']
        k = config['num_selector_qubits']
        
        if not marked_indices:
            return {
                'found_indices': [],
                'most_measured_index': None,
                'success': False,
                'success_probability': 0.0,
                'item_scores': prepared['item_scores'],
                'tier': api_response['tier']
            }
            
        # Construir oráculo
        diag_values = np.ones(2**k)
        for i in marked_indices:
            diag_values[i] = -1
        oracle_op = Diagonal(diag_values)
        
        # Construir difusor
        diffuser_op = self._build_grover_diffuser(list(range(k)))
        
        # Circuito completo
        qc = QuantumCircuit(k, k)
        qc.h(range(k))
        
        for _ in range(config['optimal_iterations']):
            qc.append(oracle_op, range(k))
            qc.append(diffuser_op, range(k))
            
        qc.measure(range(k), range(k))
        
        # Ejecutar
        backend = AerSimulator()
        transpiled = transpile(qc, backend)
        job = backend.run(transpiled, shots=config['num_shots'])
        result = job.result()
        counts = result.get_counts()
        
        # Interpretar resultados
        decimal_counts = {}
        for bitstring, count in counts.items():
            index = int(bitstring, 2) # Conversión directa
            decimal_counts[index] = count
            
        most_measured = max(decimal_counts, key=decimal_counts.get) if decimal_counts else None
        success = most_measured in marked_indices if most_measured is not None else False
        success_prob = sum(decimal_counts.get(idx, 0) for idx in marked_indices) / config['num_shots']
        
        output = {
            'found_indices': marked_indices,
            'most_measured_index': most_measured,
            'counts': counts,
            'decimal_counts': decimal_counts,
            'success': success,
            'success_probability': success_prob,
            'item_scores': prepared['item_scores'],
            'tier': api_response['tier']
        }
        
        if include_ancilla_benchmark:
            output['ancilla_efficiency'] = self._benchmark_ancillas(k)
            
        return output

    def close(self):
        pass

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.close()

