"""Gain stage controller for the aodsv2_eval board.
"""
from copy import deepcopy

from naludaq.communication import LOGGER
from naludaq.communication.analog_registers import AnalogRegisters
from naludaq.communication.control_registers import ControlRegisters

from .aodsv2 import GainStageController as AODSV2GainStageController


class GainStageController(AODSV2GainStageController):
    """Controls the AODS gain stages

    Gain stages are controlled by the following registers:
    - sel_0: ch0 input buffer
    - sel_1: ch1 input buffer
    - sel_2: ch2 input buffer
    - sel_3: ch3 input buffer
    - sel_4: external output
    - sel_5: ch1 BAMP or x8 of ch0 output
    - sel_6: select either BAMP or 8x gain for ch3
    - sel_7: select ch3 input either ch0 output or ch2 output

    There are several functions in the form `ch{X}_{some possible configuration}`
    such as `ch3_external_input()` that serve as presets and a human-readable
    way to configure each stage. Note that the configurations for each channel
    are mutually exclusive; e.g. you cannot have ch1 external input and ch1 8x ch0
    simulatenously.
    """

    def __init__(self, board, chip_number=0):
        """Initialize the controller"""
        super().__init__(board)
        self.chip_num = chip_number

        # Since we don't support multiple chips in the registers yet, we need to
        # keep track of the values of the registers ourselves.
        self._sel_values = deepcopy(
            board.registers["analog_registers"]
        )  # TODO: fix this when multichip registers are fixed

    @property
    def sel_values(self) -> "list[bool]":
        """Get a list of sel values from the internally stored registers"""
        return [
            self._sel_values[f"sel_{i}"][
                "value"
            ]  # TODO: fix this when multichip registers are fixed
            for i in range(self._num_stages)
        ]

    def _write_reg(self, reg, value):
        """Write a register.

        We need to override the chip mask manually since we don't support
        multichip registers yet.
        """
        LOGGER.debug("Selecting chip %d", self.chip_num)
        self._set_chip_tx_enabled(
            [self.chip_num]
        )  # TODO: fix this when multichip registers are fixed
        AnalogRegisters(self.board).write(reg, value)
        self._sel_values[reg][
            "value"
        ] = value  # TODO: fix this when multichip registers are fixed
        self._set_chip_tx_enabled(
            range(self.board.params.get("num_chips", 2))
        )  # TODO: fix this when multichip registers are fixed

    def _set_chip_tx_enabled(
        self, chips: list = [0, 1]
    ):  # TODO: fix this when multichip registers are fixed
        """Set whether the TX lines are enabled for each chip.
        Can be used to write a register on a single chip instead of both.
        """
        ControlRegisters(self.board).write("ard_tx_en", sum([1 << x for x in chips]))
