from itertools import chain
import pymoku
from pymoku import batched
from pymoku.moku import Moku
from pymoku.sources import ADC, DAC, Node, SlotOutput
from pymoku.registers import RegisterAccess


class _DAC(DAC):
    def __init__(self, name, routing, dac_idx, coeff):
        super().__init__(name, routing, dac_idx, coeff, bus_width=4)

    def set_idx(self, idx):
        if idx is None:
            idx = (0xFF,)
        self.routing.dacs[self.dac_idx] = idx[0]

    def get_idx(self):
        return (self.routing.dacs[self.dac_idx],)


class _DIO(DAC):
    def __init__(self, name, routing, dio_idx, coeff):
        super().__init__(name, routing, dio_idx, coeff)

    def set_idx(self, idx):
        if idx is None:
            idx = (0xFF,)
        self.routing.dacs[self.dac_idx] = idx[0]

    def get_idx(self):
        return (self.routing.dacs[self.dac_idx],)


class _SFP(DAC):
    def set_idx(self, idx):
        if idx is None:
            idx = (0b11111,)
        self.routing.sfp[self.dac_idx:self.dac_idx + self.bus_width] = idx

    def get_idx(self):
        return tuple(self.routing.sfp[self.dac_idx:self.dac_idx + self.bus_width])


class _SlotInputMux(RegisterAccess, Node):
    def __init__(self, parent, regs):
        RegisterAccess.__init__(self, parent, 0)
        Node.__init__(self)
        self.inp_regs = self.reg_props(regs)
        self.reg_length = regs[0].length

    def set_idx(self, idx):
        if idx is None:
            idx = [2**self.reg_length - 1] * self.bus_width
        self.inp_regs[:] = idx

    def get_idx(self):
        if self.inp_regs[0] == 2**self.reg_length - 1:
            return None
        return tuple(self.inp_regs[:])


class _SlotControl(RegisterAccess):
    LOOP_OFF = 17

    def __init__(self, parent=None, offset=0, prefix=None, slot_id=None):
        super().__init__(parent, offset, prefix)
        self.slot_id = slot_id
        self.moku = self.reg_control()
        regs = [
            self.reg_slv(reg=0, offset=0, length=8),
            self.reg_slv(reg=0, offset=8, length=8),
            self.reg_slv(reg=0, offset=16, length=8),
            self.reg_slv(reg=0, offset=24, length=8),
            self.reg_slv(reg=1, offset=0, length=8),
            self.reg_slv(reg=1, offset=8, length=8),
            self.reg_slv(reg=1, offset=16, length=8),
            self.reg_slv(reg=1, offset=24, length=8),
            self.reg_slv(reg=2, offset=0, length=8),
            self.reg_slv(reg=2, offset=8, length=8),
            self.reg_slv(reg=2, offset=16, length=8),
            self.reg_slv(reg=2, offset=24, length=8),
        ][:4 if self.moku.num_slots == 8 else 8]  # TODO get slot_ains correcly in manifest

        self.mux_1s = [self._slot_mux_1(r) for r in regs]
        self.mux_4s = [self._slot_mux_4(rs) for rs in batched(regs, 4) if len(rs) == 4]

    def _slot_mux_1(self, reg):
        # overlapping regs
        mux = _SlotInputMux(self, [reg])

        for idx, adc in enumerate(self.moku.adcs):
            mux.add(adc, idx=idx)
        mux.add(self.moku.ext, idx=8)
        mux.add(self.moku.dio0, idx=9)
        mux.add(self.moku.dio1, idx=10)
        # mux.add(self.moku.sfp[2], idx=(11, 12, 13, 14))  # TODO we can use the QSFP as 4 inputs
        mux.add(self.moku.sfp_in[0], idx=15)
        mux.add(self.moku.sfp_in[1], idx=16)

        for idx, outp in enumerate(list(chain(*self.moku._slot_outputs))):
            mux.add(node=outp, idx=idx + self.LOOP_OFF)

        return mux

    def _slot_mux_4(self, regs):
        def q(a):
            b = 0x80 | (a << 2)
            return (a, b, b, b)

        mux = _SlotInputMux(self, regs)

        for idx, adc in enumerate(self.moku.adcs):
            mux.add(adc, idx=q(idx), bus_width=4)
        mux.add(self.moku.ext, idx=q(32), bus_width=4)
        mux.add(self.moku.sfp_in[2], idx=(11, 12, 13, 14), bus_width=4)

        for idx, outp in enumerate(list(chain(*self.moku._slot_outputs))[::4]):
            edge_idx = tuple(range(idx * 4 + self.LOOP_OFF, (idx + 1) * 4 + self.LOOP_OFF, 1))
            mux.add(node=outp, idx=edge_idx, bus_width=4)

        return mux


# 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 _MultiSlot(RegisterAccess):
    def __init__(self, parent=None, offset=0, prefix=None):
        super().__init__(parent, offset, prefix)
        self.slots = [_SlotControl(self, offset=2 + 3 * i, slot_id=i)
                      for i in range(parent.num_slots)]
        self.dacs = self.reg_prop_list(reg=0, offset=0, length=8,
                                       count=8, mkreg=self.reg_slv)
        # self.dio = self.reg_prop_list(reg=0, offset=0, length=8,
        #                               count=8, mkreg=self.reg_slv)
        self.sfp = self.reg_prop_list(reg=26, offset=0, length=5,
                                      count=6, mkreg=self.reg_slv)


class MokuDelta(Moku):
    DEV_NAME = 'mokudelta'
    AXI_BUS_WIDTH = 16
    ADC_BIT_DEPTH = 14
    CBUF_RAM_LENGTH = 0x180000000
    URAM_LENGTH = 0x200000 // 32

    def setup_routing(self):
        # coeff: 50ohm, 0dB 30MHz
        self.adcs = [ADC(f'In {i + 1}', coeff=lambda: 38010.88, bus_width=4) for i in range(8)]
        self.ladcs = [ADC(f'LADC {i + 1}', coeff=lambda: 2**16) for i in range(8)]
        self.dio0 = ADC('DIO 1', coeff=lambda: 2**16, dtype=pymoku.dtype.digital)
        self.dio1 = ADC('DIO 2', coeff=lambda: 2**16, dtype=pymoku.dtype.digital)
        self.ext = ADC('Trigger', coeff=lambda: 2**16, bus_width=4)
        self.sfp_in = [ADC('SFP0', coeff=lambda: 2**16),
                       ADC('SFP1', coeff=lambda: 2**16),
                       ADC('QSFP', coeff=lambda: 2**16, bus_width=4)]

        self._slot_outputs = [[SlotOutput(slot_idx, outp_idx)
                              for outp_idx in range(self.num_slot_outputs)]
                              for slot_idx in range(self.num_slots)]

        self.routing = _MultiSlot(self, offset=64)
        self.dacs = [_DAC(f'Out {i + 1}', self.routing, dac_idx=i, coeff=lambda: 21845.33) for i in range(8)]
        self.dio_out0 = _DIO('DIO 1', self.routing, 0, coeff=lambda: 2**16)
        self.dio_out1 = _DIO('DIO 2', self.routing, 1, coeff=lambda: 2**16)
        self.sfp_out = [_SFP('SFP0', self.routing, 4),
                        _SFP('SFP1', self.routing, 5),
                        _SFP('QSFP', self.routing, 0, bus_width=4)]

        self._sources = self.adcs + [self.ext, self.dio0, self.dio1] + self.sfp_in  # + self.ladcs
        self._outputs = self.dacs + [self.dio_out0, self.dio_out1] + self.sfp_out

        self._loop_muxes_1 = self._create_loop_muxes(width=1)
        self._loop_muxes_2 = self._create_loop_muxes(width=2)
        self._loop_muxes_4 = self._create_loop_muxes(width=4)

        for dac in self.dacs:
            for idx, outp in enumerate(list(chain(*self._slot_outputs))):
                dac.add(outp, idx=idx)

            # TODO width 2
            for idx, outp in enumerate(list(chain(*self._slot_outputs))[::4]):
                dac.add(outp, idx=0x80 | idx, bus_width=4)

        for sfp in self.sfp_out:
            for idx, outp in enumerate(list(chain(*self._slot_outputs))):
                sfp.add(node=outp, idx=(idx,))

            for idx, outp in enumerate(list(chain(*self._slot_outputs))[::4]):
                edge_idx = tuple(range(idx * 4, (idx + 1) * 4, 1))
                sfp.add(node=outp, idx=edge_idx, bus_width=4)

        for idx, outp in enumerate(list(chain(*self._slot_outputs))):
            self.dio_out0.add(outp, idx=idx)
            self.dio_out1.add(outp, idx=idx)

    def attach_slot(self, slot_inst, slot):
        for slot_outp in self._slot_outputs[slot]:
            slot_outp.clear()

        for idx, outp in enumerate(slot_inst.outputs()):
            slot_outp = self._slot_outputs[slot][idx * outp.bus_width]
            slot_outp.add(outp, bus_width=outp.bus_width)
            # TODO sub-phases

        if self.num_slots == 1:
            self._attach_slot_single(slot_inst, slot)
        else:
            self._attach_slot_multi(slot_inst, slot)

    def _attach_slot_single(self, slot_inst, slot):
        for inp in slot_inst.inputs():
            inp.add(self.adcs[0], idx=(0, 1, 2, 3), bus_width=4)
            inp.add(self.adcs[1], idx=(4, 5, 6, 7), bus_width=4)
            inp.add(self.adcs[2], idx=(8, 9, 10, 11), bus_width=4)
            inp.add(self.adcs[3], idx=(12, 13, 14, 15), bus_width=4)
            inp.add(self.adcs[4], idx=(16, 17, 18, 19), bus_width=4)
            inp.add(self.adcs[5], idx=(20, 21, 22, 23), bus_width=4)
            inp.add(self.adcs[6], idx=(24, 25, 26, 27), bus_width=4)
            inp.add(self.adcs[7], idx=(28, 29, 30, 31), bus_width=4)
            inp.add(self.ext, idx=(32, 33, 34, 35), bus_width=4)

            inp.add(self.adcs[0], idx=(0, 2), bus_width=2)
            inp.add(self.adcs[1], idx=(4, 6), bus_width=2)
            inp.add(self.adcs[2], idx=(8, 10), bus_width=2)
            inp.add(self.adcs[3], idx=(12, 14), bus_width=2)
            inp.add(self.adcs[4], idx=(16, 18), bus_width=2)
            inp.add(self.adcs[5], idx=(20, 22), bus_width=2)
            inp.add(self.adcs[6], idx=(24, 26), bus_width=2)
            inp.add(self.adcs[7], idx=(28, 30), bus_width=2)
            inp.add(self.ext, idx=(32, 34), bus_width=2)

            inp.add(self.adcs[0], idx=0)
            inp.add(self.adcs[1], idx=4)
            inp.add(self.adcs[2], idx=8)
            inp.add(self.adcs[3], idx=12)
            inp.add(self.adcs[4], idx=16)
            inp.add(self.adcs[5], idx=20)
            inp.add(self.adcs[6], idx=24)
            inp.add(self.adcs[7], idx=28)
            inp.add(self.ext, idx=32)

            inp.add(self.dio0, idx=36)
            inp.add(self.dio1, idx=37)

            inp.add(self.sfp_in[2], idx=(38, 39, 40, 41), bus_width=4)
            inp.add(self.sfp_in[0], idx=42)
            inp.add(self.sfp_in[1], idx=43)

            inp.attach(inp.get_attached())  # init node refs

    def _attach_slot_multi(self, slot_inst, slot):
        routes = self.routing.slots[slot]
        for inp in slot_inst.inputs():
            for idx, mux in enumerate(routes.mux_4s):
                inp.add(mux, idx=tuple(range(idx * 4, (idx + 1) * 4)), bus_width=4)

            # TODO 2X

            for idx, mux in enumerate(routes.mux_1s):
                inp.add(mux, idx=idx, bus_width=1)

            inp.attach(inp.get_attached())

    def adc_coeffs(self, ch):
        return 2**16

    def _update_coeffs(self, hwstate):
        pass

    def dac_coeffs(self, ch):
        return 2**16

    def set_frontend(self, ch, fiftyr=None, gain=None, dc=None, rf=None):
        p = {}
        r = {}

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

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

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

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

        p[f'adc{ch:d}'] = r

        self.modify_hardware(**p)

    def hardware_defaults(self):
        pass
