"""Gate-level simulation of the multiplier.

Timing and delays are simulated using a tick system. Each
tick corresponds to a single delta-t (equivalent to the propagation delay of a
basic AND or OR gate).

In each tick, two stages are run. In the first stage, all components compute
their new state for the tick from the state of their inputs, but they do not
update their "public" state until the first stage completes. Once all
components have finished, the second stage is run, where components update
their public state to the newly-computed state.

This two-stage system effectively buffers the component states. The reason for
this is because states for the current tick ought to be based on states from
the previous tick. If those states aren't buffered, it's likely that some
states will be computed based on other new states that were computed in the
same tick, which is not desired.
"""

class Component:
    """A component that updates on every tick."""

    def __init__(self, sim):
        self.sim = sim
        sim.register(self)

    def tick(self):
        """Compute new state, but do not update the value of output()."""
        raise NotImplementedError('Component.tick')

    def tock(self):
        """Update the value of output() based on the state from tick()."""
        raise NotImplementedError('Component.tock')

    def output(self):
        """The "output" state of this component, which other components can
        depend upon for their calculations.

        This is most commonly a single boolean value, but it can sometimes be
        an array of boolean values (for multi-bit values like the multiplier).
        """
        raise NotImplementedError('Component.output')

    def bit(self, index):
        return Split(self.sim, index, self)

    def bits(self, start, end):
        return SplitRange(self.sim, start, end, self)


class Sim:
    def __init__(self):
        self.components = []
        self.tick_counter = 0

    def register(self, component):
        self.components.append(component)

    def tick(self):
        self.tick_counter += 1
        for com in self.components:
            com.tick()
        for com in self.components:
            com.tock()

    def const(self, c):
        return Constant(self, c)

    def link(self):
        return Link(self)

    def and2(self, a, b):
        return And(self, a, b)

    def or2(self, a, b):
        return Or(self, a, b)

    def nand2(self, a, b):
        return Not(self, And(self, a, b))

    def nor2(self, a, b):
        return Not(self, Or(self, a, b))

    def xor2(self, a, b):
        return Or(self, And(self, Not(self, a), b), And(self, a, Not(self, b)))

    def and_n(self, *args):
        if len(args) == 0:
            return self.const(True)
        elif len(args) == 1:
            return args[0]
        else:
            split = len(args) // 2
            return self.and2(self.and_n(*args[:split]), self.and_n(*args[split:]))

    def or_n(self, *args):
        if len(args) == 0:
            return self.const(False)
        elif len(args) == 1:
            return args[0]
        else:
            split = len(args) // 2
            return self.or2(self.or_n(*args[:split]), self.or_n(*args[split:]))

    def mux(self, sel, when_zero, when_one):
        return self.or2(self.and2(Not(self, sel), when_zero), self.and2(sel, when_one))


class Constant(Component):
    def __init__(self, sim, c):
        super().__init__(sim)
        self.c = c

    def set(self, c):
        self.c = c

    def tick(self):
        pass

    def tock(self):
        pass

    def output(self):
        return self.c


class And(Component):
    """Boolean AND-gate with two inputs."""

    def __init__(self, sim, a, b):
        super().__init__(sim)
        self.a = a
        self.b = b
        self.state = False
        self.buffer = False

    def tick(self):
        a = self.a.output()
        b = self.b.output()
        if isinstance(a, list) and isinstance(b, list):
            assert len(a) == len(b)
            self.state = [x & y for x, y in zip(a, b)]
        elif isinstance(a, list):
            self.state = [x & b for x in a]
        elif isinstance(b, list):
            self.state = [a & y for y in b]
        else:
            self.state = a & b

    def tock(self):
        self.buffer = self.state

    def output(self):
        return self.buffer


class Or(Component):
    """Boolean OR-gate with two inputs."""

    def __init__(self, sim, a, b):
        super().__init__(sim)
        self.a = a
        self.b = b
        self.state = False
        self.buffer = False

    def tick(self):
        a = self.a.output()
        b = self.b.output()
        if isinstance(a, list) and isinstance(b, list):
            assert len(a) == len(b)
            self.state = [x | y for x, y in zip(a, b)]
        elif isinstance(a, list):
            self.state = [x | b for x in a]
        elif isinstance(b, list):
            self.state = [a | y for y in b]
        else:
            self.state = a | b

    def tock(self):
        self.buffer = self.state

    def output(self):
        return self.buffer

class Not(Component):
    """A transparent component that immediately inverts its input.

    While a hardware NOT-gate would technically cause some delay, often it can
    be integrated into the gate immediately before or after it in the signal path.
    Therefore, it is usually modeled in software logic as zero-delay.
    """

    def __init__(self, sim, a):
        super().__init__(sim)
        self.a = a

    def tick(self):
        # Instantaneous update
        pass

    def tock(self):
        # Instantaneous update
        pass

    def output(self):
        return not self.a.output()


class Merge(Component):
    """A transparent component that joins single-bit component outputs into a
    multi-bit output."""

    def __init__(self, sim, inputs):
        super().__init__(sim)
        self.inputs = []
        for inp in inputs:
            if isinstance(inp, list):
                self.inputs.extend(inp)
            else:
                self.inputs.append(inp)

    def tick(self):
        # This component is zero-tick / transparent; it doesn't do any
        # gate-level logic. It just changes the logical grouping of bits in the
        # simulation. So the input state is reflected immediately in the output.
        pass

    def tock(self):
        # Same as `tick`
        pass

    def output(self):
        acc = []
        for inp in self.inputs:
            output = inp.output()
            if isinstance(output, list):
                acc.extend(output)
            else:
                acc.append(output)
        return acc


class Split(Component):
    """A transparent component that takes a single bit from a multi-bit output."""

    def __init__(self, sim, index, component):
        super().__init__(sim)
        self.index = index
        self.component = component

    def tick(self):
        # See `Merge.tick`
        pass

    def tock(self):
        # See `Merge.tick`
        pass

    def output(self):
        return self.component.output()[self.index]


class SplitRange(Component):
    """A transparent component that takes a range of bits from a multi-bit output."""
    def __init__(self, sim, start, end, component):
        super().__init__(sim)
        self.start = start
        self.end = end
        self.component = component

    def tick(self):
        # See `Merge.tick`
        pass

    def tock(self):
        # See `Merge.tick`
        pass

    def output(self):
        out = self.component.output()
        if isinstance(out, list):
            return out[self.start:self.end]
        else:
            # prevent errors from uninitialized outputs that aren't yet lists
            return out


class Link(Component):
    """A component that wraps another component, which gets connected after
    construction. Useful for creating cycles in component references.

    Typical usage: The link is created and a reference is stored locally. It's
    passed as an input to another component, in lieu of the actual input
    component which is not yet constructed. Once the input component is
    constructed and available, it is passed to the link via `connect()`.
    """

    def __init__(self, sim):
        super().__init__(sim)
        self.component = None

    def connect(self, component):
        self.component = component

    def tick(self):
        pass
    
    def tock(self):
        pass

    def output(self):
        return self.component.output()


class Mul4x4(Component):
    """The given 4x4 multiplier component.

    Produces 8-bit output from two 4-bit inputs, with a propagation delay of 21
    ticks.

    This also assumes a worst-case throughput of 21 ticks, meaning the input
    must be held steady until the output is ready or else the result is
    invalidated.
    """

    def __init__(self, sim, a, b):
        super().__init__(sim)
        self.a = a
        self.b = b

        # Implement a delay / throughput of 21 ticks by keeping a timer of
        # how long the input is held at the same state.
        self.a_state = None
        self.b_state = None
        self.timer = 0

        self.buffer = unpack(8, 0)

    def tick(self):
        a_state = self.a.output()
        b_state = self.b.output()
        if a_state == self.a_state and b_state == self.b_state:
            # If the input is steady, increment the timer.
            self.timer += 1
        else:
            # If the input changes, reset.
            self.timer = 0
            self.a_state = a_state
            self.b_state = b_state

    def tock(self):
        # If we've reached the target delay, then we can update the buffered state.
        if self.timer >= 21:
            a = pack(self.a_state)
            b = pack(self.b_state)
            self.buffer = unpack(8, a * b)

    def output(self):
        return self.buffer


def pack(bits):
    """Packs the bits into an integer, LSB-first."""

    if isinstance(bits, list):
        acc = 0
        for bit in reversed(bits):
            acc <<= 1
            if bit:
                acc += 1
        return acc
    else:
        return int(bits)


def unpack(n, x):
    """Unpacks the integer x into an array of n bits, LSB-first, truncating on overflow."""

    bits = []
    for _ in range(n):
        bits.append((x & 1) != 0)
        x >>= 1
    return bits


class HalfAdder:
    """Single-bit half adder. Not a component itself, but can construct
    components for sum/carry/propagate/generate from the input bits."""

    def __init__(self, sim, a, b):
        self.sim = sim
        self.a = a
        self.b = b

    def sum(self):
        return self.sim.xor2(self.a, self.b)

    def carry(self):
        return self.sim.and2(self.a, self.b)

    def propagate(self):
        return self.sim.or2(self.a, self.b)

    def generate(self):
        return self.carry()


class FullAdder1:
    def __init__(self, sim, a, b, c):
        self.sim = sim
        self.a = a
        self.b = b
        self.c = c

    def sum(self):
        return self.sim.xor2(self.sim.xor2(self.a, self.b), self.c)
    
    def carry(self):
        sim = self.sim
        return sim.or_n(
            sim.and2(self.a, self.b),
            sim.and2(self.b, self.c),
            sim.and2(self.a, self.c),
        )

class Lookahead:
    def __init__(self, sim, carry_in, ha0, ha1, ha2, ha3):
        self.sim = sim
        self.carry_in = carry_in
        self.ha0 = ha0
        self.ha1 = ha1
        self.ha2 = ha2
        self.ha3 = ha3

    def carry0(self):
        return self.carry_in

    def carry1(self):
        sim = self.sim
        return sim.or_n(
            self.ha0.generate(),
            sim.and_n(
                self.ha0.propagate(),
                self.carry_in
            ),
        )

    def carry2(self):
        sim = self.sim
        return sim.or_n(
            self.ha1.generate(),
            sim.and_n(
                self.ha1.propagate(),
                self.ha0.generate(),
            ), 
            sim.and_n(
                self.ha1.propagate(),
                self.ha0.propagate(),
                self.carry_in,
            ),
        )

    def carry3(self):
        sim = self.sim
        return sim.or_n(
            self.ha2.generate(),
            sim.and_n(
                self.ha2.propagate(),
                self.ha1.generate()
            ),
            sim.and_n(
                self.ha2.propagate(), 
                self.ha1.propagate(),
                self.ha0.generate(),
            ),
            sim.and_n(
                self.ha2.propagate(),
                self.ha1.propagate(),
                self.ha0.propagate(),
                self.carry_in,
            ),
        )

    def carry_out(self):
        sim = self.sim
        return sim.or_n(
            self.ha3.generate(),
            sim.and_n(
                self.ha3.propagate(),
                self.ha2.generate(),
            ),
            sim.and_n(
                self.ha3.propagate(),
                self.ha2.propagate(),
                self.ha1.generate(),
            ),
            sim.and_n(
                self.ha3.propagate(),
                self.ha2.propagate(),
                self.ha1.propagate(),
                self.ha0.generate(),
            ),
            sim.and_n(
                self.ha3.propagate(),
                self.ha2.propagate(),
                self.ha1.propagate(),
                self.ha0.propagate(),
                self.carry_in,
            ),
        )

    def propagate(self):
        sim = self.sim
        return sim.and_n(
            self.ha3.propagate(),
            self.ha2.propagate(),
            self.ha1.propagate(),
            self.ha0.propagate(),
        )

    def generate(self):
        return self.carry_out()


class FullAdder4:
    def __init__(self, sim, a, b, carry_in=None):
        if carry_in is None:
            carry_in = sim.const(False)

        self.sim = sim
        self.ha0 = HalfAdder(sim, Split(sim, 0, a), Split(sim, 0, b))
        self.ha1 = HalfAdder(sim, Split(sim, 1, a), Split(sim, 1, b))
        self.ha2 = HalfAdder(sim, Split(sim, 2, a), Split(sim, 2, b))
        self.ha3 = HalfAdder(sim, Split(sim, 3, a), Split(sim, 3, b))
        self.cla = Lookahead(sim, carry_in, self.ha0, self.ha1, self.ha2, self.ha3)
        self.s0 = sim.xor2(self.ha0.sum(), self.cla.carry0())
        self.s1 = sim.xor2(self.ha1.sum(), self.cla.carry1())
        self.s2 = sim.xor2(self.ha2.sum(), self.cla.carry2())
        self.s3 = sim.xor2(self.ha3.sum(), self.cla.carry3())

    def sum(self):
        return Merge(self.sim, [self.s0, self.s1, self.s2, self.s3])

    def carry_out(self):
        return self.cla.carry_out()

    def propagate(self):
        return self.cla.propagate()

    def generate(self):
        return self.cla.generate()


class FullAdder16:
    """A 16-bit full adder using a tree of lookahead units, with a propagation
    delay of 12 ticks and throughput interval of 9 ticks."""
    def __init__(self, sim, a, b, carry_in=None):
        if carry_in is None:
            carry_in = sim.const(False)

        self.sim = sim
        c0 = sim.link();
        c1 = sim.link();
        c2 = sim.link();
        c3 = sim.link();
        self.fa0 = FullAdder4(sim, SplitRange(sim, 0, 4, a), SplitRange(sim, 0, 4, b), c0)
        self.fa1 = FullAdder4(sim, SplitRange(sim, 4, 8, a), SplitRange(sim, 4, 8, b), c1)
        self.fa2 = FullAdder4(sim, SplitRange(sim, 8, 12, a), SplitRange(sim, 8, 12, b), c2)
        self.fa3 = FullAdder4(sim, SplitRange(sim, 12, 16, a), SplitRange(sim, 12, 16, b), c3)
        self.cla = Lookahead(sim, carry_in, self.fa0, self.fa1, self.fa2, self.fa3)
        c0.connect(self.cla.carry0())
        c1.connect(self.cla.carry1())
        c2.connect(self.cla.carry2())
        c3.connect(self.cla.carry3())

    def sum(self):
        return Merge(self.sim, [self.fa0.sum(), self.fa1.sum(), self.fa2.sum(), self.fa3.sum()])

    def carry_out(self):
        return self.cla.carry_out()


class Mul8x8:
    """8-bit multiplier, made of 4x 4-bit multipliers."""
    def __init__(self, sim, a, b):
        self.sim = sim
        a_low = SplitRange(sim, 0, 4, a)
        a_high = SplitRange(sim, 4, 8, a)
        b_low = SplitRange(sim, 0, 4, b)
        b_high = SplitRange(sim, 4, 8, b)

        self.al_bl = Mul4x4(sim, a_low, b_low)
        self.al_bh = Mul4x4(sim, a_low, b_high)
        self.ah_bl = Mul4x4(sim, a_high, b_low)
        self.ah_bh = Mul4x4(sim, a_high, b_high)

        # The output from the four multipliers is shaped like this:
        #
        # ah_bh[7:4] ah_bh[3:0] al_bl[7:4] al_bl[3:0]
        #            al_bh[7:4] al_bh[3:0]
        #            ah_bl[7:4] ah_bl[3:0]
        #
        # There are 8 bits in the center that have 3 sum terms. We're using an
        # array of 1-bit full adders to compress that into 2 terms so we can
        # use a 16-bit fast adder in the final stage, at the cost of 4 ticks of delay.
        fa_array = [
            FullAdder1(sim, self.al_bl.bit(4), self.al_bh.bit(0), self.ah_bl.bit(0)),
            FullAdder1(sim, self.al_bl.bit(5), self.al_bh.bit(1), self.ah_bl.bit(1)),
            FullAdder1(sim, self.al_bl.bit(6), self.al_bh.bit(2), self.ah_bl.bit(2)),
            FullAdder1(sim, self.al_bl.bit(7), self.al_bh.bit(3), self.ah_bl.bit(3)),
            FullAdder1(sim, self.ah_bh.bit(0), self.al_bh.bit(4), self.ah_bl.bit(4)),
            FullAdder1(sim, self.ah_bh.bit(1), self.al_bh.bit(5), self.ah_bl.bit(5)),
            FullAdder1(sim, self.ah_bh.bit(2), self.al_bh.bit(6), self.ah_bl.bit(6)),
            FullAdder1(sim, self.ah_bh.bit(3), self.al_bh.bit(7), self.ah_bl.bit(7)),
        ]

        self.final_stage = FullAdder16(sim,
            Merge(sim, [
                self.al_bl.bits(0, 4),
                [fa.sum() for fa in fa_array],
                self.ah_bh.bits(4, 8),
            ]),
            Merge(sim, [
                [sim.const(False)] * 5,
                [fa.carry() for fa in fa_array],
                [sim.const(False)] * 3,
            ]),
        )

    def low_product(self):
        """The 8-bit intermediate product of the lower 4 bits of the two
        inputs. Propagation delay: 21 ticks."""
        return self.al_bl

    def high_product(self):
        """The 8-bit intermediate product of the higher 4 bits of the two
        inputs. Propagation delay: 21 ticks."""
        return self.ah_bh

    def product(self):
        """The full 16-bit product of the two inputs. Propagation delay: 37
        ticks; throughput interval: 21 ticks."""
        return self.final_stage.sum()


class MultiMultiplier:
    """The high-level multiplier we are trying to build.

    Detects the size of its input and applies different algorithms for 4-bit,
    8-bit, and 16-bit numbers, so that multiplication of smaller words is
    faster."""

    def __init__(self, sim, a, b):
        self.sim = sim
        #TODO 16 bit mult logic

        # Input is 8-bit if the higher 8 bits of both inputs are zeroes.
        is_8bit = sim.nor2(
            sim.or_n(*[a.bit(x) for x in range(8, 16)]),
            sim.or_n(*[b.bit(x) for x in range(8, 16)]),
        )
        # Input is 4-bit if the higher 12 bits of both inputs are zeroes.
        is_4bit = sim.nor2(
            sim.or_n(*[a.bit(x) for x in range(4, 16)]),
            sim.or_n(*[b.bit(x) for x in range(4, 16)]),
        )

        # TODO plumb selectors to 16-bit sequencer
        a_sel = sim.const(False)
        b_sel = sim.const(False)
        a_mux = sim.mux(a_sel, a.bits(0, 8), a.bits(8, 16))
        b_mux = sim.mux(b_sel, b.bits(0, 8), b.bits(8, 16))
        mult = Mul8x8(sim, a_mux, b_mux)
        self.out_mux = sim.mux(is_4bit, mult.product(), Merge(sim, [mult.low_product(), [sim.const(False)] * 8]))

    def product(self):
        return self.out_mux


TEST_CASES = [
    (0xe, 0xf),
    (0x5, 0x5),
    (0x3f, 0x3f),
    (0x2e, 0x37),
    (0x1f, 0x55),
    (0x77, 0x33),
    (0x78, 0x77),
    #TODO 16-bit
    (0x155, 0x155),
    (0x555, 0x555),
    (0x3e7, 0x0ff),
    (0x5555, 0x5555),
    (0x3e70, 0x0ff0),
    (0x9bbd, 0xc383),
    (0xd557, 0xd555),
]
HOLD_TIME = 50
TIMEOUT = 500


def main():
    sim = Sim()
    x = sim.const(unpack(16, 0))
    y = sim.const(unpack(16, 0))
    mult = MultiMultiplier(sim, x, y)
    product = mult.product()

    for new_x, new_y in TEST_CASES:
        x.set(unpack(16, new_x))
        y.set(unpack(16, new_y))
        expected = new_x * new_y

        held = 0
        ticks = 0
        state = None
        while held < HOLD_TIME and ticks < TIMEOUT:
            sim.tick()
            ticks += 1
            new_state = pack(product.output())
            if new_state == state:
                held += 1
            else:
                held = 0
                state = new_state

        result = pack(product.output())
        print('{} * {} = {}'.format(hex(new_x), hex(new_y), hex(result)))
        if ticks >= TIMEOUT:
            print('ERROR timeout')
        elif result != expected:
            print('ERROR mismatch, expected {}'.format(hex(expected)))
    
        print('Ticks: {}'.format(ticks - held))
        print()


def test_main():
    sim = Sim()
    three = sim.const(unpack(4, 3))
    five = sim.const(unpack(4, 5))
    fa4 = FullAdder4(sim, three, five)
    fa4_sum = fa4.sum()
    mul = Mul4x4(sim, three, five)

    import random
    a = sim.const(unpack(16, random.randrange(32768)))
    b = sim.const(unpack(16, random.randrange(32768)))
    fa16 = FullAdder16(sim, a, b)
    fa16_sum = fa16.sum()

    a2 = sim.const(unpack(8, random.randrange(256)))
    b2 = sim.const(unpack(8, random.randrange(256)))
    mul8 = Mul8x8(sim, a2, b2)
    mul8_product = mul8.product()

    for i in range(50):
        sim.tick()
    print('3 + 5 =', pack(fa4_sum.output()))
    print('3 * 5 =', pack(mul.output()))
    print('{} + {} = {}'.format(pack(a.output()), pack(b.output()), pack(fa16_sum.output())))
    print('{} * {} = {}'.format(pack(a2.output()), pack(b2.output()), pack(mul8_product.output())))
    assert pack(a.output()) + pack(b.output()) == pack(fa16_sum.output())
    assert pack(a2.output()) * pack(b2.output()) == pack(mul8_product.output())


if __name__ == '__main__':
    main()
