from itertools import product as iproduct
from tea3.tea3model import Tea3Model
from tea3.pretty_print import pretty_print


def count_monomials(poly):
    s = pretty_print(poly).strip()
    if not s or s == "0":
        return 0
    return sum(1 for term in s.split("+") if term.strip())


def apply_variable_change(model, coeffs):
    """
    Variable change inside each 8-bit register:
        x_{r,0} = y_{r,0} + a1*y_{r,1} + ... + a7*y_{r,7}
        x_{r,j} = y_{r,j} for j = 1..7
    for every register r in {0,...,4}.
    """
    orig_x = [list(model.v[r * 8 : (r + 1) * 8]) for r in range(5)]
    orig_y = model.y_bits

    subs = {}

    for r in range(5):
        for j in range(1, 8):
            subs[orig_x[r][j]] = orig_y[r][j]

        new_expr = orig_y[r][0]
        for j, ai in enumerate(coeffs, start=1):
            if ai:
                new_expr += orig_y[r][j]
        subs[orig_x[r][0]] = new_expr

    new_R = []
    for i in range(8):
        row = [poly.subs(subs) for poly in model.R_bits[i]]
        new_R.append(row)

    return new_R


def run_exhaustive(steps: int, target_reg: int = 0, target_bit: int = -1):
    model = Tea3Model()
    for _ in range(steps):
        model.step()
    snapshot = model

    if target_bit == -1:
        bits = list(range(8))
    else:
        if not (0 <= target_bit < 8):
            raise ValueError("target_bit must be in [0, 7] or -1")
        bits = [target_bit]

    best_coeffs: tuple | None = None
    best_total: int = -1
    best_R: list | None = None

    for idx, coeffs in enumerate(iproduct([0, 1], repeat=7)):
        new_R = apply_variable_change(snapshot, coeffs)
        label = "".join(map(str, coeffs))
        print(f"\n[{idx:03d}] (a1,a2,a3,a4,a5,a6,a7) = {label}")

        total = 0
        for j in bits:
            poly = new_R[target_reg][j]
            pp = pretty_print(poly)
            print(f"  R[{target_reg}][{j}] = {pp}")
            total += count_monomials(poly)

        if best_total == -1 or total < best_total:
            best_total = total
            best_coeffs = coeffs
            best_R = new_R

    print("\nBest variable change (fewest total monomials)")
    print(f"  (a1,a2,a3,a4,a5,a6,a7) = ({', '.join(map(str, best_coeffs))})")
    print(f"  Total monomials: {best_total}")
    print("  Polynomial:\n")
    for j in bits:
        poly = best_R[target_reg][j]
        print(f"R[{target_reg}][{j}] = {pretty_print(poly)}")