from migen import *
from sync_serdes import MultiLineRX, MultiLineTX
from migen.genlib.fifo import SyncFIFO
from migen.build.platforms.sinara import kasli, efc
from kasli_crg import TransceiverCRG
from eem_helpers import generate_pads
from uart import UART
from io_loopback import SingleIOLoopback, IOLoopBack


class MultiSerDesLoopBack(Module):
    def __init__(self, io_pads, sys_clk_freq):
        self.uart_rx = Signal()
        self.uart_tx = Signal()

        self.submodules.uart = UART(round((115200/sys_clk_freq)*2**32))
        self.comb += [
            self.uart.phy_rx.eq(self.uart_rx),
            self.uart_tx.eq(self.uart.phy_tx),
        ]

        self.submodules.tx = MultiLineTX()
        self.submodules.rx = MultiLineRX()

        # The actual channel
        self.submodules.channel = IOLoopBack(io_pads)

        # Attach FIFO to UART TX, send rate is too slow w.r.t sysclk
        self.submodules.tx_fifo = SyncFIFO(8, 64)

        self.comb += [
            # Repetitively send 0b00100
            # Note: Replicate() doesn't work, 0b00100 is lowered into 3'd4
            # I need 5'd4 for the replicate operator in Verilog
            self.tx.txdata.eq(0b00100001000010000100),

            # Loopback channel
            self.channel.i.eq(self.tx.ser_out),
            self.channel.t.eq(self.tx.t_out),
            self.rx.ser_in_no_dly.eq(self.channel.o),

            # TX path
            self.uart.tx_data.eq(self.tx_fifo.dout),
            self.uart.tx_stb.eq(self.tx_fifo.readable),
            self.tx_fifo.re.eq(self.uart.tx_ack),

            # Just start RX alignment, no reason to wait
            self.rx.start.eq(1),
        ]

        self.submodules.rx_fsm = FSM(reset_state="WAIT_GROUP_ALIGN")

        sampled_rxdata = Array(Signal(20) for _ in range(16))
        sample_idx = Signal(4)

        self.rx_fsm.act("WAIT_GROUP_ALIGN",
            If(self.rx.err,
                NextState("WRITE_ERR_UPPER"),
            ).Elif(self.rx.rxdata == 0b11111111111111111111,
                NextValue(sampled_rxdata[0], self.rx.rxdata),
                NextValue(sample_idx, 1),
                NextState("SAMPLE_RXDATA"),
            ),
        )

        self.rx_fsm.act("WRITE_ERR_UPPER",
            If(self.tx_fifo.writable,
                self.tx_fifo.we.eq(1),
                self.tx_fifo.din.eq(0b01010101),
                NextState("WRITE_ERR_LOWER"),
            ),
        )

        self.rx_fsm.act("WRITE_ERR_LOWER",
            If(self.tx_fifo.writable,
                self.tx_fifo.we.eq(1),
                self.tx_fifo.din.eq(0b10101010),
                NextState("TERMINATE"),
            ),
        )

        self.rx_fsm.act("SAMPLE_RXDATA",
            If(sample_idx == 15,
                NextValue(sample_idx, 0),
                NextState("WRITE_PATTERN_FIRST_UPPER"),
            ).Else(
                NextValue(sampled_rxdata[sample_idx], self.rx.rxdata),
                NextValue(sample_idx, sample_idx + 1),
            ),
        )

        self.rx_fsm.act("WRITE_PATTERN_FIRST_UPPER",
            If(sample_idx == 15,
                NextState("TERMINATE"),
            ).Elif(self.tx_fifo.writable,
                self.tx_fifo.we.eq(1),
                self.tx_fifo.din.eq(sampled_rxdata[sample_idx][8:10]),
                NextState("WRITE_PATTERN_FIRST_LOWER"),
            ),
        )

        self.rx_fsm.act("WRITE_PATTERN_FIRST_LOWER",
            If(self.tx_fifo.writable,
                self.tx_fifo.we.eq(1),
                self.tx_fifo.din.eq(sampled_rxdata[sample_idx][:8]),
                NextState("WRITE_PATTERN_SECOND_UPPER"),
            ),
        )

        self.rx_fsm.act("WRITE_PATTERN_SECOND_UPPER",
            If(self.tx_fifo.writable,
                self.tx_fifo.we.eq(1),
                self.tx_fifo.din.eq(sampled_rxdata[sample_idx][18:20]),
                NextState("WRITE_PATTERN_SECOND_LOWER"),
            ),
        )

        self.rx_fsm.act("WRITE_PATTERN_SECOND_LOWER",
            If(self.tx_fifo.writable,
                self.tx_fifo.we.eq(1),
                self.tx_fifo.din.eq(sampled_rxdata[sample_idx][10:18]),
                NextValue(sample_idx, sample_idx + 1),
                NextState("WRITE_PATTERN_FIRST_UPPER"),
            ),
        )

        self.rx_fsm.act("TERMINATE",
            NextState("TERMINATE"),
        )

        self.submodules.tx_fsm = FSM(reset_state="SEND_TRAINING")
        
        self.tx_fsm.act("SEND_TRAINING",
            self.tx.txdata.eq(0b00100001000010000100),
            If(self.rx.align_done,
                NextState("SEND_ZERO"),
            ),
        )

        send_zero_duration = Signal(2)

        self.tx_fsm.act("SEND_ZERO",
            self.tx.txdata.eq(0),
            If(send_zero_duration == 0b11,
                NextState("SEND_PULSE"),
            ).Else(
                NextValue(send_zero_duration, send_zero_duration + 1),
            ),
        )

        self.tx_fsm.act("SEND_PULSE",
            self.tx.txdata.eq(0b11111111111111111111),
            NextState("WAIT_GROUP_ALIGN"),
        )

        self.tx_fsm.act("WAIT_GROUP_ALIGN",
            self.tx.txdata.eq(0),
            If(self.rx.delay_done,
                NextState("SEND_ARB_DATA1"),
            ),
        )

        self.tx_fsm.act("SEND_ARB_DATA1",
            self.tx.txdata.eq(0b00111001110011100111),
            NextState("SEND_ARB_DATA2"),
        )

        self.tx_fsm.act("SEND_ARB_DATA2",
            self.tx.txdata.eq(0),
            NextState("SEND_ARB_DATA3"),
        )

        self.tx_fsm.act("SEND_ARB_DATA3",
            self.tx.txdata.eq(0xDEADB),
            NextState("SEND_ARB_DATA4"),
        )

        self.tx_fsm.act("SEND_ARB_DATA4",
            self.tx.txdata.eq(0xBCAFE),
            NextState("TERMINATE"),
        )

        self.tx_fsm.act("TERMINATE",
            self.tx.txdata.eq(0),
            NextState("TERMINATE"),
        )


if __name__ == "__main__":
    import argparse

    parser = argparse.ArgumentParser()
    parser.add_argument("platform")
    args = parser.parse_args()

    platform_dict = {
        "kasli": kasli.Platform(hw_rev="v2.0"),
        "efc": efc.Platform(),
    }

    sysclk_name = {
        "kasli": "clk125_gtp",
        "efc": "gtp_clk",
    }

    platform = platform_dict[args.platform]
    sysclk = platform.request(sysclk_name[args.platform])

    # Generate pads for the I/O blocks
    # Using EEM1 for both as both EFC and Kasli has EEM1
    # EEM1 are not interconnected
    eem = 1
    generate_pads(platform, eem)
    pads = [
        platform.request("dio{}".format(eem), i+4) for i in range(4)
    ]
    # pad = platform.request("dio{}".format(eem), 0)

    crg = TransceiverCRG(platform, sysclk)
    top = MultiSerDesLoopBack(pads, crg.sys_clk_freq)

    # Wire up UART core to the pads
    uart_pads = platform.request("serial")
    top.comb += [
        top.uart_rx.eq(uart_pads.rx),
        uart_pads.tx.eq(top.uart_tx),
    ]

    top.submodules += crg
    platform.build(top)