from amaranth import *

from toggle_button import ToggleButton
from pulse_button import PulseButton

class ReBbaRb(Elaboratable):
    def setup_pll(self, module, platform):
        cd_sync = ClockDomain("sync")
        module.domains += cd_sync
        platform.lookup(platform.default_clk).attrs['GLOBAL'] = False

        locked = Signal()
        pllout = Signal()

        module.submodules.pll = Instance("SB_PLL40_PAD",
            p_FEEDBACK_PATH = "SIMPLE",
            p_DIVR          = 0,
            p_DIVF          = 5,
            p_DIVQ          = 1,
            p_FILTER_RANGE  = 1,
            
            o_LOCK          = locked,
            i_RESETB        = 1,
            i_BYPASS        = 0,
            i_PACKAGEPIN    = platform.request("clk12", dir="-").io,
            o_PLLOUTCORE    = cd_sync.clk
        )

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

        self.setup_pll(m, platform)

        ledr0 = platform.request("led_r")
        ledg0 = platform.request("led_g")

        ledr1 = platform.request("led_r", 1)
        ledg1 = platform.request("led_g", 1)
        ledg2 = platform.request("led_g", 2)
        ledg3 = platform.request("led_g", 3)
        ledg4 = platform.request("led_g", 4)
        
        btn1 = platform.request("button", 1)
        btn2 = platform.request("button", 2)
        btn3 = platform.request("button", 3)

        disabled = Signal()

        toggle1 = ToggleButton()
        m.d.comb += toggle1.i.eq(btn2.i)
        m.d.comb += disabled.eq(toggle1.o)
        m.submodules += toggle1

        pulse1 = PulseButton()
        m.d.comb += pulse1.i.eq(btn1.i)
        m.submodules += pulse1

        pulse2 = PulseButton()
        m.d.comb += pulse2.i.eq(btn3.i)
        m.submodules += pulse2

        init_half_freq = int(36e6 // 2)
        half_freq = Signal(range((init_half_freq + 1) * 32), init=init_half_freq // 2)
        timer = Signal(range((init_half_freq + 1) * 32))
    
        init = Signal()

        with m.If(init == 0):
            m.d.sync += ledg0.o.eq(255)
            m.d.sync += half_freq.eq(init_half_freq)
            m.d.sync += init.eq(1)
        
        with m.If(timer == half_freq):
            m.d.sync += ledr0.o.eq(~ledr0.o)
            m.d.sync += ledg0.o.eq(~ledg0.o)
            m.d.sync += timer.eq(0)
        with m.Else():
            with m.If(~disabled):
                m.d.sync += timer.eq(timer + 1)

        with m.If(pulse1.o):
            m.d.sync += half_freq.eq(half_freq << 1)
            m.d.sync += timer.eq(0)
        
        with m.If(pulse2.o):
            m.d.sync += half_freq.eq(half_freq >> 1)
            m.d.sync += timer.eq(0)

        return m

import os
import subprocess
from amaranth.build import *
from amaranth.vendor import LatticeICE40Platform

def _SplitResources(*args, pins, invert=False, conn=None, attrs=None, default_name, dir):
    assert isinstance(pins, (str, list, dict))

    if isinstance(pins, str):
        pins = pins.split()
    if isinstance(pins, list):
        pins = dict(enumerate(pins))

    resources = []
    for number, pin in pins.items():
        ios = [Pins(pin, dir=dir, invert=invert, conn=conn)]
        if attrs is not None:
            ios.append(attrs)
        resources.append(Resource.family(*args, number, default_name=default_name, ios=ios))
    return resources

def UARTResource(*args, rx, tx, rts=None, cts=None, dtr=None, dsr=None, dcd=None, ri=None,
                 conn=None, attrs=None, role=None):
    if any(line is not None for line in (rts, cts, dtr, dsr, dcd, ri)):
        assert role in ("dce", "dte")
    if role == "dte":
        dce_to_dte = "i"
        dte_to_dce = "o"
    else:
        dce_to_dte = "o"
        dte_to_dce = "i"

    io = []
    io.append(Subsignal("rx", Pins(rx, dir="i", conn=conn, assert_width=1)))
    io.append(Subsignal("tx", Pins(tx, dir="o", conn=conn, assert_width=1)))
    if rts is not None:
        io.append(Subsignal("rts", Pins(rts, dir=dte_to_dce, conn=conn, assert_width=1)))
    if cts is not None:
        io.append(Subsignal("cts", Pins(cts, dir=dce_to_dte, conn=conn, assert_width=1)))
    if dtr is not None:
        io.append(Subsignal("dtr", Pins(dtr, dir=dte_to_dce, conn=conn, assert_width=1)))
    if dsr is not None:
        io.append(Subsignal("dsr", Pins(dsr, dir=dce_to_dte, conn=conn, assert_width=1)))
    if dcd is not None:
        io.append(Subsignal("dcd", Pins(dcd, dir=dce_to_dte, conn=conn, assert_width=1)))
    if ri is not None:
        io.append(Subsignal("ri", Pins(ri, dir=dce_to_dte, conn=conn, assert_width=1)))
    if attrs is not None:
        io.append(attrs)
    return Resource.family(*args, default_name="uart", ios=io)

def SPIFlashResources(*args, cs_n, clk, copi, cipo, wp_n=None, hold_n=None,
                      conn=None, attrs=None):
    resources = []

    io_all = []
    if attrs is not None:
        io_all.append(attrs)
    io_all.append(Subsignal("cs",  PinsN(cs_n, dir="o", conn=conn)))
    io_all.append(Subsignal("clk", Pins(clk, dir="o", conn=conn, assert_width=1)))

    io_1x = list(io_all)
    io_1x.append(Subsignal("copi", Pins(copi, dir="o", conn=conn, assert_width=1)))
    io_1x.append(Subsignal("cipo", Pins(cipo, dir="i", conn=conn, assert_width=1)))
    if wp_n is not None and hold_n is not None:
        io_1x.append(Subsignal("wp",   PinsN(wp_n,   dir="o", conn=conn, assert_width=1)))
        io_1x.append(Subsignal("hold", PinsN(hold_n, dir="o", conn=conn, assert_width=1)))
    resources.append(Resource.family(*args, default_name="spi_flash", ios=io_1x,
                                     name_suffix="1x"))

    io_2x = list(io_all)
    io_2x.append(Subsignal("dq", Pins(" ".join([copi, cipo]), dir="io", conn=conn,
                                      assert_width=2)))
    resources.append(Resource.family(*args, default_name="spi_flash", ios=io_2x,
                                     name_suffix="2x"))

    if wp_n is not None and hold_n is not None:
        io_4x = list(io_all)
        io_4x.append(Subsignal("dq", Pins(" ".join([copi, cipo, wp_n, hold_n]), dir="io", conn=conn,
                                          assert_width=4)))
        resources.append(Resource.family(*args, default_name="spi_flash", ios=io_4x,
                                         name_suffix="4x"))

    return resources

def LEDResources(*args, **kwargs):
    return _SplitResources(*args, **kwargs, default_name="led", dir="o")

def ButtonResources(*args, **kwargs):
    return _SplitResources(*args, **kwargs, default_name="button", dir="i")

class IceBreakerPlatform(LatticeICE40Platform):
    device = "iCE40UP5K"
    package = "SG48"
    default_clk = "clk12"
    resources = [
        Resource("clk12", 0, Pins("35", dir="i"), Clock(12e6), Attrs(GLOBAL=True, IO_STANDARD="SB_LVCMOS")),

        *LEDResources(pins="11 37", invert=True, attrs=Attrs(IO_STANDARD="SB_LVCMOS")),
        # Semantic aliases
        Resource("led_r", 0, PinsN("11", dir="o"), Attrs(IO_STANDARD="SB_LVCMOS")),
        Resource("led_g", 0, PinsN("37", dir="o"), Attrs(IO_STANDARD="SB_LVCMOS")),

        *ButtonResources(pins="10", invert=True, attrs=Attrs(IO_STANDARD="SB_LVCMOS")),

        UARTResource(0, rx="6", tx="9", attrs=Attrs(IO_STANDARD="SB_LVTTL", PULLUP=1)),

        *SPIFlashResources(0, cs_n="16", clk="15", copi="14", cipo="17", wp_n="12", hold_n="13", attrs=Attrs(IO_STANDARD="SB_LVCMOS")),

        *LEDResources(pins={2: "7", 3: "1", 4: "2", 5: "8", 6: "3"}, conn=("pmod", 0), attrs=Attrs(IO_STANDARD="SB_LVCMOS")),
        # Semantic aliases
        Resource("led_r", 1, Pins("7", dir="o", conn=("pmod", 0)), Attrs(IO_STANDARD="SB_LVCMOS")),
        Resource("led_g", 1, Pins("1", dir="o", conn=("pmod", 0)), Attrs(IO_STANDARD="SB_LVCMOS")),
        Resource("led_g", 2, Pins("2", dir="o", conn=("pmod", 0)), Attrs(IO_STANDARD="SB_LVCMOS")),
        Resource("led_g", 3, Pins("8", dir="o", conn=("pmod", 0)), Attrs(IO_STANDARD="SB_LVCMOS")),
        Resource("led_g", 4, Pins("3", dir="o", conn=("pmod", 0)), Attrs(IO_STANDARD="SB_LVCMOS")),
        *ButtonResources(pins={1: "9", 2: "4", 3: "10"}, conn=("pmod", 0), attrs=Attrs(IO_STANDARD="SB_LVCMOS")),
    ]

    connectors = [
        Connector("pmod", 0, " 4  2 47 45 - -  3 48 46 44 - -"), # PMOD1A
        Connector("pmod", 1, "43 38 34 31 - - 42 36 32 28 - -"), # PMOD1B
        Connector("pmod", 2, "27 25 21 19 - - 26 23 20 18 - -"), # PMOD2
    ]

    def toolchain_program(self, products, name):
        iceprog = os.environ.get("ICEPROG", "iceprog")
        with products.extract("{}.bin".format(name)) as bitstream_filename:
            subprocess.check_call([iceprog, bitstream_filename])

import os
import subprocess

platform = IceBreakerPlatform()
platform.build(ReBbaRb(), do_program=True, verbose=True)