from sage.all import GF, PolynomialRing

from tea3.constants import TEA3_SBOX, T_F1, T_F2
from tea3.pretty_print import pretty_print, pretty_print_vec


class Tea3Model:
    def __init__(self):
        self.F = GF(2)

        names = (
            [f"x{i}{j}" for i in range(5) for j in range(8)] +
            [f"r{i}{j}" for i in range(5) for j in range(8)] +
            [f"R{i}{j}" for i in range(8) for j in range(8)]
        )

        self.S = PolynomialRing(self.F, names)
        self.v = self.S.gens()

        self.x_bits = [list(self.v[i*8:(i+1)*8]) for i in range(5)]
        self.r_bits = [list(self.v[40 + i*8 : 40 + (i+1)*8]) for i in range(5)]
        self.R_bits = [list(self.v[80 + i*8 : 80 + (i+1)*8]) for i in range(8)]

    def step(self):
        R = self.R_bits.copy()
        x = self.x_bits.copy()
        r = self.r_bits.copy()

        R0, R1, R2, R3, R4, R5, R6, R7 = R
        x0, x1, x2, x3, x4 = x
        r0, r1, r2, r3, r4 = r

        # update r register
        self.r_bits[0] = r1
        self.r_bits[1] = r2
        self.r_bits[2] = r3
        self.r_bits[3] = r4
        self.r_bits[4] = xor_vec(r0, S(r2))

        # update x register
        self.x_bits[0] = x1
        self.x_bits[1] = x2
        self.x_bits[2] = x3
        self.x_bits[3] = x4
        self.x_bits[4] = xor_vec(x0, r0)

        # update R registers
        self.R_bits[7] = R6
        self.R_bits[6] = R5
        self.R_bits[5] = xor_vec(R4, G31(R6, R5))
        self.R_bits[4] = R3
        self.R_bits[3] = R2
        self.R_bits[2] = R1
        self.R_bits[1] = R0
        self.R_bits[0] = xor_vec(x0, xor_vec(R7, xor_vec(BP(R4), G32(R2, R1))))

        return R7

def S(r):
    # placeholder
    return r

def xor_vec(a, b):
    return [ai + bi for ai, bi in zip(a, b)]

def BP(r):
    return [r[2], r[7], r[3], r[5], r[6], r[1], r[0], r[4]]

def G31(X, Y):
    x3,x4,x5,x6,x7,x0,x1,x2 = X
    y3,y4,y5,y6,y7,y0,y1,y2 = Y

    return [
        x1 + y0 + x0*y0 + y1 + x0*y1 + x0*x1*y1 + y0*y1 + x0*y0*y1 + x1*y0*y1,
        1 + y1 + y2 + x1*y2 + x2*y2 + y1*y2 + x1*y1*y2,
        x2 + x3 + x2*y2 + x3*y2 + y3 + x3*y3 + x3*y2*y3,
        x4*y3 + y4 + x3*y3*y4 + x4*y3*y4,
        1 + y4 + x4*y4 + x5*y4 + y5 + y4*y5 + x5*y4*y5,
        1 + x5 + x6 + y5 + x5*y5 + x6*y5 + y6 + x6*y6 + x6*y5*y6,
        1 + x6 + x6*y6 + x7*y7 + y6*y7 + x7*y6*y7,
        1 + x7 + y0 + x0*y7 + x7*y7 + x0*x7*y7 + x0*y0*y7 + x7*y0*y7,
    ]

def G32(X, Y):
    x3,x4,x5,x6,x7,x0,x1,x2 = X
    y3,y4,y5,y6,y7,y0,y1,y2 = Y

    return [
        1 + x1 + x0*x1 + y0 + x0*y0 + x1*y0 + y1 + x0*x1*y1 + x0*y0*y1 + x1*y0*y1,
        1 + y1 + x2*y2 + x1*y1*y2,
        x3 + x3*y2 + y3 + x2*y3 + x3*y3 + x2*y2*y3 + x3*y2*y3,
        x4*y3 + y4 + y3*y4 + x3*y3*y4 + x4*y3*y4,
        x4*y4 + y5 + y4*y5 + x4*y4*y5 + x5*y4*y5,
        x5 + x6 + x5*y5 + x6*y5 + y6 + x6*y6 + x6*y5*y6,
        y6 + x6*y6 + x7*y6 + y7 + y6*y7 + x7*y6*y7,
        x0*x7 + y0 + x7*y0 + y7 + x0*y7 + x7*y7 + x0*x7*y7 + x0*y0*y7 + x7*y0*y7,
    ]



t = Tea3Model()
for i in range(3):
    print("step "+str(i))
    t.step()
    print(pretty_print_vec(t.R_bits[0]))