"""Generation of GTKWave documents with nmutil.gtkw"""

from nmigen import Elaboratable, Signal, Module, Cat
from nmigen.back.pysim import Simulator
from nmigen.cli import rtlil
from math import log2

from vcd.gtkw import GTKWSave, GTKWColor
from nmutil.gtkw import write_gtkw


class Shifter(Elaboratable):
    """Simple sequential shifter

    * "Prev" port:

        * ``p_data_i``: value to be shifted

        * ``p_shift_i``: shift amount

        * ``op__sdir``: shift direction (0 = left, 1 = right)

        * ``p_valid_i`` and ``p_ready_o``: handshake

    * "Next" port:

        * ``n_data_o``: shifted value

        * ``n_valid_o`` and ``n_ready_i``: handshake
    """
    def __init__(self, width):
        self.width = width
        """data width"""
        self.p_data_i = Signal(width)
        self.p_shift_i = Signal(width)
        self.op__sdir = Signal()
        self.p_valid_i = Signal()
        self.p_ready_o = Signal()
        self.n_data_o = Signal(width)
        self.n_valid_o = Signal()
        self.n_ready_i = Signal()

    def elaborate(self, _):
        m = Module()

        # the control signals
        load = Signal()
        shift = Signal()
        direction = Signal()
        # the data flow
        shift_in = Signal(self.width)
        shift_left_by_1 = Signal(self.width)
        shift_right_by_1 = Signal(self.width)
        next_shift = Signal(self.width)
        # the register
        shift_reg = Signal(self.width, reset_less=True)
        # build the data flow
        m.d.comb += [
            # connect input and output
            shift_in.eq(self.p_data_i),
            self.n_data_o.eq(shift_reg),
            # generate shifted views of the register
            shift_left_by_1.eq(Cat(0, shift_reg[:-1])),
            shift_right_by_1.eq(Cat(shift_reg[1:], 0)),
        ]
        # choose the next value of the register according to the
        # control signals
        # default is no change
        m.d.comb += next_shift.eq(shift_reg)
        with m.If(load):
            m.d.comb += next_shift.eq(shift_in)
        with m.Elif(shift):
            with m.If(direction):
                m.d.comb += next_shift.eq(shift_right_by_1)
            with m.Else():
                m.d.comb += next_shift.eq(shift_left_by_1)

        # register the next value
        m.d.sync += shift_reg.eq(next_shift)

        # Shift counter
        shift_width = int(log2(self.width)) + 1
        next_count = Signal(shift_width)
        count = Signal(shift_width, reset_less=True)
        m.d.sync += count.eq(next_count)

        with m.FSM():
            with m.State("IDLE"):
                m.d.comb += [
                    # keep p.ready_o active on IDLE
                    self.p_ready_o.eq(1),
                    # keep loading the shift register and shift count
                    load.eq(1),
                    next_count.eq(self.p_shift_i),
                ]
                # capture the direction bit as well
                m.d.sync += direction.eq(self.op__sdir)
                with m.If(self.p_valid_i):
                    # Leave IDLE when data arrives
                    with m.If(next_count == 0):
                        # short-circuit for zero shift
                        m.next = "DONE"
                    with m.Else():
                        m.next = "SHIFT"
            with m.State("SHIFT"):
                m.d.comb += [
                    # keep shifting, while counter is not zero
                    shift.eq(1),
                    # decrement the shift counter
                    next_count.eq(count - 1),
                ]
                with m.If(next_count == 0):
                    # exit when shift counter goes to zero
                    m.next = "DONE"
            with m.State("DONE"):
                # keep n_valid_o active while the data is not accepted
                m.d.comb += self.n_valid_o.eq(1)
                with m.If(self.n_ready_i):
                    # go back to IDLE when the data is accepted
                    m.next = "IDLE"

        return m

    def __iter__(self):
        yield self.op__sdir
        yield self.p_data_i
        yield self.p_shift_i
        yield self.p_valid_i
        yield self.p_ready_o
        yield self.n_ready_i
        yield self.n_valid_o
        yield self.n_data_o

    def ports(self):
        return list(self)


def write_gtkw_direct():
    """Write a formatted GTKWave "save" file, using vcd.gtkw directly"""
    # hierarchy path, to prepend to signal names
    dut = "top.shf."
    # color styles
    style_input = GTKWColor.orange
    style_output = GTKWColor.yellow
    style_debug = GTKWColor.red
    with open("test_shifter_direct.gtkw", "wt") as gtkw_file:
        gtkw = GTKWSave(gtkw_file)
        gtkw.comment("Auto-generated by " + __file__)
        gtkw.dumpfile("test_shifter.vcd")
        # set a reasonable zoom level
        # also, move the marker to an interesting place
        gtkw.zoom_markers(-22.9, 10500000)
        gtkw.trace(dut + "clk")
        # place a comment in the signal names panel
        gtkw.blank("Shifter Demonstration")
        with gtkw.group("prev port"):
            gtkw.trace(dut + "op__sdir", color=style_input)
            # demonstrates using decimal base (default is hex)
            gtkw.trace(dut + "p_data_i[7:0]", color=style_input,
                       datafmt='dec')
            gtkw.trace(dut + "p_shift_i[7:0]", color=style_input,
                       datafmt='dec')
            gtkw.trace(dut + "p_valid_i", color=style_input)
            gtkw.trace(dut + "p_ready_o", color=style_output)
        with gtkw.group("debug"):
            gtkw.blank("Some debug statements")
            # change the displayed name in the panel
            gtkw.trace("top.zero", alias='zero delay shift',
                       color=style_debug)
            gtkw.trace("top.interesting", color=style_debug)
            gtkw.trace("top.test_case", alias="test case", color=style_debug)
            gtkw.trace("top.msg", color=style_debug)
        with gtkw.group("internal"):
            gtkw.trace(dut + "fsm_state")
            gtkw.trace(dut + "count[3:0]")
            gtkw.trace(dut + "shift_reg[7:0]", datafmt='dec')
        with gtkw.group("next port"):
            gtkw.trace(dut + "n_data_o[7:0]", color=style_output,
                       datafmt='dec')
            gtkw.trace(dut + "n_valid_o", color=style_output)
            gtkw.trace(dut + "n_ready_i", color=style_input)


def test_shifter():
    """Simulate the Shifter to generate some traces,
     as well as the GTKWave documents"""
    m = Module()
    m.submodules.shf = dut = Shifter(8)
    print("Shifter port names:")
    for port in dut:
        print("-", port.name)
    # generate RTLIL
    # try "proc; show" in yosys to check the data path
    il = rtlil.convert(dut, ports=dut.ports())
    with open("test_shifter.il", "w") as f:
        f.write(il)

    # write the GTKWave project file, directly
    write_gtkw_direct()

    # Describe a GTKWave document

    # Style for signals, classes and groups
    gtkwave_style = {
        # Root selector. Gives default attributes for every signal.
        '': {'base': 'dec'},
        # color the traces, according to class
        # class names are not hardcoded, they are just strings
        'in': {'color': 'orange'},
        'out': {'color': 'yellow'},
        # signals in the debug group have a common color and module path
        'debug': {'module': 'top', 'color': 'red'},
        # display a different string replacing the signal name
        'test_case': {'display': 'test case'},
    }

    # DOM style description for the trace pane
    gtkwave_desc = [
        # simple signal, without a class
        # even so, it inherits the top-level root attributes
        'clk',
        # comment
        {'comment': 'Shifter Demonstration'},
        # collapsible signal group
        ('prev port', [
            # attach a class style for each signal
            ('op__sdir', 'in'),
            ('p_data_i[7:0]', 'in'),
            ('p_shift_i[7:0]', 'in'),
            ('p_valid_i', 'in'),
            ('p_ready_o', 'out'),
        ]),
        # Signals in a signal group inherit the group attributes.
        # In this case, a different module path and color.
        ('debug', [
            {'comment': 'Some debug statements'},
            # inline attributes, instead of a class name
            ('zero', {'display': 'zero delay shift'}),
            'interesting',
            'test_case',
            'msg',
        ]),
        ('internal', [
            'fsm_state',
            'count[3:0]',
            'shift_reg[7:0]',
        ]),
        ('next port', [
            ('n_data_o[7:0]', 'out'),
            ('n_valid_o', 'out'),
            ('n_ready_i', 'in'),
        ]),
    ]

    write_gtkw("test_shifter.gtkw", "test_shifter.vcd",
               gtkwave_desc,  gtkwave_style,
               module="top.shf", loc=__file__, marker=10500000)

    sim = Simulator(m)
    sim.add_clock(1e-6)

    # demonstrates adding extra debug signal traces
    # they end up in the top module
    #
    zero = Signal()  # mark an interesting place
    #
    # demonstrates string traces
    #
    # display a message when the signal is high
    # the low level is just an horizontal line
    interesting = Signal(decoder=lambda v: 'interesting!' if v else '')
    # choose between alternate strings based on numerical value
    test_cases = ['', '13>>2', '3<<4', '21<<0']
    test_case = Signal(8, decoder=lambda v: test_cases[v])
    # hack to display arbitrary strings, like debug statements
    msg = Signal(decoder=lambda _: msg.str)
    msg.str = ''

    def send(data, shift, direction):
        # present input data and assert valid_i
        yield dut.p_data_i.eq(data)
        yield dut.p_shift_i.eq(shift)
        yield dut.op__sdir.eq(direction)
        yield dut.p_valid_i.eq(1)
        yield
        # wait for p.ready_o to be asserted
        while not (yield dut.p_ready_o):
            yield
        # show current operation operation
        if direction:
            msg.str = f'{data}>>{shift}'
        else:
            msg.str = f'{data}<<{shift}'
        # force dump of the above message by toggling the
        # underlying signal
        yield msg.eq(0)
        yield msg.eq(1)
        # clear input data and negate p.valid_i
        yield dut.p_valid_i.eq(0)
        yield dut.p_data_i.eq(0)
        yield dut.p_shift_i.eq(0)
        yield dut.op__sdir.eq(0)

    def receive(expected):
        # signal readiness to receive data
        yield dut.n_ready_i.eq(1)
        yield
        # wait for n.valid_o to be asserted
        while not (yield dut.n_valid_o):
            yield
        # read result
        result = yield dut.n_data_o
        # negate n.ready_i
        yield dut.n_ready_i.eq(0)
        # check result
        assert result == expected
        # finish displaying the current operation
        msg.str = ''
        yield msg.eq(0)
        yield msg.eq(1)

    def producer():
        # 13 >> 2
        yield from send(13, 2, 1)
        # 3 << 4
        yield from send(3, 4, 0)
        # 21 << 0
        # use a debug signal to mark an interesting operation
        # in this case, it is a shift by zero
        yield interesting.eq(1)
        yield from send(21, 0, 0)
        yield interesting.eq(0)

    def consumer():
        # the consumer is not in step with the producer, but the
        # order of the results are preserved
        # 13 >> 2 = 3
        yield test_case.eq(1)
        yield from receive(3)
        # 3 << 4 = 48
        yield test_case.eq(2)
        yield from receive(48)
        # 21 << 0 = 21
        yield test_case.eq(3)
        # you can look for the rising edge of this signal to quickly
        # locate this point in the traces
        yield zero.eq(1)
        yield from receive(21)
        yield zero.eq(0)
        yield test_case.eq(0)

    sim.add_sync_process(producer)
    sim.add_sync_process(consumer)
    sim_writer = sim.write_vcd(
        "test_shifter.vcd",
        # include additional signals in the trace dump
        traces=[zero, interesting, test_case, msg],
    )
    with sim_writer:
        sim.run()


if __name__ == "__main__":
    test_shifter()
