from functools import partial
from pymoku.moku import Moku
from pymoku.sources import ADC, Node, DAC
from pymoku.registers import RegisterAccess


class _SlotControl(RegisterAccess):
    def __init__(self, parent=None, offset=0, prefix=None, slot_id=None):
        super().__init__(parent, offset, prefix)
        self.slot_id = slot_id
        self.loops = self.reg_prop_list(reg=0, offset=0, length=8,
                                        count=4, mkreg=self.reg_slv)


# RegisterAccess' weren't designed to be modified during their lifetime.  This
# gives us a class seperate from the Moku object that can be discarded.
class _Routing(RegisterAccess):
    def __init__(self, parent, offset=0, prefix=None):
        super().__init__(parent, offset, prefix)
        self.slots = [_SlotControl(self, offset=18 + i, slot_id=i)
                      for i in range(parent.num_slots)]
        self.dacs = self.reg_prop_list(reg=7, offset=0, length=8,
                                       count=4, mkreg=self.reg_slv)


class _ADCMux(Node):
    def __init__(self, moku):
        super().__init__(name='HADC')
        self.moku = moku

    def set_idx(self, idx: tuple):
        self.moku.adc_mode(idx[0])

    def get_idx(self) -> tuple:
        return (self.moku.adc_mode(),)


class MokuPro(Moku):
    DEV_NAME = 'mokupro'
    AXI_BUS_WIDTH = 16
    ADC_BIT_DEPTH = 10
    CBUF_RAM_LENGTH = 0x180000000
    URAM_LENGTH = 0

    def setup_routing(self):
        self.routing = _Routing(self)

        self.adc0 = ADC('InA', partial(self.adc_coeffs, 0), bus_width=4)
        self.adc1 = ADC('InB', partial(self.adc_coeffs, 1), bus_width=4)
        self.adc2 = ADC('InC', partial(self.adc_coeffs, 2), bus_width=4)
        self.adc3 = ADC('InD', partial(self.adc_coeffs, 3), bus_width=4)

        # single instance without edges used as shared resource
        self.adc_mux = _ADCMux(self)
        self.adc_mux.add(node=self.adc0, idx='1.25', bus_width=1)
        self.adc_mux.add(node=self.adc1, idx='1.25', bus_width=1)
        self.adc_mux.add(node=self.adc2, idx='1.25', bus_width=1)
        self.adc_mux.add(node=self.adc3, idx='1.25', bus_width=1)

        # A -> D are in reverse order 0 -> 1
        # TODO config AC will conflict AD when it should transition
        self.adc_mux.add(node=self.adc1, idx='2.5AC', bus_width=2)
        self.adc_mux.add(node=self.adc3, idx='2.5AC', bus_width=2)
        self.adc_mux.add(node=self.adc1, idx='2.5BC', bus_width=2)
        self.adc_mux.add(node=self.adc2, idx='2.5BC', bus_width=2)
        self.adc_mux.add(node=self.adc0, idx='2.5AD', bus_width=2)
        self.adc_mux.add(node=self.adc3, idx='2.5AD', bus_width=2)
        self.adc_mux.add(node=self.adc0, idx='2.5BD', bus_width=2)
        self.adc_mux.add(node=self.adc2, idx='2.5BD', bus_width=2)

        self.adc_mux.add(node=self.adc0, idx='5D', bus_width=4)
        self.adc_mux.add(node=self.adc1, idx='5C', bus_width=4)
        self.adc_mux.add(node=self.adc2, idx='5B', bus_width=4)
        self.adc_mux.add(node=self.adc3, idx='5A', bus_width=4)

        self.ext = ADC('Ext', coeff=lambda: 2**14, bus_width=4)

        self.dac0 = DAC('OutA', self.routing, dac_idx=0, coeff=2**15 - 1)
        self.dac1 = DAC('OutB', self.routing, dac_idx=1, coeff=2**15 - 1)
        self.dac2 = DAC('OutC', self.routing, dac_idx=2, coeff=2**15 - 1)
        self.dac3 = DAC('OutD', self.routing, dac_idx=3, coeff=2**15 - 1)

        self._sources = [self.adc0, self.adc1, self.adc2, self.adc3, self.ext]
        self._outputs = [self.dac0, self.dac1, self.dac2, self.dac3]

        super().setup_routing()

    def attach_slot(self, slot_inst, slot):
        # attach the platform sources to the slot inputs
        for inp in slot_inst.inputs():
            inp.add(self.adc_mux, limit=[self.adc0], idx=0)
            inp.add(self.adc_mux, limit=[self.adc1], idx=1)
            inp.add(self.adc_mux, limit=[self.adc2], idx=2)
            inp.add(self.adc_mux, limit=[self.adc3], idx=3)
            inp.add(self.adc_mux, limit=[self.adc0, self.adc1], idx=(1, 0), bus_width=2)
            inp.add(self.adc_mux, limit=[self.adc2, self.adc3], idx=(3, 2), bus_width=2)
            # these odd index orderings are from the HADC itself
            inp.add(self.adc_mux, idx=(3, 1, 2, 0), bus_width=4)
            inp.add(self.ext, idx=4)
            inp.add(self.ext, idx=(4, 6), bus_width=2)
            inp.add(self.ext, idx=(4, 5, 6, 7), bus_width=4)

        # attach the slout outputs to the actual slot
        for idx, outp in enumerate(slot_inst.outputs()):
            slot_outp = self._slot_outputs[slot][idx]
            slot_outp.clear()
            slot_outp.add(outp, bus_width=1)
            slot_outp.add(outp, bus_width=2)
            slot_outp.add(outp, bus_width=4)

        # attach all the loops to other slots
        loop_off = 8
        for inp in slot_inst.inputs():
            for idx, loop in enumerate(self._loop_muxes_1[slot]):
                # for wider inputs we just duplicate (sample & hold)
                if inp.bus_width == 1:
                    inp.add(loop, idx=idx + loop_off)
                if inp.bus_width == 2:
                    inp.add(loop, idx=(idx + loop_off, idx + loop_off))
                if inp.bus_width == 4:
                    inp.add(loop, idx=(idx + loop_off, idx + loop_off, idx + loop_off, idx + loop_off))

            for idx, loop in enumerate(self._loop_muxes_2[slot]):
                if inp.bus_width == 2:
                    inp.add(loop, idx=(2 * idx + loop_off, 2 * (idx + 1) + loop_off), bus_width=2)
                if inp.bus_width == 4:
                    inp.add(loop, idx=(2 * idx + loop_off, 2 * idx + loop_off,
                                       2 * (idx + 1) + loop_off, 2 * (idx + 1) + loop_off), bus_width=2)

            for idx, loop in enumerate(self._loop_muxes_4[slot]):
                inp.add(loop, idx=tuple(range(4 * idx + loop_off, 4 * (idx + 1) + loop_off, 1)), bus_width=4)

    def _update_coeffs(self, hwstate):
        self._cal = dict(
            adc=[hwstate['adc0']['_coeff'],
                 hwstate['adc1']['_coeff'],
                 hwstate['adc2']['_coeff'],
                 hwstate['adc3']['_coeff']],
            dac=[hwstate['dac0']['_coeff'],
                 hwstate['dac1']['_coeff'],
                 hwstate['dac2']['_coeff'],
                 hwstate['dac3']['_coeff']])

    def adc_coeffs(self, ch):
        return self._cal['adc'][ch]

    def dac_coeffs(self, ch):
        return self._cal['dac'][ch]

    def set_frontend(self, ch, enable=None, fiftyr=None, gain=None,
                     dc=None, bw=None, source=None):
        # 'enable' has no effect
        p = {}
        r = {}

        if fiftyr is not None:
            r['fiftyr'] = fiftyr

        if gain is not None:
            assert gain in [0, -20, -40]
            r['gain'] = gain

        if bw is not None:
            assert bw in [300, 600]
            r['bw'] = bw

        if dc is not None:
            r['dc'] = dc

        if source is not None:
            r['source'] = source

        p['adc' + str(ch)] = r

        self.modify_hardware(**p)

    def adc_mode(self, mode=None):
        return self.modify_hardware(adc_mode=mode)['adc_mode']

    def hardware_defaults(self):
        self.set_frontend(0, fiftyr=True, gain=0, dc=True)
        self.set_frontend(1, fiftyr=True, gain=0, dc=True)
        self.set_frontend(2, fiftyr=True, gain=0, dc=True)
        self.set_frontend(3, fiftyr=True, gain=0, dc=True)

        self.modify_hardware(adc_mode='1.25')
        # TODO not sure why the default aren't sticking in the firmware
        self.modify_hardware(dac0=dict(gain=0),
                             dac1=dict(gain=0),
                             dac2=dict(gain=0),
                             dac3=dict(gain=0))
