from amaranth import *
from amaranth.lib import data

class DoubleDabble(Elaboratable):
    def __init__(self):
        # Ports
        self.i = Signal(6)
        self.outLo = Signal(4)
        self.outHi = Signal(4)

        self.done = Signal(1, reset = 0)

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

        ddlayout = data.StructLayout({
            "base": 6,
            "low": 4, 
            "high": 4
        })

        scratchpad_storage = Signal(ddlayout)
        scratchpad = data.View(ddlayout, scratchpad_storage)

        ddInput = self.i
        prevInput = Signal(6)
        shifts = Signal(range(ddInput.width))

        with m.If(self.done == 0):
            with m.If(prevInput != ddInput):
                m.d.sync += ()
                # m.d.sync += [
                #     shifts.eq(0),
                #     prevInput.eq(ddInput),
                #     scratchpad.base.eq(ddInput)]

            with m.Else():
                with m.If(shifts != ddInput.width):
                    with m.If(scratchpad.high > 5):
                        m.d.sync += scratchpad.high.eq(scratchpad.high + 3)

                    with m.If(scratchpad.low > 5):
                        m.d.sync += scratchpad.low.eq(scratchpad.low + 3)

                    m.d.sync += [
                        shifts.eq(shifts + 1),
                        scratchpad.eq(scratchpad.as_value() << 1)]

                with m.Else():
                    m.d.comb += [
                        self.done.eq(1),
                        self.outHi.eq(scratchpad.high),
                        self.outLo.eq(scratchpad.low)]

        with m.Else():
            with m.If(prevInput != ddInput):
                m.d.comb += self.done.eq(0)

        return m

if __name__ == "__main__":
    from amaranth.sim import *

    dut = DoubleDabble()

    def dd_ut(bcd, bin):
        yield bcd.i.eq(bin)

        for _ in range(300):
            yield Tick()
            yield Settle()

    def proc():
        yield from dd_ut(dut, 12)

    sim = Simulator(dut)
    sim.add_clock(1e-6)
    sim.add_sync_process(proc)

    with sim.write_vcd("../sim/doubledabble.vcd", 'w'):
        sim.run()