from sage.all import BooleanPolynomialRing
from sage.crypto.sbox import SBox

from tea3.tea3model import F31, F32


def _vars():
    R = BooleanPolynomialRing(16, [f"u{i}" for i in range(16)])
    u = R.gens()
    return list(u[:8]), list(u[8:])


def _poly_masks(p, vars_):
    idx = {v: i for i, v in enumerate(vars_)}
    masks = []
    for monom in p:
        if monom:
            mask = 0
            for v in monom.variables():
                mask |= 1 << idx[v]
            masks.append(mask)
    return int(p.constant_coefficient()), masks


def _truth_table(vecfun):
    x, y = _vars()
    polys = vecfun(x, y)
    vars_ = x + y
    coords = [_poly_masks(p, vars_) for p in polys]

    tt = [0] * (1 << 16)
    for n in range(1 << 16):
        out = 0
        for j, (c, masks) in enumerate(coords):
            bit = c
            for mask in masks:
                if (n & mask) == mask:
                    bit ^= 1
            out |= bit << j
        tt[n] = out

    return tt


def _sbox(vecfun):
    return SBox(_truth_table(vecfun), big_endian=False)

def run_f31f32():
    f31 = _sbox(F31)
    f32 = _sbox(F32)
    print("\nF31")
    print("differential uniformity:", f31.differential_uniformity())
    print("max DDT coefficient:", f31.maximal_difference_probability_absolute())
    print("max difference probability:", f31.maximal_difference_probability())
    print("max LAT coefficient:", f31.maximal_linear_bias_absolute())
    print("relative bias:", f31.maximal_linear_bias_relative())

    print("\nF32")
    print("differential uniformity:", f32.differential_uniformity())
    print("max DDT coefficient:", f32.maximal_difference_probability_absolute())
    print("max difference probability:", f32.maximal_difference_probability())
    print("max LAT coefficient:", f32.maximal_linear_bias_absolute())
    print("relative bias:", f32.maximal_linear_bias_relative())