import magma as m
m.set_mantle_target("ice40")
import mantle

N = 4
Register4 = mantle.DefineRegister(4)
T = m.Bits(N)

class ShiftRegister(m.Circuit):
    name = "ShiftRegister"
    IO = ["I", m.In(T), "O", m.Out(T), "CLK", m.In(m.Clock)]
    @classmethod
    def definition(io):
        regs = [Register4() for _ in range(N)]
        [m.wire(io.CLK, reg.CLK) for reg in regs]
        m.wire(io.I, getattr(regs[0], "I"))
        m.braid(regs, foldargs={"I":"O"})
        m.wire(regs[-1].O, io.O)


"""
Explanation: Graphviz
Let's play around with the graphviz backend.
Shift register
Let's create a Shift register. This example is taken from the shift_register notebook.
End of explanation
"""


ShiftRegister

"""
Explanation: If Python returns a Circuit in a jupyter notebook,
the circuit is displayed using graphviz.
Now let's visualize our ShiftRegister type.
End of explanation
"""


from functools import reduce

def one_hot_mux(conds, inputs):
    outputs = []
    for cond, inp in zip(conds, inputs):
        outputs.append(inp & m.uint([cond for _ in range(len(inp))]))
    return reduce(lambda x, y: x | y, outputs)


class SimpleALU(m.Circuit):
    name = "SimpleALU"
    IO = ["a", m.In(m.UInt(4)), "b", m.In(m.UInt(4)),
          "opcode", m.In(m.UInt(2)), "out", m.Out(m.UInt(4))]
    
    @classmethod
    def definition(io):
        is_op0 = io.opcode == m.uint(0, n=2)
        is_op1 = io.opcode == m.uint(1, n=2)
        is_op2 = io.opcode == m.uint(2, n=2)
        is_op3 = io.opcode == m.uint(3, n=2)
        op0_out = io.a + io.b
        op1_out = io.a - io.b
        op2_out = io.a
        op3_out = io.b
        m.wire(io.out, one_hot_mux([is_op0, is_op1, is_op2, is_op3], [op0_out, op1_out, op2_out, op3_out]))

SimpleALU

"""
Explanation: Simple ALU
Another example, taken from the Chisel examples folder.
End of explanation
"""


