# Copyright 2023 Qilimanjaro Quantum Tech
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""QubitControl class."""
from abc import abstractmethod
from dataclasses import asdict, dataclass
from typing import Sequence

from qililab.constants import RUNCARD
from qililab.instruments.awg_settings.awg_sequencer import AWGSequencer
from qililab.instruments.awg_settings.typings import AWGTypes
from qililab.instruments.instrument import Instrument
from qililab.pulse import PulseBusSchedule
from qililab.utils.asdict_factory import dict_factory


class AWG(Instrument):
    """Abstract base class defining all instruments used to control or readout the qubits."""

    @dataclass(kw_only=True)
    class AWGSettings(Instrument.InstrumentSettings):
        """Contains the settings of a AWG.

        Args:
            num_sequencers (int): Number of sequencers (physical I/Q pairs)
            awg_sequencers (Sequence[AWGSequencer]): Properties of each AWG sequencer
        """

        num_sequencers: int
        awg_sequencers: Sequence[AWGSequencer]

        def __post_init__(self):
            """build AWGSequencers and IQ channels"""
            super().__post_init__()
            if self.num_sequencers <= 0:
                raise ValueError(f"The number of sequencers must be greater than 0. Received: {self.num_sequencers}")
            if len(self.awg_sequencers) != self.num_sequencers:
                raise ValueError(
                    f"The number of sequencers: {self.num_sequencers} does not match"
                    + f" the number of AWG Sequencers settings specified: {len(self.awg_sequencers)}"
                )
            self.awg_sequencers = [
                AWGSequencer(**sequencer) if isinstance(sequencer, dict) else sequencer  # pylint: disable=not-a-mapping
                for sequencer in self.awg_sequencers
            ]

        def to_dict(self):
            """Return a dict representation of an AWG instrument."""
            result = asdict(self, dict_factory=dict_factory)
            result.pop(AWGTypes.AWG_SEQUENCERS.value)

            return result | {AWGTypes.AWG_SEQUENCERS.value: [sequencer.to_dict() for sequencer in self.awg_sequencers]}

    settings: AWGSettings

    @abstractmethod
    def compile(
        self, pulse_bus_schedule: PulseBusSchedule, nshots: int, repetition_duration: int, num_bins: int
    ) -> list:
        """Compiles the ``PulseBusSchedule`` into an assembly program.

        Args:
            pulse_bus_schedule (PulseBusSchedule): the list of pulses to be converted into a program
            nshots (int): number of shots / hardware average
            repetition_duration (int): repetition duration
            num_bins (int): number of bins

        Returns:
            list: list of compiled assembly programs
        """

    @abstractmethod
    def run(self, port: str):
        """Run the uploaded program"""

    @abstractmethod
    def upload(self, port: str):
        """Upload compiled program."""

    @property
    def num_sequencers(self):
        """Number of sequencers in the AWG

        Returns:
            int: number of sequencers
        """
        return self.settings.num_sequencers

    @property
    def awg_sequencers(self):
        """AWG 'awg_sequencers' property."""
        return self.settings.awg_sequencers

    @property
    def intermediate_frequencies(self):
        """AWG 'intermediate_frequencies' property."""
        return [sequencer.intermediate_frequency for sequencer in self.awg_sequencers]

    def to_dict(self):
        """Return a dict representation of an AWG instrument."""
        return {RUNCARD.NAME: self.name.value} | self.settings.to_dict()

    def get_sequencers_from_chip_port_id(self, chip_port_id: str):
        """Get sequencer ids from the chip port identifier

        Args:
            chip_port_id (str): chip port identifier

        Returns:
            list[AWGSequencer]: list of integers containing the indices of the sequencers connected to the chip port
        """
        if seqs := [sequencer for sequencer in self.awg_sequencers if sequencer.chip_port_id == chip_port_id]:
            return seqs
        raise IndexError(
            f"No sequencer found connected to port {chip_port_id}. Please make sure the `chip_port_id` "
            "attribute is correct."
        )

    def get_sequencer(self, sequencer_id: int) -> AWGSequencer:
        """Get sequencer from the sequencer identifier

        Args:
            sequencer_id (int): sequencer identifier

        Returns:
            AWGSequencer: sequencer associated with the sequencer_id
        """
        sequencer_identifiers = [
            sequencer.identifier for sequencer in self.awg_sequencers if sequencer.identifier == sequencer_id
        ]

        if len(sequencer_identifiers) != 1:
            raise ValueError(
                f"Each sequencer should have a unique id. Found {len(sequencer_identifiers)} sequencers "
                f"with id {sequencer_id}."
            )

        return self.awg_sequencers[sequencer_identifiers[0]]
