from pymoku.moku import Moku
from pymoku.sources import ADC, DAC


class MokuLab(Moku):
    DEV_NAME = 'moku20'
    AXI_BUS_WIDTH = 8  # bytes
    ADC_BIT_DEPTH = 12
    CBUF_RAM_LENGTH = 0x8000000
    URAM_LENGTH = 0

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

    def setup_routing(self):
        self.adc0 = ADC('InA', lambda: self.adc_coeffs(0))
        self.adc1 = ADC('InB', lambda: self.adc_coeffs(1))
        self.ext = ADC('Ext', lambda: 2.0**15)

        self.dac0 = DAC('OutA', moku=self, dac_idx=0, coeff=1.0)
        self.dac1 = DAC('OutB', moku=self, dac_idx=1, coeff=1.0)

        self._sources = [self.adc0, self.adc1, self.ext]
        self._outputs = [self.dac0, self.dac1]
        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.ext, idx=2)

            for idx, loop in enumerate(self._loop_muxes[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 _update_coeffs(self, hwstate):
        self._cal = dict(
            adc=[hwstate['adc0']['_coeff'],
                 hwstate['adc1']['_coeff']],
            dac=[hwstate['dac0']['_coeff'],
                 hwstate['dac1']['_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, atten=None, dc=None):
        p = {}
        r = {}

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

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

        if atten is not None:
            assert atten in [0, 20]
            r['atten'] = atten

        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, fiftyr=True, atten=0, dc=True)
        self.set_frontend(1, fiftyr=True, atten=0, dc=True)
