import pymoku
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_props([self.reg_unsg(reg=0, offset=0, length=4),
                                     self.reg_unsg(reg=0, offset=4, length=4),
                                     self.reg_unsg(reg=0, offset=8, length=4),
                                     self.reg_unsg(reg=0, offset=12, length=4)])


class _Routing(RegisterAccess):
    def __init__(self, parent, offset=0, prefix=None):
        super().__init__(parent, offset, prefix)
        self.slots = [_SlotControl(self, offset=6 + i, slot_id=i)
                      for i in range(parent.num_slots)]
        self.dacs = self.reg_props([self.reg_unsg(reg=4, offset=0, length=4),
                                    self.reg_unsg(reg=4, offset=4, length=4),
                                    self.reg_unsg(reg=4, offset=8, length=4),
                                    self.reg_unsg(reg=4, offset=12, length=4)])


class MokuGo(Moku):
    DEV_NAME = 'mokugo'
    AXI_BUS_WIDTH = 8  # bytes
    ADC_BIT_DEPTH = 12
    DAC_BIT_DEPTH = 12
    CBUF_RAM_LENGTH = 0x8000000
    URAM_LENGTH = 0

    class DIO(DAC):
        def __init__(self, name, moku, dio_key):
            Node.__init__(self, name)
            self.moku = moku
            self.dio_key = dio_key
            self.bus_width = 1

        def set_idx(self, idx):
            self.moku.dio_output(**{self.dio_key: idx})

        def get_idx(self):
            return self.moku.dio_output()[self.dio_key]

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)

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

        self.adc0 = ADC('In 1', lambda: self.adc_coeffs(0))
        self.adc1 = ADC('In 2', lambda: self.adc_coeffs(1))
        self.dio_in = ADC('DIO In', lambda: 1.0, dtype=pymoku.dtype.digital)

        self.dac0 = DAC('Out 1', self.routing, dac_idx=0, coeff=1.0)
        self.dac1 = DAC('Out 2', self.routing, dac_idx=1, coeff=1.0)
        self.dio_out_a = MokuGo.DIO('DIO OutA', moku=self, dio_key='source0')
        self.dio_dir_a = MokuGo.DIO('DIO DirA', moku=self, dio_key='dir_source0')
        self.dio_out_b = MokuGo.DIO('DIO OutB', moku=self, dio_key='source1')
        self.dio_dir_b = MokuGo.DIO('DIO DirB', moku=self, dio_key='dir_source1')

        self._sources = [self.adc0, self.adc1, self.dio_in]
        self._outputs = [self.dac0, self.dac1, self.dio_out_a, self.dio_dir_a, self.dio_out_b, self.dio_dir_b]

        super().setup_routing()

    def attach_slot(self, slot_inst, slot):
        for inp in slot_inst.inputs():
            inp.add(self.adc0, idx=0)
            inp.add(self.adc1, idx=1)
            inp.add(self.dio_in, idx=2)

            for idx, loop in enumerate(self._loop_muxes_1[slot]):
                inp.add(loop, idx=idx + 3)

        for idx, outp in enumerate(slot_inst.outputs()):
            slot_outp = self._slot_outputs[slot][idx]
            slot_outp.clear()
            slot_outp.add(outp)

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

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

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

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

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

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

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

        self.modify_hardware(**p)

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

    def set_power_supply(self, ch, V=None, I=None):
        ppsu = 'ppsu{:d}'.format(ch)
        req = {ppsu: dict(enable=True, setpoint=dict(I=I, V=V))}
        self.modify_hardware(**req)

    def dio_output(self, source0=None, source1=None,
                   dir_source0=None, dir_source1=None,
                   ctrl=None, dir_ctrl=None):
        dio = {}
        if source0 is not None:
            dio['source0'] = source0

        if source1 is not None:
            dio['source1'] = source1

        if ctrl is not None:
            dio['ctrl'] = ctrl

        if dir_source0 is not None:
            dio['dir_source0'] = dir_source0

        if dir_source1 is not None:
            dio['dir_source1'] = dir_source1

        if dir_ctrl is not None:
            dio['dir_ctrl'] = dir_ctrl

        data = {'dio': dio} if dio else {}
        state = self.modify_hardware(**data)
        return state['dio']
