from dataclasses import dataclass, field

import random

STREAM_LENGTH = 100


@dataclass
class Photon(object):
    basis: int
    value: int
    _measured: bool = False
    _measured_value: bool = None

    def measure(self, basis):
        self._measured = True
        if self.basis == basis:
            return self.value
        else:
            if not self._measured:
                self._measured_value = random.randint(0, 1)
            return self._measured_value

    def __repr__(self):
        return '{}:{}'.format(self.basis, self.value)


@dataclass
class Alice:
    _otp: list = None
    bases: list = field(default_factory=list)
    _values: list = field(default_factory=list)

    def encode_qbit(self, value):
        basis = random.randint(0, 1)
        self.bases.append(basis)
        return Photon(basis, value)

    def generate_stream(self, stream_size=10):
        for _ in range(stream_size):
            qbit = random.randint(0, 1)
            self._values.append(qbit)
            yield self.encode_qbit(qbit)

    def set_otp(self, other_correct):
        self._otp = [v for c, v in zip(other_correct, self._values) if c]

    def get_otp_length(self):
        return len(self._otp)


@dataclass
class Bob:
    correct: list = None
    _otp: list = None
    _bases: list = field(default_factory=list)
    _values: list = field(default_factory=list)

    def ingest_stream(self, stream):
        for photon in stream:
            basis = random.randint(0, 1)
            self._bases.append(basis)
            qbit_value = photon.measure(basis)
            self._values.append(qbit_value)

    def insecure_get_bases(self, other_bases):
        self.correct = [not (a ^ b) for a, b in zip(self._bases, other_bases)]
        self._otp = [v for c, v in zip(self.correct, self._values) if c]


@dataclass
class Eve:
    _bases: list = field(default_factory=list)
    _values: list = field(default_factory=list)

    def ingest_stream(self, stream):
        for photon in stream:
            basis = random.randint(0, 1)
            self._bases.append(basis)
            qbit_value = photon.measure(basis)
            self._values.append(qbit_value)
            yield photon


def main():
    alice, bob, eve = Alice(), Bob(), Eve()

    # Step 1. Generate a stream of random photons polarizations in a random
    # basis
    stream = alice.generate_stream(STREAM_LENGTH)

    # Step 1.5 EVE ALSO INGESTS AND COPIES THE PHOTONS
    # stream = eve.ingest_stream(stream)

    # Step 2. Bob ingests the stream of photons
    bob.ingest_stream(stream)

    bob.insecure_get_bases(alice.bases)
    alice.set_otp(bob.correct)

    assert alice._otp == bob._otp
    print("OTP Length: {}".format(alice.get_otp_length()))
    print("OTP Length: {}".format(alice._otp))


if __name__ == "__main__":
    main()