from turtle import width
from typing import Callable
from myhdl import *
from constants import RESET_ACTIVE
from myhdl_wrap import Myhdl_Wrapper

import random
from random import randint

random.seed(63)

class Multiply(Myhdl_Wrapper):
    def __init__(self):
        super().__init__()

    # Main code, this is the actual logic
    @staticmethod
    @block
    def Multiply(clk : Signal, reset : Signal, a : Signal, b : Signal, prod : Signal, is_signed : Signal, ready : Signal, valid : Signal):  # this must be the same name as the class name

        width_a = len(a)
        width_b = len(b)
        width_prod = len(a) + len(b)

        a0 = Signal(modbv(0)[width_a:])
        b0 = Signal(modbv(0)[width_b:])
        a1 = Signal(modbv(0)[width_a:])
        b1 = Signal(modbv(0)[width_b:])
        prod_raw = Signal(modbv(0)[width_prod:])

        @instance
        def logic():
            while True:
                yield valid.posedge
                a0.next = a
                b0.next = b

                yield clk.posedge
                ready.next = False

                if not is_signed:
                    a1.next = a0
                    b1.next = b0
                else:
                    if(a0.signed() < 0):
                        a1.next = -a0.signed()
                    else:
                        a1.next = a0
                    
                    if(b0.signed() < 0):
                        b1.next = -b0.signed()
                    else:
                        b1.next = b0

                prod_raw.next = a1 * b1

                if is_signed:
                    if (a0[width_a-1] ^ b0[width_b-1]) == 1:
                        prod.next = -prod_raw
                    else:
                        prod.next = prod_raw
                    pass
                else:
                    prod.next = prod_raw

                ready.next = True

        return logic


    @block
    def tb(self, func: Callable):
        reset = Signal(False)
        clk = Signal(bool(0))
        a = Signal(intbv(0)[32:])
        b = Signal(intbv(0)[32:])
        prod = Signal(modbv(0)[64:])
        is_signed = Signal(bool(0))
        ready = Signal(bool(1))
        valid = Signal(bool(0))

        a_signed = intbv(0, -int(a.max/2), int(a.max/2)-1)
        b_signed = intbv(0, -int(b.max/2), int(b.max/2)-1)


        dut = func(clk=clk, reset=reset, a=a, b=b, prod=prod, is_signed=is_signed, ready=ready, valid=valid)

        @always(delay(2))
        def clock_gen():
            clk.next = not clk

        @instance
        def monitor():
            # print(f"ready: {ready}, valid: {valid}")
            # print("-"*20)
            while True:
                yield ready, valid
                # print(f"ready: {ready}, valid: {valid}")

        @instance
        def stimulus():
            yield clk.negedge

            # unsigned test
            for i in range(30):
                a_int = randint(a.min, a.max-1)
                b_int = randint(b.min, b.max-1)
                a.next = a_int
                b.next = b_int

                valid.next = True
                yield ready.negedge
                valid.next = False

                yield ready.posedge
                yield delay(1)
                assert int(prod) == a_int * b_int, "Unsigned multiplication error random"

            # unsigned extremes test
            for a_int in (a.min, a.max-1):
                for b_int in (b.min, b.max-1):
                    a.next = a_int
                    b.next = b_int

                    valid.next = True
                    yield ready.negedge
                    valid.next = False

                    yield ready.posedge
                    yield delay(1)

                    assert int(prod) == a_int * b_int, "Unsigned multiplication error extreme"


            is_signed.next = True
            yield clk.negedge

            for i in range(30):
                a_int = randint(a_signed.min, a_signed.max-1)
                b_int = randint(b_signed.min, b_signed.max-1)
                
                a_signed[:] = a_int
                b_signed[:] = b_int
                a.next = a_signed.unsigned()
                b.next = b_signed.unsigned()

                valid.next = True
                yield ready.negedge
                valid.next = False

                yield ready.posedge
                yield delay(1)

                # print('-'*30)
                # print(prod)
                # print(a_int * b_int)

                assert int(prod.signed()) == a_int * b_int, "Signed multiplication error random"


            for a_int in (a_signed.min, 0, a_signed.max-1):
                for b_int in (b_signed.min, 0, b_signed.max-1):

                    a_signed[:] = a_int
                    b_signed[:] = b_int
                    a.next = a_signed.unsigned()
                    b.next = b_signed.unsigned()

                    valid.next = True
                    yield ready.negedge
                    valid.next = False

                    yield ready.posedge
                    yield delay(1)

                    assert int(prod.signed()) == a_int * b_int, "Signed multiplication error extreme"


            raise StopSimulation


        return dut, clock_gen, monitor, stimulus

    def export(self):
        reset = Signal(False)
        clk = Signal(False)
        a = Signal(intbv(0)[32:])
        b = Signal(intbv(0)[32:])
        prod = Signal(modbv(0)[64:])
        is_signed = Signal(False)
        ready = Signal(True)
        valid = Signal(False)

        # assigning signals, kargs only
        self._export(clk=clk, reset=reset, a=a, b=b, prod=prod, is_signed=is_signed, ready=ready, valid=valid)


def test_template_sim():
    hdl = Multiply()
    hdl.sim()

def test_template_cosim():
    hdl = Multiply()
    hdl.export()
    hdl.cosim()