from functools import partial
from pymoku.applets import doasync
from PySide6.QtWidgets import (QWidget, QVBoxLayout, QHBoxLayout,
                               QPushButton, QCheckBox, QGroupBox)
from PySide6.QtCore import QTimer, Qt, Signal
from pyqtgraph.parametertree import Parameter, ParameterTree


class _StatusCheck(QCheckBox):
    def __init__(self, parent=None, label=None):
        super().__init__(parent)
        self.setStyleSheet('QCheckBox::indicator:checked { background-color: green; } '
                           'QCheckBox::indicator:unchecked { background-color: red; }')
        self.setTristate(True)
        self.setCheckState(Qt.PartiallyChecked)
        self.setEnabled(False)
        self.setToolTip(label)


class RfdcDialog(QWidget):
    _state_update = Signal(dict)

    def __init__(self, moku):
        super().__init__()
        self.setWindowTitle("RFDC Config")
        self.moku = moku

        tree = ParameterTree(self, showHeader=False)
        # tree.addParameters(self._adc_mixer_params())
        # tree.addParameters(self._dac_mixer_params())
        tree.addParameters(self._adc_qmc_params())
        tree.addParameters(self._dac_qmc_params())
        tree.addParameters(self._threshold_params())
        tree.addParameters(self._dsa_params())

        llyt = QHBoxLayout()
        llyt.addWidget(self._adc_status())
        llyt.addWidget(self._dac_status())
        llyt.addWidget(self._rts_panel())

        lyt = QVBoxLayout()
        lyt.addWidget(tree)
        lyt.addLayout(llyt)
        self.setLayout(lyt)

        self._timer = QTimer(self)
        self._timer.timeout.connect(self._poll_state)
        self._poll_state()

    def _adc_mixer_params(self):
        params = Parameter.create(name='ADC Mixers', type='group')

        def _update_mixer(mixer, idx, state):
            for k, v in state['adc_mix'][idx].items():
                mixer.child(k).setValue(v, blockSignal=True)

        def _set_mixer(idx, mixer, changes):
            # d = {c.name(): c.value() for c in mixer.children()}
            # TODO mercury API
            pass

        for idx in range(8):
            mixer = Parameter.create(
                name=f'ADC{idx:d}', type='group', expanded=False,
                children=[
                    dict(name='Freq', type='float', value=0.0),
                    dict(name='PhaseOffset', type='float', value=0.0),
                    dict(name='EventSource', type='int', value=2),
                    dict(name='CoarseMixFreq', type='int', value=16),
                    dict(name='MixerMode', type='int', value=4),
                    dict(name='FineMixerScale', type='int', value=0),
                    dict(name='MixerType', type='int', value=1)
                ])
            params.addChild(mixer)
            mixer.sigTreeStateChanged.connect(partial(_set_mixer, idx))
            self._state_update.connect(partial(_update_mixer, mixer, idx))

        return params

    def _dac_mixer_params(self):
        params = Parameter.create(name='DAC Mixers', type='group')

        def _update_mixer(mixer, idx, state):
            for k, v in state['dac_mix'][idx].items():
                mixer.child(k).setValue(v, blockSignal=True)

        def _set_mixer(idx, mixer, changes):
            # d = {c.name(): c.value() for c in mixer.children()}
            # TODO mercury API
            pass

        for idx in range(8):
            mixer = Parameter.create(
                name=f'DAC{idx:d}', type='group', expanded=False,
                children=[
                    dict(name='Freq', type='float', value=0.0),
                    dict(name='PhaseOffset', type='float', value=0.0),
                    dict(name='EventSource', type='int', value=2),
                    dict(name='CoarseMixFreq', type='int', value=16),
                    dict(name='MixerMode', type='int', value=4),
                    dict(name='FineMixerScale', type='int', value=0),
                    dict(name='MixerType', type='int', value=1)
                ])
            params.addChild(mixer)
            mixer.sigTreeStateChanged.connect(partial(_set_mixer, idx))
            self._state_update.connect(partial(_update_mixer, mixer, idx))

        return params

    def _adc_qmc_params(self):
        params = Parameter.create(name='ADC QMC', type='group')

        def _update_qmc(qmc, idx, state):
            for k, v in state['adc_qmc'][idx].items():
                qmc.child(k).setValue(v, blockSignal=True)

        def _set_qmc(idx, qmc, changes):
            # d = {c.name(): c.value() for c in qmc.children()}
            # TODO mercury API
            pass

        for idx in range(8):
            qmc = Parameter.create(
                name=f'ADC{idx:d}', type='group', expanded=False,
                children=[
                    dict(name='EnablePhase', type='bool', value=False),
                    dict(name='EnableGain', type='bool', value=False),
                    dict(name='GainCorrectionFactor', type='float', value=0.0),
                    dict(name='PhaseCorrectionFactor', type='float', value=0.0),
                    dict(name='OffsetCorrectionFactor', type='int', value=0),
                    dict(name='EventSource', type='int', value=0),
                ])
            params.addChild(qmc)
            qmc.sigTreeStateChanged.connect(partial(_set_qmc, idx))
            self._state_update.connect(partial(_update_qmc, qmc, idx))

        return params

    def _dac_qmc_params(self):
        params = Parameter.create(name='DAC QMC', type='group')

        def _update_qmc(qmc, idx, state):
            for k, v in state['dac_qmc'][idx].items():
                qmc.child(k).setValue(v, blockSignal=True)

        def _set_qmc(idx, qmc, changes):
            # d = {c.name(): c.value() for c in qmc.children()}
            # TODO mercury API
            pass

        for idx in range(8):
            qmc = Parameter.create(
                name=f'DAC{idx:d}', type='group', expanded=False,
                children=[
                    dict(name='EnablePhase', type='bool', value=False),
                    dict(name='EnableGain', type='bool', value=False),
                    dict(name='GainCorrectionFactor', type='float', value=0.0),
                    dict(name='PhaseCorrectionFactor', type='float', value=0.0),
                    dict(name='OffsetCorrectionFactor', type='int', value=0),
                    dict(name='EventSource', type='int', value=0),
                ])
            params.addChild(qmc)
            qmc.sigTreeStateChanged.connect(partial(_set_qmc, idx))
            self._state_update.connect(partial(_update_qmc, qmc, idx))

        return params

    def _threshold_params(self):
        params = Parameter.create(name='Thresholds', type='group')
        {
            'UpdateThreshold': 4,
            'ThresholdMode': [
                {'real': 0, 'imag': 0, 'numerator': 0, 'denominator': 1},
                {'real': 0, 'imag': 0, 'numerator': 0, 'denominator': 1}
            ],
            'ThresholdAvgVal': [
                {'real': 0, 'imag': 0, 'numerator': 0, 'denominator': 1},
                {'real': 0, 'imag': 0, 'numerator': 0, 'denominator': 1}
            ],
            'ThresholdUnderVal': [
                {'real': 0, 'imag': 0, 'numerator': 0, 'denominator': 1},
                {'real': 0, 'imag': 0, 'numerator': 0, 'denominator': 1}
            ],
            'ThresholdOverVal': [
                {'real': 0, 'imag': 0, 'numerator': 0, 'denominator': 1},
                {'real': 0, 'imag': 0, 'numerator': 0, 'denominator': 1}
            ]
        },
        return params

    def _dsa_params(self):
        params = Parameter.create(name='DSA', type='group')
        # 'dsa': [
        #     {'DisableRTS': 0, 'Attenuation': 0.0},
        #     {'DisableRTS': 0, 'Attenuation': 0.0},
        #     {'DisableRTS': 0, 'Attenuation': 0.0},
        #     {'DisableRTS': 0, 'Attenuation': 0.0},
        #     {'DisableRTS': 0, 'Attenuation': 0.0},
        #     {'DisableRTS': 0, 'Attenuation': 0.0},
        #     {'DisableRTS': 0, 'Attenuation': 0.0},
        #     {'DisableRTS': 0, 'Attenuation': 0.0}
        # ],
        return params

    def _rts_panel(self, state=None):
        group = QGroupBox('Over Voltage')
        lyt = QHBoxLayout()

        def _update(wdgt, mask, state):
            ov = bool(state.get('rts', 0) & mask)
            wdgt.setChecked(not ov)

        labels = [(f'ADC{idx:d}', 1 << (4 * idx)) for idx in range(8)]
        for label, mask in labels:
            wdgt = _StatusCheck(self, label)
            lyt.addWidget(wdgt)
            self._state_update.connect(partial(_update, wdgt, mask))

        def _reset():
            self._poll_state(dict(rts=0x11111111))

        btn = QPushButton("Reset All")
        btn.clicked.connect(_reset)
        lyt.addWidget(btn)

        group.setLayout(lyt)
        return group

    def _adc_status(self, state=None):
        group = QGroupBox('ADC Tiles')
        lyt = QHBoxLayout()

        def _update(wdgt, idx, state):
            enabled = bool(state['ip_status']['ADCTileStatus'][idx]['IsEnabled'])
            wdgt.setChecked(enabled)

        for idx in range(4):
            label = f'ADC Tile {idx:d}'
            wdgt = _StatusCheck(self, label)
            lyt.addWidget(wdgt)
            self._state_update.connect(partial(_update, wdgt, idx))

        group.setLayout(lyt)
        return group

    def _dac_status(self, state=None):
        group = QGroupBox('DAC Tiles')
        lyt = QHBoxLayout()

        def _update(wdgt, idx, state):
            enabled = bool(state['ip_status']['DACTileStatus'][idx]['IsEnabled'])
            wdgt.setChecked(enabled)

        for idx in range(4):
            label = f'DAC Tile {idx:d}'
            wdgt = _StatusCheck(self, label)
            lyt.addWidget(wdgt)
            self._state_update.connect(partial(_update, wdgt, idx))

        group.setLayout(lyt)
        return group

    @doasync
    def _poll_state(self, msg={}):
        state = self.moku.modify_hardware(rfdc=msg)['rfdc']
        self._state_update.emit(state)

    def closeEvent(self, e):
        self._timer.stop()
        super().closeEvent(e)

    def showEvent(self, e):
        self._timer.start(1000)
        super().showEvent(e)


def main():
    from PySide6.QtWidgets import QApplication

    app = QApplication()

    f = RfdcDialog(None)
    f.show()
    app.exec()


if __name__ == '__main__':
    main()
