from .base import QuantumProtocol
from qiskit_aer import Aer
from qiskit import QuantumCircuit, QuantumRegister, ClassicalRegister
import numpy as np
import random
from typing import Dict, List, Tuple
from qiskit.visualization import circuit_drawer



class E91Protocol(QuantumProtocol):
    """
    Implementação do protocolo E91 (Ekert91) baseado em emaranhamento.
    
    O protocolo E91 usa pares de qubits emaranhados e medições em diferentes ângulos
    para estabelecer uma chave segura e verificar a presença de espionagem através
    da violação das desigualdades de Bell.
    """
    
    def __init__(self, key_length: int = 128):
        super().__init__(key_length)
        self.simulator = Aer.get_backend('qasm_simulator')
        self.current_circuit = None
        
        # Ângulos de medição para Alice e Bob
        self.alice_angles = [0, np.pi/4, np.pi/2]
        self.bob_angles = [np.pi/4, np.pi/2, 3*np.pi/4]
    
    def _prepare_entangled_pair(self) -> QuantumCircuit:
        """
        Prepara um par de qubits emaranhados no estado de Bell Φ+.
        
        Returns:
            QuantumCircuit: Circuito com o par emaranhado
        """
        qr = QuantumRegister(2, 'q')
        cr = ClassicalRegister(2, 'c')
        qc = QuantumCircuit(qr, cr)
        
        # Criar estado de Bell Φ+ = (|00⟩ + |11⟩)/√2
        qc.h(qr[0])
        qc.cx(qr[0], qr[1])
        
        self.current_circuit = qc
        return qc
    
    def _measure_qubit(self, circuit: QuantumCircuit, angle: float, qubit: int) -> QuantumCircuit:
        """
        Adiciona medição em um ângulo específico para um qubit.
        
        Args:
            circuit (QuantumCircuit): Circuito a ser medido
            angle (float): Ângulo de medição
            qubit (int): Índice do qubit a ser medido
            
        Returns:
            QuantumCircuit: Circuito com medição adicionada
        """
        # Rotação para o ângulo de medição
        circuit.ry(2*angle, qubit)
        circuit.measure([qubit], [qubit])
        return circuit

    def _check_bell_inequality(self, correlations: List[float]) -> bool:
        """
        Verifica a violação da desigualdade de CHSH.
        
        Args:
            correlations (List[float]): Lista de correlações medidas
            
        Returns:
            bool: True se a desigualdade for violada (indicando ausência de espião)
        """
        # Cálculo do parâmetro S da desigualdade CHSH
        S = abs(sum(correlations))
        # O valor teórico máximo é 2√2 ≈ 2.82
        # Um valor > 2 indica violação da desigualdade
        return S > 2
    
    def _calculate_correlation(self, alice_results: List[int], 
                             bob_results: List[int]) -> float:
        """
        Calcula a correlação entre as medições de Alice e Bob.
        
        Args:
            alice_results (List[int]): Resultados das medições de Alice
            bob_results (List[int]): Resultados das medições de Bob
            
        Returns:
            float: Valor da correlação (-1 a 1)
        """
        if not alice_results or not bob_results:
            return 0
        
        return np.mean([a == b for a, b in zip(alice_results, bob_results)])

    def generate_key(self) -> Dict[str, List[int]]:
        """
        Executa o protocolo E91 completo para gerar uma chave compartilhada.
        
        Returns:
            Dict[str, List[int]]: Dicionário contendo as chaves finais de Alice e Bob
        """
        self.logger.info(f"Iniciando protocolo E91 para gerar chave de {self.key_length} bits")
        
        alice_key = []
        bob_key = []
        correlations = []
        
        # Precisamos de mais pares para compensar descartes
        num_pairs = self.key_length * 4
        
        for _ in range(num_pairs):
            # Preparar par emaranhado
            circuit = self._prepare_entangled_pair()
            
            # Alice e Bob escolhem ângulos aleatoriamente
            alice_angle = random.choice(self.alice_angles)
            bob_angle = random.choice(self.bob_angles)
            
            # Adicionar medições
            circuit = self._measure_qubit(circuit, alice_angle, 0)  # Alice
            circuit = self._measure_qubit(circuit, bob_angle, 1)    # Bob
            
            # Executar circuito
            job = self.simulator.run(circuit, shots=1)
            result = job.result()
            counts = list(result.get_counts().keys())[0]
            alice_result = int(counts[0])
            bob_result = int(counts[1])
            
            # Verificar se os ângulos são compatíveis para a chave
            if alice_angle == bob_angle:
                alice_key.append(alice_result)
                bob_key.append(bob_result)
            else:
                # Usar para teste de Bell
                correlation = self._calculate_correlation([alice_result], [bob_result])
                correlations.append(correlation)
        
        # Verificar violação da desigualdade de Bell
        if not self._check_bell_inequality(correlations):
            self.logger.error("Possível espionagem detectada - violação de Bell não observada!")
            return {"alice": [], "bob": []}
        
        # Cortar para o tamanho desejado
        alice_key = alice_key[:self.key_length]
        bob_key = bob_key[:self.key_length]
        
        self.logger.info(f"Protocolo E91 concluído com sucesso. "
                        f"Chave final de {len(alice_key)} bits gerada.")
        
        return {
            "alice": alice_key,
            "bob": bob_key
        }

    def get_circuit_visualization(self) -> str:
        """
        Retorna a visualização do circuito atual em ASCII.
        
        Returns:
            str: Representação ASCII do circuito
        """
        if self.current_circuit:
            return circuit_drawer(self.current_circuit, output='text')
        return "Nenhum circuito disponível"