from amaranth import *
from amaranth.sim import Simulator
from amaranth.back import rtlil, verilog
import amaranth_playground


class ClockSwitcher(Elaboratable):
    """Dummy Amaranth clock multiplexer module."""

    def __init__(self, registered=False):
        self.registered = registered
        self.sel = Signal(1)

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

        m.domains += ClockDomain("clk3")

        if self.registered:
            pass
        else:
            with m.If(self.sel):
                m.d.comb += ClockSignal("clk3").eq(ClockSignal("clk1"))
            with m.Else():
                m.d.comb += ClockSignal("clk3").eq(ClockSignal("clk2"))

        return m


sw = DomainRenamer({
    "clk1": "foo",
    "clk2": "bar",
    "clk3": "baz"
})(ClockSwitcher())


async def testbench(ctx):
    await ctx.delay(18e-6)
    ctx.set(sw.sel, 1)
    await ctx.delay(20e-6)

sim = Simulator(sw)
sim.add_clock(1e-6, domain="foo")
sim.add_clock(1.05e-6, domain="bar")
sim.add_testbench(testbench)

with amaranth_playground.show_waveforms(sim):
    sim.run()

amaranth_playground.show_verilog(verilog.convert(sw, ports=[sw.sel, ClockSignal("baz")]))