from migen import *
from sync_serdes import *
from migen.genlib.fifo import SyncFIFO
from migen.build.platforms.sinara import kasli
from migen.genlib.misc import WaitTimer
from kasli_crg import KasliCRG
from eem_helpers import generate_pads
from uart import UART
from io_loopback import SingleIOLoopback


SEPARATOR = Constant(0b0101)

class SingleSerDesLoopBack(Module):
    def __init__(self, io_pad, sys_clk_freq, debug):
        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 = SingleLineTX()
        self.submodules.rx = SingleLineRX()

        # Primary adjustment to master-slave bitslip
        self.submodules.slave_aligner = SlaveAligner()

        # Optimal delay solver
        self.submodules.delay_solver = DelayOptimizer()

        # Debugging readers
        if debug:
            self.submodules.bitslip_reader = BitSlipReader()
            self.submodules.post_align_reader = BitSlipReader()
            self.submodules.phase_reader = PhaseReader()

        # The actual channel
        self.submodules.channel = SingleIOLoopback(io_pad)

        # 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
            self.tx.txdata.eq(0b00100),

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

            # 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),
        ]

        # Route deserializer to phase_reader & the delay tap optimizer
        self.comb += [
            self.slave_aligner.start.eq(0),
            self.delay_solver.start.eq(0),

            # RXDATA for aligner & optimzer
            self.slave_aligner.loopback_rxdata.eq(self.rx.rxdata),
            self.delay_solver.loopback_rxdata.eq(self.rx.rxdata),

            # Delay tap value
            self.delay_solver.delay_tap.eq(self.rx.cnt_out),

            # IDELAY delay tap control, such that phase_reader can
            # change tap value after delay measurement
            self.rx.ce.eq(self.delay_solver.inc_en),
            self.rx.master_bitslip.eq(self.slave_aligner.master_bitslip),
            self.rx.slave_bitslip.eq(self.slave_aligner.slave_bitslip),
            self.rx.ld.eq(self.delay_solver.ld),
            self.rx.cnt_in.eq(self.delay_solver.opt_delay_tap),
        ]

        # Debugging logics
        if debug:
            self.comb += [
                # Start the reader initially
                self.bitslip_reader.start.eq(1),
                self.post_align_reader.start.eq(0),
                self.phase_reader.start.eq(0),

                self.bitslip_reader.loopback_rxdata.eq(self.rx.rxdata),
                self.post_align_reader.loopback_rxdata.eq(self.rx.rxdata),
                self.phase_reader.loopback_rxdata.eq(self.rx.rxdata),

                self.phase_reader.delay_tap.eq(self.rx.cnt_out),

                self.rx.ce.eq(
                    self.phase_reader.inc_en |
                    self.delay_solver.inc_en
                ),
                self.rx.master_bitslip.eq(
                    self.bitslip_reader.bitslip |
                    self.slave_aligner.master_bitslip |
                    self.post_align_reader.bitslip
                ),
                self.rx.slave_bitslip.eq(
                    self.bitslip_reader.bitslip |
                    self.slave_aligner.slave_bitslip |
                    self.post_align_reader.bitslip
                ),
            ]

        if debug:
            delay_tap_count = Signal(6)
            bitslip_count = Signal(3)
            post_align_bitslip_count = Signal(3)

        rx_intra_aligned = Signal()
        select_odd = Signal()
        rxdata_decimated = Signal(5)
        rxdata_pulse_read = Signal(10)

        self.comb += [
            If(select_odd,
                rxdata_decimated.eq(self.rx.rxdata[1::2]),
            ).Else(
                rxdata_decimated.eq(self.rx.rxdata[::2]),
            ),
        ]

        if not debug:
            release_fsm = FSM(reset_state="WAIT_ALIGNER")
            self.submodules += release_fsm

            release_fsm.act("WAIT_ALIGNER",
                self.slave_aligner.start.eq(1),
                If(self.slave_aligner.done,
                    NextState("WAIT_DELAY_OPT"),
                ),
            )

            release_fsm.act("WAIT_DELAY_OPT",
                self.delay_solver.start.eq(1),
                If(self.delay_solver.done,
                    NextValue(select_odd, self.delay_solver.select_odd),
                    NextState("INTRA_ALIGN_DONE"),
                ),
            )

            release_fsm.act("INTRA_ALIGN_DONE",
                rx_intra_aligned.eq(1),
                NextState("WAIT_TX_ZERO"),
            )

            release_fsm.act("WAIT_TX_ZERO",
                If(rxdata_decimated == 0,
                    NextState("WAIT_PULSE"),
                ),
            )

            release_fsm.act("WAIT_PULSE",
                If(rxdata_decimated != 0,
                    NextValue(rxdata_pulse_read, rxdata_decimated),
                    NextState("NEXT_RXDATA"),
                ),
            )

            next_rxdata = Signal(5)

            release_fsm.act("NEXT_RXDATA",
                NextValue(next_rxdata, rxdata_decimated),
                NextState("WRITE_PULSE_UPPER"),
            )

            release_fsm.act("WRITE_PULSE_UPPER",
                If(self.tx_fifo.writable,
                    self.tx_fifo.we.eq(1),
                    self.tx_fifo.din.eq(0),
                    NextState("WRITE_PULSE_LOWER"),
                ),
            )

            release_fsm.act("WRITE_PULSE_LOWER",
                If(self.tx_fifo.writable,
                    self.tx_fifo.we.eq(1),
                    self.tx_fifo.din.eq(rxdata_pulse_read),
                    NextState("WRITE_NEXT_UPPER"),
                ),
            )

            release_fsm.act("WRITE_NEXT_UPPER",
                If(self.tx_fifo.writable,
                    self.tx_fifo.we.eq(1),
                    self.tx_fifo.din.eq(0),
                    NextState("WRITE_NEXT_LOWER"),
                ),
            )

            release_fsm.act("WRITE_NEXT_LOWER",
                If(self.tx_fifo.writable,
                    self.tx_fifo.we.eq(1),
                    self.tx_fifo.din.eq(next_rxdata),
                    NextState("TERMINATE"),
                ),
            )

            release_fsm.act("TERMINATE",
                NextState("TERMINATE"),
            )

        else:
            fsm = FSM(reset_state="WAIT_DONE")
            self.submodules += fsm

            fsm.act("WAIT_DONE",
                If(self.bitslip_reader.done,
                    NextState("WRITE_UPPER"),
                ),
            )

            fsm.act("WRITE_UPPER",
                # Exist state if all results are sent
                If(bitslip_count == 5,
                    NextState("START_SLAVE_ALIGNER"),
                ).Elif(self.tx_fifo.writable,
                    self.tx_fifo.we.eq(1),
                    self.tx_fifo.din.eq(self.bitslip_reader.data_result[bitslip_count][8:]),
                    NextState("WRITE_LOWER"),
                ),
            )

            fsm.act("WRITE_LOWER",
                self.tx_fifo.we.eq(1),
                self.tx_fifo.din.eq(self.bitslip_reader.data_result[bitslip_count][:8]),
                NextValue(bitslip_count, bitslip_count + 1),
                NextState("WRITE_UPPER"),
            )

            fsm.act("START_SLAVE_ALIGNER",
                self.slave_aligner.start.eq(1),
                # self.rx.ce.eq(self.delay_optimizer.inc_en),
                If(self.slave_aligner.done,
                    NextState("WRITE_DONE_UPPER"),
                ).Else(
                    NextState("START_SLAVE_ALIGNER")
                )
            )

            fsm.act("WRITE_DONE_UPPER",
                self.post_align_reader.start.eq(1),
                If(self.tx_fifo.writable,
                    self.tx_fifo.we.eq(1),
                    self.tx_fifo.din.eq(0xFF),
                    NextState("WRITE_DONE_LOWER")
                )
            )

            fsm.act("WRITE_DONE_LOWER",
                self.post_align_reader.start.eq(1),
                If(self.tx_fifo.writable,
                    self.tx_fifo.we.eq(1),
                    self.tx_fifo.din.eq(0xFF),
                    NextState("REREAD_BITSLIP")
                )
            )

            fsm.act("REREAD_BITSLIP",
                If(self.post_align_reader.done,
                    NextState("REWRITE_UPPER"),
                ).Else(
                    NextState("REREAD_BITSLIP"),
                ),
            )

            fsm.act("REWRITE_UPPER",
                # Exist state if all results are sent
                If(post_align_bitslip_count == 5,
                    NextState("WRITE_B2P_DIVIDER_UPPER"),
                ).Elif(self.tx_fifo.writable,
                    self.tx_fifo.we.eq(1),
                    self.tx_fifo.din.eq(self.post_align_reader.data_result[post_align_bitslip_count][8:]),
                    NextState("REWRITE_LOWER"),
                ),
            )

            fsm.act("REWRITE_LOWER",
                self.tx_fifo.we.eq(1),
                self.tx_fifo.din.eq(self.post_align_reader.data_result[post_align_bitslip_count][:8]),
                NextValue(post_align_bitslip_count, post_align_bitslip_count + 1),
                NextState("REWRITE_UPPER"),
            )

            fsm.act("WRITE_B2P_DIVIDER_UPPER",
                self.phase_reader.start.eq(1),
                If(self.tx_fifo.writable,
                    self.tx_fifo.we.eq(1),
                    self.tx_fifo.din.eq(0xFF),
                    NextState("WRITE_B2P_DIVIDER_LOWER"),
                )
            )

            fsm.act("WRITE_B2P_DIVIDER_LOWER",
                self.phase_reader.start.eq(1),
                If(self.tx_fifo.writable,
                    self.tx_fifo.we.eq(1),
                    self.tx_fifo.din.eq(0xFF),
                    NextState("WAIT_PHASE_READER"),
                )
            )

            fsm.act("WAIT_PHASE_READER",
                If(self.phase_reader.done,
                    NextState("WRITE_DELAY_UPPER"),
                ).Else(
                    NextState("WAIT_PHASE_READER"),
                )
            )

            fsm.act("WRITE_DELAY_UPPER",
                If(delay_tap_count == 32,
                    NextState("DELAY_SOLVER_DIVIDER_UPPER")
                ).Elif(self.tx_fifo.writable,
                    self.tx_fifo.we.eq(1),
                    self.tx_fifo.din.eq(self.phase_reader.data_result[delay_tap_count][8:]),
                    NextState("WRITE_DELAY_LOWER"),
                ),
            )

            fsm.act("WRITE_DELAY_LOWER",
                If(self.tx_fifo.writable,
                    self.tx_fifo.we.eq(1),
                    self.tx_fifo.din.eq(self.phase_reader.data_result[delay_tap_count][:8]),
                    NextValue(delay_tap_count, delay_tap_count + 1),
                    NextState("WRITE_DELAY_UPPER")
                ),
            )

            fsm.act("DELAY_SOLVER_DIVIDER_UPPER",
                self.delay_solver.start.eq(1),
                If(self.tx_fifo.writable,
                    self.tx_fifo.we.eq(1),
                    self.tx_fifo.din.eq(0xFF),
                    NextState("DELAY_SOLVER_DIVIDER_LOWER"),
                )
            )

            fsm.act("DELAY_SOLVER_DIVIDER_LOWER",
                self.delay_solver.start.eq(1),
                If(self.tx_fifo.writable,
                    self.tx_fifo.we.eq(1),
                    self.tx_fifo.din.eq(0xFF),
                    NextState("WAIT_DELAY_SOLVER"),
                )
            )

            fsm.act("WAIT_DELAY_SOLVER",
                If(self.delay_solver.done,
                    NextValue(select_odd, self.delay_solver.select_odd),
                    NextState("WRITE_UPPER_ZERO"),
                ).Else(
                    NextState("WAIT_DELAY_SOLVER"),
                )
            )

            fsm.act("WRITE_UPPER_ZERO",
                If(self.tx_fifo.writable,
                    self.tx_fifo.we.eq(1),
                    self.tx_fifo.din.eq(0x00),
                    NextState("WRITE_LOWER_OPT"),
                )
            )

            fsm.act("WRITE_LOWER_OPT",
                If(self.tx_fifo.writable,
                    self.tx_fifo.we.eq(1),
                    self.tx_fifo.din.eq(self.delay_solver.opt_delay_tap),
                    NextState("RESAMPLE_RXDATA_UPPER"),
                ),
            )

            fsm.act("RESAMPLE_RXDATA_UPPER",
                If(self.tx_fifo.writable,
                    self.tx_fifo.we.eq(1),
                    self.tx_fifo.din.eq(self.rx.rxdata[8:]),
                    NextState("RESAMPLE_RXDATA_LOWER"),
                )
            )

            fsm.act("RESAMPLE_RXDATA_LOWER",
                If(self.tx_fifo.writable,
                    self.tx_fifo.we.eq(1),
                    self.tx_fifo.din.eq(self.rx.rxdata[:8]),
                    NextState("INTRA_ALIGN_DONE"),
                )
            )

            fsm.act("INTRA_ALIGN_DONE",
                rx_intra_aligned.eq(1),
                NextState("WAIT_TX_ZERO"),
            )

            fsm.act("WAIT_TX_ZERO",
                If(rxdata_decimated == 0,
                    NextState("WAIT_PULSE"),
                ),
            )

            fsm.act("WAIT_PULSE",
                If(rxdata_decimated != 0,
                    NextValue(rxdata_pulse_read, rxdata_decimated),
                    NextState("WRITE_PULSE_UPPER"),
                )
            )

            fsm.act("WRITE_PULSE_UPPER",
                If(self.tx_fifo.writable,
                    self.tx_fifo.we.eq(1),
                    self.tx_fifo.din.eq(rxdata_pulse_read[8:]),
                    NextState("WRITE_PULSE_LOWER"),
                )
            )

            fsm.act("WRITE_PULSE_LOWER",
                If(self.tx_fifo.writable,
                    self.tx_fifo.we.eq(1),
                    self.tx_fifo.din.eq(rxdata_pulse_read[:8]),
                    NextState("TERMINATE"),
                )
            )

            fsm.act("TERMINATE",
                NextState("TERMINATE"),
            )

        tx_fsm = FSM(reset_state="SEND_TRAINING")
        self.submodules += tx_fsm

        tx_fsm.act("SEND_TRAINING",
            self.tx.txdata.eq(0b00100),
            If(rx_intra_aligned,
                NextState("TX_ZERO"),
            ),
        )

        # Intra-ISERDES alignment done, investigate group delay
        # TX is first set zero for around 16 cycles.
        tx_zero_counter = Signal(5)

        tx_fsm.act("TX_ZERO",
            self.tx.txdata.eq(0b00000),
            If(tx_zero_counter == 15,
                NextState("TX_HIGH"),
            ).Else(
                NextValue(tx_zero_counter, tx_zero_counter + 1),
            ),
        )

        tx_fsm.act("TX_HIGH",
            self.tx.txdata.eq(0b11111),
            NextState("TX_LOW"),
        )

        tx_fsm.act("TX_LOW",
            self.tx.txdata.eq(0b00000),
        )


if __name__ == "__main__":
    platform = kasli.Platform(hw_rev="v2.0")

    # Generate pads for the I/O blocks
    eem = 3
    generate_pads(platform, eem)
    # pads = [
    #     platform.request("dio{}".format(eem), i) for i in range(4)
    # ]
    pad = platform.request("dio{}".format(eem), 0)

    crg = KasliCRG(platform)
    top = SingleSerDesLoopBack(pad, crg.sys_clk_freq, False)

    # 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)