import cirq
import numpy as np


def bitstring(bits):
    """Преобразует список битов в строку, состоящую из символов '0' и '1'.

    Аргументы:
    bits -- список или массив битов (0 или 1)

    Возвращает:
    строку, например, '0101101'
    """
    return "".join(str(int(b)) for b in bits)


class BB84(object):

    def __init__(self, repetitions=1, random_seed=None):
        """
        Инициализация протокола BB84.

        repetitions -- число повторений измерений (обычно 1, для статистики можно больше)
        random_seed -- зерно генератора случайных чисел для воспроизводимости (если None, сгенерируется случайно)
        """
        super().__init__()
        if random_seed is None:
            random_seed = np.random.randint(0, 2**31 - 1)  # Генерируем случайное зерно
        np.random.seed(random_seed)  # Настраиваем генератор случайных чисел
        self.repetitions = repetitions  # Количество повторов симуляции схемы
        self.circuits = {}  # Словарь для хранения схем с разными названиями

    def setup_demo(self, num_qubits):
        """
        Генерация случайных параметров для одной сессии BB84.

        Выбираются:
        - Базисы Алисы (0 - вычислительный, 1 - базис Адамара)
        - Состояния Алисы для каждого кубита (0 или 1)
        - Базисы Боба (случайные 0 или 1)
        - Базисы Евы (случайные 0 или 1)
        Также вычисляется ожидаемый секретный ключ для совпадающих базисов Алисы и Боба
        """
        self.alice_basis = [
            np.random.randint(0, 2) for _ in range(num_qubits)
        ]  # 0 или 1
        self.alice_state = [
            np.random.randint(0, 2) for _ in range(num_qubits)
        ]  # 0 или 1
        self.bob_basis = [np.random.randint(0, 2) for _ in range(num_qubits)]
        self.eve_basis = [np.random.randint(0, 2) for _ in range(num_qubits)]

        # Формируем ожидаемый ключ - отбираем биты Алисы там, где базисы совпали с Бобом
        self.expected_key = bitstring(
            [
                self.alice_state[i]
                for i in range(num_qubits)
                if self.alice_basis[i] == self.bob_basis[i]
            ]
        )

    def make_bb84_circuit(
        self, num_qubits, sender_basis, receiver_basis, sender_state, name="1"
    ):
        """
        Построение квантовой схемы для протокола BB84.

        Параметры:
        num_qubits -- число кубитов
        sender_basis -- базисы отправителя (0 - вычислительный, 1 - Адамара)
        receiver_basis -- базисы получателя (0 или 1)
        sender_state -- биты состояния отправителя (0 или 1)
        name -- название схемы (ключ в словаре circuits)

        Алгоритм:
        1) Отправитель подготавливает кубиты:
           - Применяет X к кубитам с состоянием 1 (переключение |0> -> |1>)
           - Применяет H к кубитам, если базис Адамара (1)
        2) Получатель измеряет кубиты:
           - Применяет H перед измерением, если базис Адамара (1)
        3) Все кубиты измеряются
        """
        self.qubits = cirq.LineQubit.range(num_qubits)  # Создание линии кубитов
        self.circuits[name] = cirq.Circuit()  # Инициализация пустой схемы

        sender_flips = []  # Операции X для изменения состояния с |0> на |1>
        sender_rotations = []  # Операции H для переключения базиса

        # Цикл по кубитам на стороне отправителя
        for index in range(num_qubits):
            if sender_state[index] == 1:
                sender_flips.append(
                    cirq.X(self.qubits[index])
                )  # Применяем X, если бит 1
            if sender_basis[index] == 1:
                sender_rotations.append(
                    cirq.H(self.qubits[index])
                )  # Применяем H, если выбран базис Адамара

        # Добавляем операции отправителя в схему как отдельные "слои" времени (Moments),
        # что соответствует выполнению этих операций одновременно независимо для разных кубитов
        self.circuits[name].append(cirq.Moment(sender_flips))
        self.circuits[name].append(cirq.Moment(sender_rotations))

        # Операции получателя - подготовка базиса измерения
        recv_basis_choice = []
        for index, rbasis in enumerate(receiver_basis):
            if rbasis == 1:
                recv_basis_choice.append(
                    cirq.H(self.qubits[index])
                )  # Применяем H для базиса Адамара перед измерением

        self.circuits[name].append(
            cirq.Moment(recv_basis_choice)
        )  # Добавляем перед измерением
        self.circuits[name].append(
            cirq.Moment(cirq.measure_each(*self.qubits))
        )  # Измеряем все кубиты

    def simulate_base_protocol(self, num_qubits=8, print_legend=None):
        """
        Симуляция базового протокола BB84 без присутствия перехватчика (Евы).

        - Генерация параметров
        - Создание схемы для Алисы и Боба
        - Симуляция измерений
        - Выделение ключа на основе совпадения базисов
        - Печать схемы и результатов
        """
        self.setup_demo(num_qubits)

        # Строим схему взаимодействия Алисы и Боба
        self.make_bb84_circuit(
            num_qubits, self.alice_basis, self.bob_basis, self.alice_state, name="base"
        )

        # Запуск симуляции схемы c заданным числом повторений
        result = cirq.Simulator().run(
            program=self.circuits["base"], repetitions=self.repetitions
        )

        # Получение бит результата измерений: для каждого кубита берётся первый результат первого повторения (.item())
        result_bitstring = bitstring(
            [int(result.measurements[str(q)][0].item()) for q in self.qubits]
        )

        # Формируем ключ Боба только для тех позиций, где базисы совпали с Алисой
        obtained_key = "".join(
            [
                result_bitstring[i]
                for i in range(num_qubits)
                if self.alice_basis[i] == self.bob_basis[i]
            ]
        )

        # Выводим схему и результаты
        print("########## Печать схемы ##########")
        print(self.circuits["base"])
        print("########## Печать результатов ##########")
        self.print_results(
            self.alice_basis,
            self.bob_basis,
            self.alice_state,
            self.expected_key,
            obtained_key,
        )
        print(
            "Общий секретный ключ успешно распределен и может использоваться для шифрования."
        )

    def simulate_eavesdropped_protocol(self, num_qubits=8, print_legend=None):
        """
        Симуляция протокола BB84 с перехватом.

        - Генерация параметров
        - Симуляция перехвата Евы между Алисой и Бобом
        - Симуляция передачи от Евы к Бобу
        - Формирование ключей по разным базисам
        - Печать схемы и результатов
        """
        self.setup_demo(num_qubits)

        # Создаем схему передачи Алиса -> Ева (с их базисами)
        self.make_bb84_circuit(
            num_qubits,
            self.alice_basis,
            self.eve_basis,
            self.alice_state,
            name="alice_eve",
        )

        # Симуляция измерений Евы
        result = cirq.Simulator().run(
            program=self.circuits["alice_eve"], repetitions=self.repetitions
        )
        # Ева сохраняет измеренные биты
        eve_state = [int(result.measurements[str(q)][0].item()) for q in self.qubits]

        # Ева формирует новую схему для передачи Бобу, используя свои измерения и выбранный базис
        self.make_bb84_circuit(
            num_qubits, self.eve_basis, self.bob_basis, eve_state, name="eve_bob"
        )

        # Симуляция измерений Боба
        result = cirq.Simulator().run(
            program=self.circuits["eve_bob"], repetitions=self.repetitions
        )
        result_bitstring = bitstring(
            [int(result.measurements[str(q)][0].item()) for q in self.qubits]
        )

        # Формируем ключи для анализа совпадения (Алиса-Ева и Алиса-Боб)
        intercepted_key = "".join(
            [
                result_bitstring[i]
                for i in range(num_qubits)
                if self.alice_basis[i] == self.eve_basis[i]
            ]
        )
        obtained_key = "".join(
            [
                result_bitstring[i]
                for i in range(num_qubits)
                if self.alice_basis[i] == self.bob_basis[i]
            ]
        )

        # Объединяем цепочки схем Алиса->Ева и Ева->Боб для вывода
        circuit = self.circuits["alice_eve"] + self.circuits["eve_bob"]
        print("########## Печать схемы ##########")
        print(circuit)
        print("########## Печать результатов ##########")
        self.print_eavesdropped_results(
            self.alice_basis,
            self.bob_basis,
            self.eve_basis,
            self.alice_state,
            eve_state,
            self.expected_key,
            obtained_key,
        )

        print("Ева попыталась подслушать и повторно отправить кубиты.")
        print("Это приводит к ошибкам в ключе, которые обнаруживает Боб.")

    def print_results(
        self, alice_basis, bob_basis, alice_state, expected_key, obtained_key
    ):
        """
        Вывод результатов базового протокола без перехвата.

        Печатаются:
        - исходное сообщение Алисы,
        - базисы Алисы и Боба,
        - совпадения базисов (в виде X),
        - ключи Алисы и Боба.
        """
        num_qubits = len(alice_basis)
        basis_match = "".join(
            ["X" if alice_basis[i] == bob_basis[i] else "_" for i in range(num_qubits)]
        )
        alice_basis_str = "".join(
            ["C" if alice_basis[i] == 0 else "H" for i in range(num_qubits)]
        )
        bob_basis_str = "".join(
            ["C" if bob_basis[i] == 0 else "H" for i in range(num_qubits)]
        )

        print(f"Только Алиса знает сообщение:\t{bitstring(alice_state)}")
        print("Базисы Алисы и Боба (C - вычислительный, H - Адамара):")
        print(f"Алиса: {alice_basis_str}")
        print(f"Боб:   {bob_basis_str}")
        print(f"Совпадение базисов:   {basis_match}")
        print(f"Ключ Алисы: {expected_key}")
        print(f"Ключ Боба: {obtained_key}")

    def print_eavesdropped_results(
        self,
        alice_basis,
        bob_basis,
        eve_basis,
        alice_state,
        eve_state,
        expected_key,
        obtained_key,
    ):
        """
        Вывод результатов протокола с перехватом.

        Печатаются:
        - исходное сообщение Алисы,
        - предположение Евы (ее измеренные биты),
        - базисы Алисы, Евы и Боба,
        - совпадения базисов,
        - ключи Алисы, Евы (ставится равным ключу Алисы, тк она пытается угадать),
          и ключ Боба.
        """
        num_qubits = len(alice_basis)
        eve_basis_match = "".join(
            ["X" if alice_basis[i] == eve_basis[i] else "_" for i in range(num_qubits)]
        )
        bob_basis_match = "".join(
            ["X" if alice_basis[i] == bob_basis[i] else "_" for i in range(num_qubits)]
        )
        alice_basis_str = "".join(
            ["C" if alice_basis[i] == 0 else "H" for i in range(num_qubits)]
        )
        eve_basis_str = "".join(
            ["C" if eve_basis[i] == 0 else "H" for i in range(num_qubits)]
        )
        bob_basis_str = "".join(
            ["C" if bob_basis[i] == 0 else "H" for i in range(num_qubits)]
        )

        print(f"Только Алиса знает сообщение:\t{bitstring(alice_state)}")
        print(f"Предположение Евы о cообщении:\t{bitstring(eve_state)}")
        print("Базисы (C - вычислительный, H - Адамара):")
        print(f"Алиса: {alice_basis_str}")
        print(f"Ева:   {eve_basis_str}")
        print(f"Боб:   {bob_basis_str}")
        print(f"Совпадение базисов Алисы и Евы: {eve_basis_match}")
        print(f"Совпадение базисов Алисы и Боба: {bob_basis_match}")
        print(f"Ключ Алисы: {expected_key}")
        print(f"Ключ Евы: {expected_key}")  # По идее, Ева пытается угадать ключ Алисы
        print(f"Ключ Боба: {obtained_key}")


if __name__ == "__main__":
    bb84 = BB84()

    print("=== БАЗОВЫЙ ПРОТОКОЛ BB84 ===")
    bb84.simulate_base_protocol(num_qubits=8, print_legend=True)

    print("\n\n=== ПРОТОКОЛ BB84 С ПЕРЕХВАТОМ ===")
    bb84.simulate_eavesdropped_protocol(num_qubits=8, print_legend=False)