from myhdl import *
from myhdl_wrap import Myhdl_Wrapper

import random
from random import randint

random.seed(63)

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

    # Main code, this is the actual logic
    @staticmethod
    @block
    def ShiftReg(reset: Signal, clk: Signal, load: Signal, in0: Signal, out0: Signal, left_right: Signal):

        width = len(out0)

        @instance
        def shifter():
            while True:
                yield clk.posedge, reset.negedge

                if not reset:
                    out0.next = load
                else:
                    if not left_right:
                        out0.next[width:1] = out0[width-1:0]
                        out0.next[1:0] = in0
                    else:
                        out0.next[width-1:0] = out0[width:1]
                        out0.next[width:width-1] = in0

        return shifter


    @block
    def tb(self, func):
        reset = Signal(False)
        clk = Signal(bool(0))
        load = Signal(intbv(0xA5)[8:])
        in0 = Signal(bool(0))
        out0 = Signal(modbv(int(load))[8:])
        left_right = Signal(bool(0))

        dut = func(reset=reset, clk=clk, load=load, in0=in0, out0=out0, left_right=left_right)

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

        @instance
        def monitor():
            last_val = out0

            while True:
                yield clk.posedge, reset.negedge
                yield delay(1)

                if reset == False:
                    assert out0 == load, "Is output is not zero at reset"

                else:
                    if not left_right:
                        assert int(out0) == ((last_val << 1) & 0xFF) | int(in0), "Not shifting left correctly"
                    else:
                        assert int(out0) == ((last_val >> 1) & 0xFF) | (int(in0) << 7), "Not shifting rigth correctly"

                last_val = int(out0)

        @instance
        def reset_test():
            yield clk.negedge
            reset.next = True
            while True:
                reset.next = True
                yield delay(randint(25, 28))

                reset.next = False
                yield delay(randint(1,4))

        @instance
        def stimulus():
            for i in range(20):
                yield clk.negedge
                left_right.next = 0
                in0.next = randint(0, 1)

            for i in range(20):
                yield clk.negedge
                left_right.next = 1
                in0.next = randint(0, 1)

            raise StopSimulation


        return dut, clock_gen, monitor, stimulus, reset_test

    def export(self):
        reset = Signal(False)
        clk = Signal(bool(0))
        load = Signal(intbv(0xA5)[8:])
        in0 = Signal(bool(0))
        out0 = Signal(modbv(int(load))[8:])
        left_right = Signal(bool(0))

        self._export(reset=reset, clk=clk, load=load, in0=in0, out0=out0, left_right=left_right)


def test_shift_reg_sim():
    hdl = ShiftReg()
    hdl.sim()

def test_shift_reg_cosim():
    hdl = ShiftReg()
    hdl.export()
    hdl.cosim()