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.clk1 = Signal()
        self.clk2 = Signal()
        self.out = Signal()
        self.sel = Signal()

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

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

        return m


m = Module()
sw = ClockSwitcher()
m.submodules += sw
# Without this line, clk3 is "multiply-driven"
m.domains += ClockDomain("clk3")


m.d.comb += sw.clk1.eq(ClockSignal("clk1"))
m.d.comb += sw.clk2.eq(ClockSignal("clk2"))
m.d.comb += ClockSignal("clk3").eq(sw.out)

count = Signal(4)
m.d.clk3 += count.eq(count + 1)


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

sim = Simulator(m)
sim.add_clock(1e-6, domain="clk1")
sim.add_clock(1.05e-6, domain="clk2")
sim.add_testbench(testbench)

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

amaranth_playground.show_verilog(verilog.convert(m, ports=[sw.sel]))
# amaranth_playground.show_rtlil(rtlil.convert(m, ports=[sw.sel]))