from nmigen import *
from nmigen.back import rtlil, verilog


READ_PORTS_PER_CHANNEL  = 2
WRITE_PORTS_PER_CHANNEL = 1
CHANNELS                = 2
REGISTERS               = 32


class Gpr(Elaboratable):
    def __init__(self, rwidth=32):
        self.rwidth = rwidth
        self.rdepth = (REGISTERS - 1).bit_length()

        self.i_wr_a = Array(Array(Signal(self.rdepth) for _ in range(WRITE_PORTS_PER_CHANNEL)) for _ in range(CHANNELS))
        self.i_wr_d = Array(Array(Signal(rwidth) for _ in range(WRITE_PORTS_PER_CHANNEL)) for _ in range(CHANNELS))
        self.i_rd_a = Array(Array(Signal(self.rdepth) for _ in range(READ_PORTS_PER_CHANNEL)) for _ in range(CHANNELS))
        self.o_rd_d = Array(Array(Signal(rwidth) for _ in range(READ_PORTS_PER_CHANNEL)) for _ in range(CHANNELS))

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

        select  = Memory(width=CHANNELS, depth=self.rdepth)
        regfile = Array(Memory(width=self.rwidth, depth=self.rdepth) for _ in range(CHANNELS))
        readmux = Array(Array(Array(Signal(self.rwidth)) for _ in range(CHANNELS) for _ in range(READ_PORTS_PER_CHANNEL)) for _ in range(CHANNELS))

        for chan in range(CHANNELS):
            # Write ports
            for port in range(WRITE_PORTS_PER_CHANNEL):
                m.submodules["reg_write_c{}p{}".format(chan, port)] = reg_write_port = regfile[chan].write_port()
                m.submodules["sel_write_c{}p{}".format(chan, port)] = sel_write_port = select.write_port()
                m.d.comb += [
                    reg_write_port.addr.eq(self.i_wr_a[chan][port]),
                    reg_write_port.data.eq(self.i_wr_d[chan][port]),
                    sel_write_port.addr.eq(self.i_wr_a[chan][port]),
                    sel_write_port.data.eq(port)
                ]

            # Read ports
            for port in range(READ_PORTS_PER_CHANNEL):
                m.submodules["reg_read_c{}p{}".format(chan, port)] = reg_read_port = regfile[chan].read_port()
                m.submodules["sel_read_c{}p{}".format(chan, port)] = sel_read_port = select.read_port()

                for otherchan in range(CHANNELS):
                    m.d.comb += readmux[otherchan][port][chan].eq(reg_read_port.data)

                m.d.comb += [
                    reg_read_port.addr.eq(self.i_rd_a[chan][port]),
                    sel_read_port.addr.eq(self.i_rd_a[chan][port]),
                    self.o_rd_d[chan][port].eq(readmux[chan][port][sel_read_port.data])
                ]

        return m


dut = Gpr()

ports = [
    dut.i_wr_a, dut.i_wr_d,
    dut.i_rd_a, dut.o_rd_d
]

print(verilog.convert(dut, ports=ports))