from enum import Enum
from nmigen import *
from nmigen.hdl.rec import *
from nmigen.asserts import *
from nmigen.cli import main
from pipe_stage import PipeStage

# NEXT TASKS (IMMEDIATE)
# get the testbenches working and model checking done
# formal proof (k-induction and/or ABC-PDR)
# figure out how the Altera gearbox design works (number of pipeline stages and what each does)


class GearboxBusLayout(Layout):
    def __init__(self, *, upstream_width, downstream_width): # the * forces keyword args
        super().__init__([
            # DATA
            ("data_in",        unsigned(upstream_width)),      # FROM SOURCE
            ("data_out",       unsigned(downstream_width)),     # TO DEST

            # CONTROL
            ("upstream_valid_in",       1),                       # FROM SOURCE
            ("upstream_ready_out",      1),                       # TO SOURCE

            ("downstream_ready_in",       1),                       # FROM DEST
            ("downstream_valid_out",      1),                       # TO DEST
        ])

class GearboxBus(Record):
    def __init__(self, *, upstream_width, downstream_width):
        super().__init__(GearboxBusLayout(upstream_width=upstream_width, downstream_width=downstream_width))


class ArbitraryGearbox(Elaboratable):
        def __init__(self, *, upstream_width, downstream_width):
            self.upstream_width = upstream_width
            self.downstream_width = downstream_width

            self.bus = GearboxBus(upstream_width=upstream_width, downstream_width=downstream_width)

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

            upstream_width = self.upstream_width
            downstream_width = self.downstream_width

            shift_reg_len        = upstream_width + downstream_width
            shift_reg            = Signal(shift_reg_len)
            valid_bit_count = Signal(range(shift_reg_len+1))


            upstream_transaction = Signal(1)
            downstream_transaction = Signal(1)

            m.d.comb += upstream_transaction.eq(  self.bus.upstream_ready_out  & self.bus.upstream_valid_in)
            m.d.comb += downstream_transaction.eq(self.bus.downstream_ready_in & self.bus.downstream_valid_out)

            m.d.comb += self.bus.upstream_ready_out.eq(shift_reg_len - valid_bit_count >= upstream_width)
            m.d.comb += self.bus.downstream_valid_out.eq(              valid_bit_count >= downstream_width)

            m.d.comb += self.bus.data_out.eq(shift_reg[0:downstream_width])

            with m.If(upstream_transaction == 1):
                with m.If(downstream_transaction == 1): # upstream and downstream
                    # the simultaneous transaction is hard, but updating the valid count isn't
                    m.d.sync += valid_bit_count.eq(valid_bit_count + upstream_width - downstream_width)

                    # the updated shift register is composed of two parts, one from the old bits,
                    # and one from the new bits.
                    #
                    # First, we shift out downstream_width bits from the bottom of the register
                    # What remains is (shift_reg >> downstream_width)
                    #
                    # To fit the new bits in, we must shift them at the correct location so no valid
                    # bit is overwritten and no junk bits are left:
                    # (self.bus.data_in << (valid_bit_count - downstream_width))

                    m.d.sync += shift_reg.eq((self.bus.data_in << (valid_bit_count - downstream_width))| (shift_reg >> downstream_width))

                with m.Else():                          # upstream, no downstream
                    m.d.sync += valid_bit_count.eq(valid_bit_count + upstream_width)
                    # we shift in bits on the high side of the shift register
                    m.d.sync += shift_reg.eq((self.bus.data_in << valid_bit_count) | shift_reg)
            with m.Else():
                with m.If(downstream_transaction == 1): # no upstream, downstream
                    m.d.sync += valid_bit_count.eq(valid_bit_count - downstream_width)
                    # we get rid of the low bits that go out the bus during this transaction
                    m.d.sync += shift_reg.eq(shift_reg[downstream_width:])


            return m



class IdempotentGearbox(Elaboratable):
        def __init__(self, *, width):
            self.width = width

            self.bus = GearboxBus(upstream_width=width, downstream_width=width)

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

            m.submodules.pipe = pipe = PipeStage(width=self.width, registered_ready=False)

            m.d.comb += pipe.upstream_valid_in.eq(self.bus.upstream_valid_in)
            m.d.comb += pipe.downstream_ready_in.eq(self.bus.downstream_ready_in)
            m.d.comb += pipe.upstream_data_in.eq(self.bus.data_in)

            m.d.comb += self.bus.downstream_valid_out.eq(pipe.downstream_valid_out)
            m.d.comb += self.bus.data_out.eq(pipe.downstream_data_out)
            m.d.comb += self.bus.upstream_ready_out.eq(pipe.upstream_ready_out)


            return m



# This is non-synthesizable but is intended to provide a model for non-unitary ratio gearbox
# in order to do formal verification and model-checking against both the unitary-ratio gearbox
# and also the arbitrary ratio gearbox.

class GoldenGearboxModel(Elaboratable):
        def __init__(self, *, upstream_width, downstream_width, sim_memory_size):
            self.upstream_width = upstream_width
            self.downstream_width = downstream_width
            self.sim_memory_size = sim_memory_size

            self.memory = Signal(sim_memory_size)

            self.bus = GearboxBus(upstream_width=upstream_width, downstream_width=downstream_width)

        def elaborate(self, platform):
            m = Module()
            write_ptr = Signal(range(self.sim_memory_size))
            read_ptr = Signal(range(self.sim_memory_size))

            upstream_width = self.upstream_width
            downstream_width = self.downstream_width

            # these are to make fake flow control signals
            fake_memory = self.upstream_width + self.downstream_width
            fake_occupied = Signal(range(1 + self.upstream_width + self.downstream_width))

            upstream_transaction = Signal(1)
            downstream_transaction = Signal(1)

            m.d.comb += upstream_transaction.eq(  self.bus.upstream_ready_out  & self.bus.upstream_valid_in)
            m.d.comb += downstream_transaction.eq(self.bus.downstream_ready_in & self.bus.downstream_valid_out)

            m.d.comb += self.bus.upstream_ready_out.eq(fake_memory - fake_occupied >= upstream_width)
            m.d.comb += self.bus.downstream_valid_out.eq(              fake_occupied >= downstream_width)

            with m.If(upstream_transaction == 1):
                with m.If(downstream_transaction == 1): # upstream and downstream
                    # the simultaneous transaction is hard, but updating the valid count isn't
                    m.d.sync += fake_occupied.eq(fake_occupied + upstream_width - downstream_width)
                    m.d.sync += read_ptr.eq(read_ptr + self.downstream_width)
                    m.d.sync += write_ptr.eq(write_ptr + self.upstream_width)

                    with m.Switch(write_ptr):
                            for cand_wptr in range(0, self.sim_memory_size):
                                with m.Case(cand_wptr):

                                    m.d.sync += self.memory.eq(

                                        Cat(self.memory.bit_select(0, cand_wptr),
                                            self.bus.data_in))


                with m.Else():                          # upstream, no downstream
                    m.d.sync += fake_occupied.eq(fake_occupied + upstream_width)
                    m.d.sync += write_ptr.eq(write_ptr + self.upstream_width)

                    with m.Switch(write_ptr):
                            for cand_wptr in range(0, self.sim_memory_size):
                                with m.Case(cand_wptr):

                                    m.d.sync += self.memory.eq(

                                        Cat(self.memory.bit_select(0, cand_wptr),
                                            self.bus.data_in))

            with m.Else():
                with m.If(downstream_transaction == 1): # no upstream, downstream
                    m.d.sync += fake_occupied.eq(fake_occupied - downstream_width)
                    m.d.sync += read_ptr.eq(read_ptr + self.downstream_width)

            m.d.comb += self.bus.data_out.eq(self.memory.bit_select(read_ptr, self.downstream_width))


            return m



class DummyPlug(Elaboratable):

    #def __init__(self):



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

        upstream_width   = 1
        downstream_width = 3
        sim_memory_size = 8

        expiration_date = int(sim_memory_size/upstream_width)-1

        number_of_writes = Signal(range(1024))
        wrapped = Signal(1)

        m.submodules.gearbox = gearbox = ArbitraryGearbox(upstream_width=upstream_width, downstream_width=downstream_width)
        m.submodules.golden = golden = GoldenGearboxModel(upstream_width=upstream_width, downstream_width=downstream_width, sim_memory_size=sim_memory_size)

        m.d.comb += gearbox.bus.data_in.eq(AnySeq(upstream_width))
        m.d.comb += gearbox.bus.downstream_ready_in.eq(AnySeq(1))
        m.d.comb += gearbox.bus.upstream_valid_in.eq(AnySeq(1))

        m.d.comb += golden.bus.data_in.eq(gearbox.bus.data_in)
        m.d.comb += golden.bus.downstream_ready_in.eq(gearbox.bus.downstream_ready_in)
        m.d.comb += golden.bus.upstream_valid_in.eq(gearbox.bus.upstream_valid_in)

        with m.If(gearbox.bus.upstream_valid_in & gearbox.bus.upstream_ready_out == 1):
            m.d.sync += number_of_writes.eq(number_of_writes + 1)
            with m.If(number_of_writes > expiration_date):
                m.d.sync += wrapped.eq(1)

        with m.If(wrapped == 0):
            m.d.comb += Assert(gearbox.bus.upstream_ready_out == golden.bus.upstream_ready_out)
            m.d.comb += Assert(gearbox.bus.downstream_valid_out == golden.bus.downstream_valid_out)

            with m.If(gearbox.bus.downstream_valid_out == 1):
                m.d.comb += Assert(gearbox.bus.data_out == golden.bus.data_out)


        return m



if __name__ == '__main__':
    baka =DummyPlug()
    main(baka)
    #platform.build(DummyPlug())