from migen import *
from sync_serdes import PhaseReader, DelayOptimizer, BitSlipReader
import random


def reader_testbench(dut, rxdata_list):
    yield dut.delay_tap.eq(0)
    yield dut.start.eq(1)
    assert (yield dut.stab_timer.wait) == 0

    for i in range(32):
        yield dut.loopback_rxdata.eq(rxdata_list[i])
        yield
        yield
        assert (yield dut.stab_timer.wait) == 1

        # Keep yielding until the DUT gives CE signal
        while (yield dut.inc_en) == 0:
            yield

        # Check that inc_en is deassrted after 1 clock cycle
        yield
        assert (yield dut.inc_en) == 0

        # Load a new tap value
        yield dut.delay_tap.eq(i + 1)
        yield
        # Nothing to check in the READ_TAP state
        yield

    assert(yield dut.done) == 1

    for i in range(32):
        signal = yield dut.data_result[i]
        expected = rxdata_list[i]
        assert signal == expected

    for i in range(200):
        assert (yield dut.inc_en) == 0
        yield

    # Untouched delay: Record should be invariant
    for i in range(32):
        signal = yield dut.data_result[i]
        expected = rxdata_list[i]
        assert signal == expected


def optimal_delay_testbench(dut, pulse_list, cycles, pulse_index, min_delay, max_offset, opt_delay_tap):
    # Start the module
    yield dut.delay_tap.eq(0)
    yield dut.start.eq(1)
    assert (yield dut.stab_timer.wait) == 0

    for i in range(cycles):
        # Pass in a new rxdata for sampling
        # The stab_timer should start waiting after
        yield dut.loopback_rxdata.eq(pulse_list[i])
        yield
        yield
        assert (yield dut.stab_timer.wait) == 1

        # Eventually, the wait will end
        # Either it triggers an increment or a finished signal
        # And we will get the expected pulse location
        # inc_en is pulsed after this is found
        if i == (cycles - 1):
            while (yield dut.done) == 0:
                yield
            break
        else:
            while (yield dut.inc_en) == 0:
                yield
        
        # Then we increment the rxdata index
        yield dut.delay_tap.eq(i + 1)
        yield

    # Fast-forward to the result
    # while (yield dut.done) == 0:
    #     yield
    
    assert (yield dut.done) == 1
    assert (yield dut.expected_pulse) == pulse_index
    assert (yield dut.min_delay) == min_delay
    assert (yield dut.max_offset) == max_offset
    assert (yield dut.opt_delay_tap) == opt_delay_tap

    for _ in range(100):
        yield
    
    # Invariant test: Everything is frozen after done
    assert (yield dut.done) == 1
    assert (yield dut.expected_pulse) == pulse_index
    assert (yield dut.min_delay) == min_delay
    assert (yield dut.max_offset) == max_offset
    assert (yield dut.opt_delay_tap) == opt_delay_tap


def bitslip_reader_tb(dut, rxdata_list):
    # Start the module
    yield dut.start.eq(1)
    assert (yield dut.stab_timer.wait) == 0

    for i in range(5):
        yield dut.loopback_rxdata.eq(rxdata_list[i])
        yield
        yield
        assert (yield dut.stab_timer.wait) == 1

        # Keep yielding until the DUT gives BITSLIP signal
        while (yield dut.bitslip) == 0:
            yield
        
        # There will be 2 BITSLIP pulses
        # Both BITSLIP pulses should last for 1 cycle
        assert (yield dut.bitslip) == 1
        yield
        assert (yield dut.bitslip) == 0
        yield
        assert (yield dut.bitslip) == 1
        yield
        assert (yield dut.bitslip) == 0
    
    assert (yield dut.done) == 1
    # The result in the module should contain all rxdata
    for i, rxdata in enumerate(rxdata_list):
        assert (yield dut.data_result[i]) == rxdata
    
    yield
    yield


# # Random testing for delay reader
# for _ in range(32):
#     rxdata_list = [ random.getrandbits(10) for _ in range(32) ]
#     dut = PhaseReader()
#     run_simulation(dut, reader_testbench(dut, rxdata_list), vcd_name="phase_reader.vcd")

# # Random testing for optimal delay calculation
# # Generate a delay list
# start = random.randint(0, 9)
# start_length = random.randint(1, 10)
# offset = random.randint(4, 5)

# current_index = start
# remaining_length = start_length
# single_pulse_list = []

# expected_index = (current_index + 1) % 10
# expected_length = 10

# for tap in range(32 + offset):
#     single_pulse_list.append(1 << current_index)
#     remaining_length -= 1

#     if remaining_length == 0:
#         current_index = (current_index + 1) % 10
#         remaining_length = 10

# pulse_list = list(single_pulse_list)
# for i in range(offset, 32):
#     pulse_list[i] |= single_pulse_list[i - offset]

# found_start_edge = False
# max_offset = 0

# # Calculate min_delay
# for i, pulse in enumerate(pulse_list):
#     if (pulse & (1 << expected_index)) != 0:
#         if not found_start_edge:
#             min_delay = i
#             found_start_edge = True
#         else:
#             max_offset += 1
#     if (pulse & (1 << expected_index)) == 0 and found_start_edge:
#         cycles = i + 1
#         break

# print(min_delay)
# print(max_offset)
# print(cycles)
# opt_delay = int(min_delay + (max_offset / 2))
# print(opt_delay)

# # Simulate
# dut = DelayOptimizer()
# run_simulation(dut, optimal_delay_testbench(
#     dut, pulse_list, cycles, expected_index,
#     min_delay, max_offset, opt_delay),
#     vcd_name="delay_opt.vcd"
# )

# Random test for bitslip reader
for _ in range(32):
    rxdata_list = [ random.getrandbits(10) for _ in range(5) ]
    dut = BitSlipReader()
    run_simulation(dut, bitslip_reader_tb(dut, rxdata_list), vcd_name="bitslip_reader.vcd")