Quantum_Cryptography/model/channel.py

import random
from dataclasses import dataclass, field
import typing
import abc
import numpy as np


class Channel(abc.ABC):
    @abc.abstractmethod
    def send_classical(self, data: str, is_bob: bool = True):
        pass

    @abc.abstractmethod
    def send_bit(self, bit: bool, basis: bool):
        pass

    @abc.abstractmethod
    def recv_classical(self, n: int, is_bob: bool = True) -> str:
        pass

    @abc.abstractmethod
    def get_photon_results(self, bases: typing.Optional[typing.List[bool]]) -> typing.List[typing.Tuple[bool, bool]]:
        """
        Returns tuple of clicks of 0 and 1 for each bit
        """
        pass


@dataclass
class ChannelSym(Channel):
    # [[(value, basis)]]
    bob_photons: typing.List[typing.List[typing.Tuple[bool, bool]]] = field(default_factory=list)
    bob_classical_data_pool: str = ''
    alice_classical_data_pool: str = ''
    # probability for a photon to go to the wrong detector
    p_opt: float = 0.015
    # dark count
    p_dc: float = 0.05
    # the average number of emitted photons
    mu: float = 1
    # eta - the probability for a detector to react to a photon
    detector_sensitivity: float = 0.8
    transmittance: float = 0.9
    eve: typing.Optional[typing.Any] = None

    def get_number_of_photos(self) -> int:
        return round(np.random.poisson(self.mu))

    def send_classical(self, data: str, is_bob: bool = True):
        if is_bob:
            self.bob_classical_data_pool += data
        else:
            self.alice_classical_data_pool += data

    def send_bit(self, bit: bool, basis: bool):
        n_photons = self.get_number_of_photos()
        delta = [(bit, basis) for _i in range(n_photons)]
        if self.eve is not None:
            delta = self.eve.process_photons(delta)
        self.bob_photons.append(delta)

    def recv_classical(self, n: int, is_bob: bool = True) -> str:
        pool_name = ['alice_classical_data_pool', 'bob_classical_data_pool'][not is_bob]
        if n == -1:
            n = len(getattr(self, pool_name))
        while len(getattr(self, pool_name)) < n:
            pass
        data = getattr(self, pool_name)[:n]
        setattr(self, pool_name, getattr(self, pool_name)[n:])
        # print(f'{["Alice", "Bob"][is_bob]} received {n} characters')
        return data

    def get_photon_results(self, bases: typing.Optional[typing.List[bool]]) -> typing.List[typing.Tuple[bool, bool]]:
        res = list()
        while len(self.bob_photons) < len(bases):
            pass
        for photons, basis in zip(self.bob_photons, bases):
            clicks = [random.uniform(0, 1) <= self.p_dc, random.uniform(0, 1) <= self.p_dc]
            if random.uniform(0, 1) > self.transmittance:
                res.append(clicks)
                continue
            for ph_bit, ph_basis in photons:
                if random.uniform(0, 1) <= self.detector_sensitivity:  # We detected everything correctly
                    if ph_basis != basis:
                        if random.uniform(0, 1) > 0.5:
                            clicks[1] = True
                        else:
                            clicks[0] = True
                        continue
                    if random.uniform(0, 1) > self.p_opt:  # Photon doesn't flip
                        clicks[ph_bit] = True
                    else:
                        clicks[not ph_bit] = True
            res.append((clicks[0], clicks[1]))
        self.bob_photons = list()
        return res