from .week1 import blood_cell_compatibility_lookup, BloodType
import random


class Dealer:
    def __init__(self):
        T = [[blood_cell_compatibility_lookup(r, d)
              for d in BloodType ]
             for r in BloodType]
        self.M_B = [[random.randint(0, 1) for _ in range(8)] for _ in range(8)]

        self.r = random.randint(0, 8)
        self.s = random.randint(0, 8)

        self.M_A = [[self.M_B[i][j] ^ T[i - self.r][j - self.s] for j in range(8)] for i in range(8)]

    def rand_a(self):
        return self.r, self.M_A

    def rand_b(self):
        return self.s, self.M_B


class Alice:
    def __init__(self, x, r, M_A):
        self.x = x
        self.r = r
        self.M_A = M_A

    def send(self):
        # Should return u = x+r mod 2^n
        self.u = (self.x + self.r) % 8
        return self.u

    def receive(self, v, z_b):
        self.v = v
        self.z_b = z_b

    def output(self):
        z = self.M_A[self.u][self.v] ^ self.z_b
        return z


class Bob:
    def __init__(self, y, s, M_B):
        self.y = y
        self.s = s
        self.M_B = M_B

    def send(self):
        # Should return v = y+s mod 2^n
        v = (self.y + self.s) % 8
        z_b = self.M_B[self.u][v]

        return v, z_b

    def receive(self, u):
        self.u = u


def ottt(recipient, donor):
    dealer = Dealer()
    alice = Alice(recipient, *dealer.rand_a())
    bob = Bob(donor, *dealer.rand_b())
    bob.receive(alice.send())
    alice.receive(*bob.send())
    z = alice.output()
    return z


def main():
    green = 0
    red = 0
    for i, recipient in enumerate(BloodType):
        for j, donor in enumerate(BloodType):
            z = ottt(i, j)
            lookup = blood_cell_compatibility_lookup(recipient, donor)
            if lookup == z:
                green += 1
            else:
                print(f"'{BloodType(donor).name} -> {BloodType(recipient).name}' should be {lookup}.")
                red += 1
    print("Green:", green)
    print("Red  :", red)