import cocotb

from cocotb.clock import Clock
from cocotb.triggers import Timer, RisingEdge, FallingEdge
from cocotb.handle import Immediate

import math

import random

CLK_PERIOD = 5


@cocotb.test
async def test_sanity(dut):

    data_in = [random.randint(0, 2**32-1) for _ in range(32)]

    async def input_data():
        for word in data_in[:-1]:
            dut.i_data.value = word
            dut.i_valid.value = 1
            dut.i_last.value = 0

            await RisingEdge(dut.i_clk)
            while (not dut.o_ready.value):
                await RisingEdge(dut.i_clk)

        dut.i_data.value = data_in[-1]
        dut.i_valid.value = 1
        dut.i_last.value = 1

        await RisingEdge(dut.i_clk)
        while (not dut.o_ready.value):
            await RisingEdge(dut.i_clk)

        dut.i_valid.value = 0
        dut.i_last.value = 0

    data_out = []

    async def output_data():
        while True:
            if (dut.o_valid.value):
                data_out.append(int(dut.o_data.value))
            await RisingEdge(dut.i_clk)

    cocotb.start_soon(Clock(dut.i_clk, CLK_PERIOD, unit="ns").start())


    dut.i_rst.value = Immediate(0)
    await RisingEdge(dut.i_clk)
    await RisingEdge(dut.i_clk)
    dut.i_rst.value = 1
    await RisingEdge(dut.i_clk)
    await RisingEdge(dut.i_clk)
    dut.i_rst.value = 0
    await RisingEdge(dut.i_clk)
    await RisingEdge(dut.i_clk)

    dut.i_ready.value = 1

    cocotb.start_soon(input_data())
    cocotb.start_soon(output_data())

    await RisingEdge(dut.o_last)

    await RisingEdge(dut.i_clk)

    expected_data_out = []
    for i in range(math.ceil(len(data_in)/4)):
        vals = data_in[i*4:(i+1)*4]

        converted = 0
        for i, val in enumerate(vals):
            converted |= val << (32*i)

        expected_data_out.append(converted)


    assert data_out == expected_data_out

@cocotb.test
async def test_incomplete_last(dut):
    data_in = [random.randint(0, 2**32-1) for _ in range(30)]

    async def input_data():
        for word in data_in[:-1]:
            dut.i_data.value = word
            dut.i_valid.value = 1
            dut.i_last.value = 0

            await RisingEdge(dut.i_clk)
            while (not dut.o_ready.value):
                await RisingEdge(dut.i_clk)

        dut.i_data.value = data_in[-1]
        dut.i_valid.value = 1
        dut.i_last.value = 1

        await RisingEdge(dut.i_clk)
        while (not dut.o_ready.value):
            await RisingEdge(dut.i_clk)

        dut.i_valid.value = 0
        dut.i_last.value = 0

    data_out = []

    async def output_data():
        while True:
            if (dut.o_valid.value):
                data_out.append(int(dut.o_data.value))
            await RisingEdge(dut.i_clk)

    cocotb.start_soon(Clock(dut.i_clk, CLK_PERIOD, unit="ns").start())


    dut.i_rst.value = Immediate(0)
    await RisingEdge(dut.i_clk)
    await RisingEdge(dut.i_clk)
    dut.i_rst.value = 1
    await RisingEdge(dut.i_clk)
    await RisingEdge(dut.i_clk)
    dut.i_rst.value = 0
    await RisingEdge(dut.i_clk)
    await RisingEdge(dut.i_clk)

    dut.i_ready.value = 1

    cocotb.start_soon(input_data())
    cocotb.start_soon(output_data())

    await RisingEdge(dut.o_last)

    await RisingEdge(dut.i_clk)

    expected_data_out = []
    for i in range(math.ceil(len(data_in)/4)):
        vals = data_in[i*4:(i+1)*4]

        converted = 0
        for i, val in enumerate(vals):
            converted |= val << (32*i)

        expected_data_out.append(converted)


    assert data_out == expected_data_out


@cocotb.test
async def test_output_backpressure(dut):
    data_in = [random.randint(0, 2**32-1) for _ in range(32)]

    async def input_data():
        for word in data_in[:-1]:
            dut.i_data.value = word
            dut.i_valid.value = 1
            dut.i_last.value = 0

            await RisingEdge(dut.i_clk)
            while (not dut.o_ready.value):
                await RisingEdge(dut.i_clk)

        dut.i_data.value = data_in[-1]
        dut.i_valid.value = 1
        dut.i_last.value = 1

        await RisingEdge(dut.i_clk)
        while (not dut.o_ready.value):
            await RisingEdge(dut.i_clk)

        dut.i_valid.value = 0
        dut.i_last.value = 0

    data_out = []

    async def output_data():
        while True:
            if (dut.o_valid.value and dut.i_ready.value):
                data_out.append(int(dut.o_data.value))
            await RisingEdge(dut.i_clk)

    cocotb.start_soon(Clock(dut.i_clk, CLK_PERIOD, unit="ns").start())


    dut.i_rst.value = Immediate(0)
    await RisingEdge(dut.i_clk)
    await RisingEdge(dut.i_clk)
    dut.i_rst.value = 1
    await RisingEdge(dut.i_clk)
    await RisingEdge(dut.i_clk)
    dut.i_rst.value = 0
    await RisingEdge(dut.i_clk)
    await RisingEdge(dut.i_clk)

    async def output_backpressure():
        while True:
            await RisingEdge(dut.i_clk)

            dut.i_ready.value = random.randint(0, 100) < 75

    cocotb.start_soon(output_backpressure())
    cocotb.start_soon(input_data())
    cocotb.start_soon(output_data())

    while True:
        await RisingEdge(dut.i_clk)
        if dut.o_last.value and dut.i_ready.value:
            break
    
    await RisingEdge(dut.i_clk)


    expected_data_out = []
    for i in range(math.ceil(len(data_in)/4)):
        vals = data_in[i*4:(i+1)*4]

        converted = 0
        for i, val in enumerate(vals):
            converted |= val << (32*i)

        expected_data_out.append(converted)


    if not data_out == expected_data_out:
        print(data_out)
        print(expected_data_out)

        assert data_out == expected_data_out

@cocotb.test
async def test_output_backpressure_nonfull(dut):
    data_in = [random.randint(0, 2**32-1) for _ in range(30)]

    async def input_data():
        for word in data_in[:-1]:
            dut.i_data.value = word
            dut.i_valid.value = 1
            dut.i_last.value = 0

            await RisingEdge(dut.i_clk)
            while (not dut.o_ready.value):
                await RisingEdge(dut.i_clk)

        dut.i_data.value = data_in[-1]
        dut.i_valid.value = 1
        dut.i_last.value = 1

        await RisingEdge(dut.i_clk)
        while (not dut.o_ready.value):
            await RisingEdge(dut.i_clk)

        dut.i_valid.value = 0
        dut.i_last.value = 0

    data_out = []

    async def output_data():
        while True:
            if (dut.o_valid.value and dut.i_ready.value):
                data_out.append(int(dut.o_data.value))
            await RisingEdge(dut.i_clk)

    cocotb.start_soon(Clock(dut.i_clk, CLK_PERIOD, unit="ns").start())


    dut.i_rst.value = Immediate(0)
    await RisingEdge(dut.i_clk)
    await RisingEdge(dut.i_clk)
    dut.i_rst.value = 1
    await RisingEdge(dut.i_clk)
    await RisingEdge(dut.i_clk)
    dut.i_rst.value = 0
    await RisingEdge(dut.i_clk)
    await RisingEdge(dut.i_clk)

    async def output_backpressure():
        while True:
            await RisingEdge(dut.i_clk)

            dut.i_ready.value = random.randint(0, 100) < 75

    cocotb.start_soon(output_backpressure())
    cocotb.start_soon(input_data())
    cocotb.start_soon(output_data())

    while True:
        await RisingEdge(dut.i_clk)
        if dut.o_last.value and dut.i_ready.value:
            break
    
    await RisingEdge(dut.i_clk)


    expected_data_out = []
    for i in range(math.ceil(len(data_in)/4)):
        vals = data_in[i*4:(i+1)*4]

        converted = 0
        for i, val in enumerate(vals):
            converted |= val << (32*i)

        expected_data_out.append(converted)


    if not data_out == expected_data_out:
        print(data_out)
        print(expected_data_out)

        assert data_out == expected_data_out


@cocotb.test
async def test_input_nonvalid(dut):
    data_in = [random.randint(0, 2**32-1) for _ in range(32)]

    async def input_data():
        for word in data_in[:-1]:
            dut.i_data.value = word
            dut.i_valid.value = 1
            dut.i_last.value = 0

            await RisingEdge(dut.i_clk)
            while (not dut.o_ready.value):
                await RisingEdge(dut.i_clk)

            if random.randint(0, 10) > 7:
                dut.i_valid.value = 0
                await RisingEdge(dut.i_clk)


        dut.i_data.value = data_in[-1]
        dut.i_valid.value = 1
        dut.i_last.value = 1

        await RisingEdge(dut.i_clk)
        while (not dut.o_ready.value):
            await RisingEdge(dut.i_clk)

        dut.i_valid.value = 0
        dut.i_last.value = 0

    data_out = []

    async def output_data():
        while True:
            if (dut.o_valid.value and dut.i_ready.value):
                data_out.append(int(dut.o_data.value))
            await RisingEdge(dut.i_clk)

    cocotb.start_soon(Clock(dut.i_clk, CLK_PERIOD, unit="ns").start())


    dut.i_rst.value = Immediate(0)
    await RisingEdge(dut.i_clk)
    await RisingEdge(dut.i_clk)
    dut.i_rst.value = 1
    await RisingEdge(dut.i_clk)
    await RisingEdge(dut.i_clk)
    dut.i_rst.value = 0
    await RisingEdge(dut.i_clk)
    await RisingEdge(dut.i_clk)

    dut.i_ready.value = 1

    cocotb.start_soon(input_data())
    cocotb.start_soon(output_data())

    while True:
        await RisingEdge(dut.i_clk)
        if dut.o_last.value and dut.i_ready.value:
            break
    
    await RisingEdge(dut.i_clk)


    expected_data_out = []
    for i in range(math.ceil(len(data_in)/4)):
        vals = data_in[i*4:(i+1)*4]

        converted = 0
        for i, val in enumerate(vals):
            converted |= val << (32*i)

        expected_data_out.append(converted)


    if not data_out == expected_data_out:
        print(data_out)
        print(expected_data_out)

        assert data_out == expected_data_out

@cocotb.test
async def test_input_nonvalid_output_nonready(dut):
    data_in = [random.randint(0, 2**32-1) for _ in range(30)]

    async def input_data():
        for word in data_in[:-1]:
            dut.i_data.value = word
            dut.i_valid.value = 1
            dut.i_last.value = 0

            await RisingEdge(dut.i_clk)
            while (not dut.o_ready.value):
                await RisingEdge(dut.i_clk)

            if random.randint(0, 10) > 7:
                dut.i_valid.value = 0
                await RisingEdge(dut.i_clk)


        dut.i_data.value = data_in[-1]
        dut.i_valid.value = 1
        dut.i_last.value = 1

        await RisingEdge(dut.i_clk)
        while (not dut.o_ready.value):
            await RisingEdge(dut.i_clk)

        dut.i_valid.value = 0
        dut.i_last.value = 0

    data_out = []

    async def output_data():
        while True:
            if (dut.o_valid.value and dut.i_ready.value):
                data_out.append(int(dut.o_data.value))
            await RisingEdge(dut.i_clk)

    cocotb.start_soon(Clock(dut.i_clk, CLK_PERIOD, unit="ns").start())


    dut.i_rst.value = Immediate(0)
    await RisingEdge(dut.i_clk)
    await RisingEdge(dut.i_clk)
    dut.i_rst.value = 1
    await RisingEdge(dut.i_clk)
    await RisingEdge(dut.i_clk)
    dut.i_rst.value = 0
    await RisingEdge(dut.i_clk)
    await RisingEdge(dut.i_clk)

    async def output_backpressure():
        while True:
            await RisingEdge(dut.i_clk)

            dut.i_ready.value = random.randint(0, 100) < 75

    cocotb.start_soon(output_backpressure())
    cocotb.start_soon(input_data())
    cocotb.start_soon(output_data())

    while True:
        await RisingEdge(dut.i_clk)
        if dut.o_last.value and dut.i_ready.value:
            break
    
    await RisingEdge(dut.i_clk)


    expected_data_out = []
    for i in range(math.ceil(len(data_in)/4)):
        vals = data_in[i*4:(i+1)*4]

        converted = 0
        for i, val in enumerate(vals):
            converted |= val << (32*i)

        expected_data_out.append(converted)


    if not data_out == expected_data_out:
        print(data_out)
        print(expected_data_out)

        assert data_out == expected_data_out