#!/usr/bin/env python3

import operator, functools
from z3 import *
import my_utils

"""
rows with dots are partial products:

     aaaa
b ___....
b __...._
b _....__
b ....___

width = OUTPUT_SIZE (INPUT_SIZE*2-1)
height = INPUT_SIZE

"""

def factor(poly):
    INPUT_SIZE=size=my_utils.ceil_binlog(poly)
    OUTPUT_SIZE=INPUT_SIZE*2

    a=BitVec('a', INPUT_SIZE)
    b=BitVec('b', INPUT_SIZE)

    # partial products
    p=[BitVec('p_%d' % i, OUTPUT_SIZE) for i in range(INPUT_SIZE)]

    s=Solver()

    for i in range(INPUT_SIZE):
        # if there is a bit in b[], assign shifted a[] padded with zeroes at left/right
        # if there is no bit in b[], let p[] be zero

        # Concat() is for glueling together bitvectors (of different widths)
        # BitVecVal() is constant of specific width

        if i==0:
            s.add(p[i] == If((b>>i)&1==1, Concat(BitVecVal(0, OUTPUT_SIZE-i-INPUT_SIZE), a), 0))
        else:
            s.add(p[i] == If((b>>i)&1==1, Concat(BitVecVal(0, OUTPUT_SIZE-i-INPUT_SIZE), a, BitVecVal(0, i)), 0))

    # form expression like s.add(p[0] ^ p[1] ^ ... ^ p[OUTPUT_SIZE-1] == poly)
    # replace operator.xor to operator.add to factorize numbers:
    s.add(functools.reduce (operator.xor, p)==poly)

    # we are not interesting in outputs like these:
    s.add(a!=1)
    s.add(b!=1)

    if s.check()==unsat:
        print ("prime factor: 0x%x or %s" % (poly, my_utils.poly_to_str(poly)))
        return

    m=s.model()
    #print ("sat, a=0x%x, b=0x%x" % (m[a].as_long(), m[b].as_long()))
    # factor results recursively:
    factor(m[a].as_long())
    factor(m[b].as_long())

poly=int(sys.argv[1], 16)

print ("starting. poly=0x%x or %s" % (poly, my_utils.poly_to_str(poly)))
factor(poly)