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

class BB84Protocol(QuantumProtocol):
    """
    Implementação do protocolo BB84 para distribuição quântica de chaves.
    Herda da classe base QuantumProtocol e implementa correção de erros quânticos.
    """

    def __init__(self, key_length: int = 128, error_correction: bool = True):
        """
        Inicializa o protocolo BB84.

        Args:
            key_length (int): Tamanho desejado da chave final em bits
            error_correction (bool): Se True, utiliza correção de erros quânticos
        """
        super().__init__(key_length)  # Chama o construtor da classe base
        self.error_correction = error_correction
        self.simulator = Aer.get_backend('qasm_simulator')
        self.current_circuit = None

    def _prepare_qubit(self, bit: int, basis: str) -> QuantumCircuit:
        """
        Prepara um único qubit no estado e base especificados.

        Args:
            bit (int): Bit a ser codificado (0 ou 1)
            basis (str): Base de medição ('Z' para rectilinear ou 'X' para diagonal)

        Returns:
            QuantumCircuit: Circuito quântico preparado
        """
        if self.error_correction:
            # Circuito com correção de erros (código de repetição de 3 qubits)
            qr = QuantumRegister(3, 'q')
            cr = ClassicalRegister(1, 'c')
            qc = QuantumCircuit(qr, cr)
            
            # Prepara o estado base
            if bit == 1:
                qc.x(qr[0])
            
            # Aplica o código de repetição
            qc.cx(qr[0], qr[1])  # CNOT para copiar o estado
            qc.cx(qr[0], qr[2])
            
            # Aplica a transformação de base se necessário
            if basis == 'X':
                qc.h(qr)  # Aplica H em todos os qubits
        else:
            # Circuito simples sem correção
            qr = QuantumRegister(1, 'q')
            cr = ClassicalRegister(1, 'c')
            qc = QuantumCircuit(qr, cr)
            
            # Prepara o estado
            if bit == 1:
                qc.x(qr[0])
            if basis == 'X':
                qc.h(qr[0])

        self.current_circuit = qc
        return qc

    def _measure_qubit(self, circuit: QuantumCircuit, basis: str) -> QuantumCircuit:
        """
        Adiciona medição em uma base específica ao circuito.

        Args:
            circuit (QuantumCircuit): Circuito a ser medido
            basis (str): Base de medição ('Z' ou 'X')

        Returns:
            QuantumCircuit: Circuito com medição adicionada
        """
        if self.error_correction:
            # Medição com correção de erros
            if basis == 'X':
                circuit.h(range(3))  # Aplica H em todos os qubits para base X
            
            # Medição em todos os qubits
            circuit.measure_all()
            
        else:
            # Medição simples
            if basis == 'X':
                circuit.h(0)
            circuit.measure([0], [0])
            
        return circuit

    def _error_correction_decode(self, measurements: List[int]) -> int:
        """
        Decodifica medições usando voto majoritário para correção de erros.

        Args:
            measurements (List[int]): Lista de medições dos qubits

        Returns:
            int: Bit decodificado
        """
        return 1 if sum(measurements) > len(measurements)/2 else 0

    def generate_random_bits(self, n: int) -> List[int]:
        """
        Gera uma lista de bits aleatórios.

        Args:
            n (int): Número de bits a serem gerados

        Returns:
            List[int]: Lista de bits aleatórios
        """
        return [random.randint(0, 1) for _ in range(n)]

    def generate_random_bases(self, n: int) -> List[str]:
        """
        Gera uma lista de bases aleatórias.

        Args:
            n (int): Número de bases a serem geradas

        Returns:
            List[str]: Lista de bases aleatórias ('Z' ou 'X')
        """
        return [random.choice(['Z', 'X']) for _ in range(n)]

    def simulate_transmission(self, bits: List[int], bases: List[str]) -> Tuple[List[int], List[str]]:
        """
        Simula a transmissão dos qubits de Alice para Bob.

        Args:
            bits (List[int]): Bits a serem transmitidos
            bases (List[str]): Bases usadas para codificação

        Returns:
            Tuple[List[int], List[str]]: Medições de Bob e bases escolhidas
        """
        bob_bases = self.generate_random_bases(len(bits))
        bob_measurements = []
        
        for i in range(len(bits)):
            # Alice prepara o qubit
            circuit = self._prepare_qubit(bits[i], bases[i])
            
            # Bob mede o qubit
            circuit = self._measure_qubit(circuit, bob_bases[i])
            
            # Executa a simulação
            job = self.simulator.run(circuit, shots=1)
            result = job.result()
            counts = list(result.get_counts().keys())[0]
            
            if self.error_correction:
                # Decodifica as medições com correção de erros
                measurements = [int(bit) for bit in counts]
                measured_bit = self._error_correction_decode(measurements)
            else:
                measured_bit = int(counts)
                
            bob_measurements.append(measured_bit)
            
        return bob_measurements, bob_bases

    def sift_key(self, alice_bits: List[int], alice_bases: List[str], 
                 bob_measurements: List[int], bob_bases: List[str]) -> Tuple[List[int], List[int]]:
        """
        Realiza o processo de peneiramento da chave.

        Args:
            alice_bits (List[int]): Bits originais de Alice
            alice_bases (List[str]): Bases usadas por Alice
            bob_measurements (List[int]): Medições de Bob
            bob_bases (List[str]): Bases usadas por Bob

        Returns:
            Tuple[List[int], List[int]]: Chaves peneiradas de Alice e Bob
        """
        alice_key = []
        bob_key = []
        
        for i in range(len(alice_bits)):
            if alice_bases[i] == bob_bases[i]:
                alice_key.append(alice_bits[i])
                bob_key.append(bob_measurements[i])
                
        return alice_key, bob_key

    def check_eavesdropping(self, alice_key: List[int], bob_key: List[int], 
                           sample_size: float = 0.25) -> bool:
        """
        Verifica a presença de um espião comparando uma amostra das chaves.

        Args:
            alice_key (List[int]): Chave de Alice
            bob_key (List[int]): Chave de Bob
            sample_size (float): Proporção da chave a ser verificada

        Returns:
            bool: True se não houver evidência de espionagem, False caso contrário
        """
        if len(alice_key) != len(bob_key):
            return False
            
        sample_length = int(len(alice_key) * sample_size)
        indices = random.sample(range(len(alice_key)), sample_length)
        
        differences = 0
        for i in indices:
            if alice_key[i] != bob_key[i]:
                differences += 1
                
        error_rate = differences / sample_length
        self.logger.info(f"Taxa de erro na amostra: {error_rate:.2%}")
        
        # Define um limite aceitável de erro (5%)
        return error_rate < 0.05

    def get_circuit_visualization(self) -> str:
        """
        Implementação do método abstrato da classe base.
        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"

    def generate_key(self) -> Dict[str, List[int]]:
        """
        Implementação do método abstrato da classe base.
        Executa o protocolo BB84 completo para gerar uma chave compartilhada.

        Returns:
            Dict[str, List[int]]: Dicionário contendo as chaves finais de Alice e Bob
        """
        # Fase 1: Preparação
        raw_bits = self.generate_random_bits(self.key_length * 4)
        alice_bases = self.generate_random_bases(len(raw_bits))
        
        self.logger.info(f"Iniciando protocolo BB84 para gerar chave de {self.key_length} bits")
        self.logger.info(f"Correção de erros: {'Ativada' if self.error_correction else 'Desativada'}")
        
        # Fase 2: Transmissão
        bob_measurements, bob_bases = self.simulate_transmission(raw_bits, alice_bases)
        
        # Fase 3: Peneiramento
        alice_key, bob_key = self.sift_key(raw_bits, alice_bases, bob_measurements, bob_bases)
        
        self.logger.info(f"Chave peneirada gerada com {len(alice_key)} bits")
        
        # Fase 4: Detecção de espionagem
        if not self.check_eavesdropping(alice_key, bob_key):
            self.logger.error("Possível tentativa de espionagem detectada!")
            return {"alice": [], "bob": []}
        
        # Fase 5: Finalização
        final_key_length = min(self.key_length, len(alice_key))
        alice_final_key = alice_key[:final_key_length]
        bob_final_key = bob_key[:final_key_length]
        
        self.logger.info(f"Protocolo BB84 concluído com sucesso. "
                        f"Chave final de {final_key_length} bits gerada.")
        
        return {
            "alice": alice_final_key,
            "bob": bob_final_key
        }

    def get_key_statistics(self, alice_key: List[int], bob_key: List[int]) -> Dict:
        """
        Calcula estatísticas sobre as chaves geradas.

        Args:
            alice_key (List[int]): Chave de Alice
            bob_key (List[int]): Chave de Bob

        Returns:
            Dict: Dicionário com estatísticas das chaves
        """
        if not alice_key or not bob_key:
            return {
                "key_length": 0,
                "error_rate": 1.0,
                "matching_rate": 0.0,
                "entropy": 0.0
            }

        differences = sum(a != b for a, b in zip(alice_key, bob_key))
        error_rate = differences / len(alice_key)
        
        # Calcula a entropia da chave
        ones = sum(alice_key) / len(alice_key)
        zeros = 1 - ones
        if zeros == 0 or ones == 0:
            entropy = 0
        else:
            entropy = -(ones * np.log2(ones) + zeros * np.log2(zeros))

        return {
            "key_length": len(alice_key),
            "error_rate": error_rate,
            "matching_rate": 1 - error_rate,
            "entropy": entropy
        }

def main():
    """Função principal para demonstração do protocolo."""
    # Criar uma instância do protocolo
    bb84 = BB84Protocol(key_length=64, error_correction=True)
    
    # Gerar uma chave compartilhada
    result = bb84.generate_key()
    
    # Verificar se as chaves são idênticas
    if result["alice"] and result["bob"]:
        print("\nChave de Alice:", "".join(map(str, result["alice"])))
        print("Chave de Bob:  ", "".join(map(str, result["bob"])))
        
        # Obter e mostrar estatísticas
        stats = bb84.get_key_statistics(result["alice"], result["bob"])
        print("\nEstatísticas da chave:")
        print(f"Tamanho: {stats['key_length']} bits")
        print(f"Taxa de erro: {stats['error_rate']:.2%}")
        print(f"Taxa de correspondência: {stats['matching_rate']:.2%}")
        print(f"Entropia: {stats['entropy']:.2f} bits")
        
        # Mostrar visualização do último circuito
        print("\nÚltimo circuito gerado:")
        print(bb84.get_circuit_visualization())
    else:
        print("Falha na geração da chave - possível tentativa de espionagem detectada.")