#  Copyright (C) 2023 ASTRON (Netherlands Institute for Radio Astronomy)
#  SPDX-License-Identifier: Apache-2.0

import numpy

from integration_test.default.devices.base import AbstractTestBases
from tangostationcontrol.common.constants import (
    N_rcu,
    N_rcu_inp,
    DEFAULT_N_HBA_TILES,
    CLK_200_MHZ,
)


class TestCalibrationDevice(AbstractTestBases.TestDeviceBase):
    """Integration test class for device calibration"""

    HBA_ANTENNA_NAMES = [
        "Tile00",
        "Tile01",
        "Tile02",
        "Tile03",
        "Tile04",
        "Tile05",
        "Tile06",
        "Tile07",
        "Tile08",
        "Tile09",
        "Tile10",
        "Tile11",
        "Tile12",
        "Tile13",
        "Tile14",
        "Tile15",
        "Tile16",
        "Tile17",
        "Tile18",
        "Tile19",
        "Tile20",
        "Tile21",
        "Tile22",
        "Tile23",
        "Tile24",
        "Tile25",
        "Tile26",
        "Tile27",
        "Tile28",
        "Tile29",
        "Tile30",
        "Tile31",
        "Tile32",
        "Tile33",
        "Tile34",
        "Tile35",
        "Tile36",
        "Tile37",
        "Tile38",
        "Tile39",
        "Tile40",
        "Tile41",
        "Tile42",
        "Tile43",
        "Tile44",
        "Tile45",
        "Tile46",
        "Tile47",
    ]
    antennafield_iden = "STAT/AFH/HBA0"

    def setUp(self):
        self.stationmanager_proxy = self.setup_proxy("STAT/StationManager/1")

        super().setUp("STAT/Calibration/1")

        self.recv_proxy = self.setup_proxy("STAT/RECVH/H0")
        self.sdpfirmware_proxy = self.setup_proxy("STAT/SDPFirmware/HBA0")
        self.sdp_proxy = self.setup_proxy("STAT/SDP/HBA0")

        # antennafield properties are restored for each test
        self.antennafield_proxy = self.setup_proxy(
            self.antennafield_iden,
            restore_properties=True,
            cb=lambda x: {
                x.put_property(
                    {
                        "Power_to_RECV_mapping": [1, 1, 1, 0]
                        + [-1] * ((DEFAULT_N_HBA_TILES * 2) - 4),
                        "Antenna_Sets": ["ALL"],
                        "Antenna_Set_Masks": ["1" * DEFAULT_N_HBA_TILES],
                        "Frequency_Band_RW_default": ["HBA_110_190"]
                        * (DEFAULT_N_HBA_TILES * 2),
                    }
                )
            },
        )

        # configure the frequencies, which allows access
        # to the calibration attributes and commands
        self.sdpfirmware_proxy.clock_RW = CLK_200_MHZ
        self.recv_proxy.RCU_band_select_RW = [[1] * N_rcu_inp] * N_rcu

    def test_calibrate_recv(self):
        """Test whether RECV calibration works as expected"""
        calibration_properties = {
            "Antenna_Cables": ["50m", "80m"] * (DEFAULT_N_HBA_TILES // 2),
            "Control_to_RECV_mapping":
            # [1, 0, 1, 1, 1, 2, 1, x ... 1, 47]
            numpy.array([[1, x] for x in range(0, DEFAULT_N_HBA_TILES)]).flatten(),
            "Power_to_RECV_mapping": numpy.array(
                [[1, x + DEFAULT_N_HBA_TILES] for x in range(0, DEFAULT_N_HBA_TILES)]
            ).flatten(),
            # [1, 48, 1, 49, x ... 1, 95]
            "Frequency_Band_RW_default": ["HBA_110_190"] * (DEFAULT_N_HBA_TILES * 2),
        }

        self.antennafield_proxy = self.setup_proxy(
            self.antennafield_iden,
            cb=lambda x: {x.put_property(calibration_properties)},
        )

        self.proxy.boot()

        # calibrate
        self.proxy.calibrate_recv("STAT/AFH/HBA0")

        # check the results
        rcu_attenuator_db_pwr = self.antennafield_proxy.RCU_attenuator_dB_RW[:, 0]
        rcu_attenuator_db_ctrl = self.antennafield_proxy.RCU_attenuator_dB_RW[:, 1]

        # gather settings
        field_attenuation = self.antennafield_proxy.Field_Attenuation_R

        for mapping_name, rcu_attenuator_db in [
            ("power", rcu_attenuator_db_pwr),
            ("control", rcu_attenuator_db_ctrl),
        ]:
            # values should be the same for the same cable length
            self.assertEqual(
                1,
                len(set(rcu_attenuator_db[0::2])),
                msg=f"{mapping_name} - rcu_attenuator_db={rcu_attenuator_db}",
            )
            self.assertEqual(
                1,
                len(set(rcu_attenuator_db[1::2])),
                msg=f"{mapping_name} - rcu_attenuator_db={rcu_attenuator_db}",
            )
            # value should be larger for the shorter cable, as those signals need damping
            self.assertGreater(
                rcu_attenuator_db[0],
                rcu_attenuator_db[1],
                msg=f"{mapping_name} - rcu_attenuator_db={rcu_attenuator_db}",
            )
            # longest cable should require no damping, so only field attenuation applies
            self.assertEqual(
                field_attenuation,
                rcu_attenuator_db[1],
                msg=f"{mapping_name} - rcu_attenuator_db={rcu_attenuator_db}",
            )

    def test_calibrate_sdp(self):
        """Test whether SDP calibration works as expected"""
        calibration_properties = {
            "Antenna_Names": self.HBA_ANTENNA_NAMES,
            "Antenna_Cables": ["50m", "80m"] * (DEFAULT_N_HBA_TILES // 2),
            "Antenna_to_SDP_Mapping": [0, 1, 0, 0]
            + [-1, -1] * (DEFAULT_N_HBA_TILES - 2),
            "Control_to_RECV_mapping":
            # [1, 0, 1, 1, 1, 2, 1, x ... 1, 47]
            numpy.array([[1, x] for x in range(0, DEFAULT_N_HBA_TILES)]).flatten(),
            "Power_to_RECV_mapping": numpy.array(
                [[1, x + DEFAULT_N_HBA_TILES] for x in range(0, DEFAULT_N_HBA_TILES)]
            ).flatten(),
            # [1, 48, 1, 49, x ... 1, 95]
            "Frequency_Band_RW_default": ["HBA_110_190"] * (DEFAULT_N_HBA_TILES * 2),
        }

        self.antennafield_proxy = self.setup_proxy(
            self.antennafield_iden,
            cb=lambda x: {x.put_property(calibration_properties)},
        )

        self.proxy.boot()

        # calibrate
        self.proxy.calibrate_sdp("STAT/AFH/HBA0")

        # check the results
        # antenna #0 is on FPGA 0, input 2 and 3,
        # antenna #1 is on FPGA 0, input 0 and 1
        signal_input_samples_delay = self.sdp_proxy.FPGA_signal_input_samples_delay_RW

        # delays should be equal for both polarisations
        self.assertEqual(
            signal_input_samples_delay[0, 0], signal_input_samples_delay[0, 1]
        )
        self.assertEqual(
            signal_input_samples_delay[0, 2], signal_input_samples_delay[0, 3]
        )

        # antenna #0 is shorter, so should have a greater delay
        self.assertGreater(
            signal_input_samples_delay[0, 2],
            signal_input_samples_delay[0, 0],
            msg=f"{signal_input_samples_delay}",
        )
        # antenna #1 is longest, so should have delay 0
        self.assertEqual(0, signal_input_samples_delay[0, 0])