# 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.

import math
from collections import deque
from typing import TYPE_CHECKING, Any, Callable

from qililab.config import logger
from qililab.instruments.qdevil import QDevilQDac2
from qililab.qprogram.blocks import Average, Block, ForLoop, InfiniteLoop, Loop, Parallel
from qililab.qprogram.calibration import Calibration
from qililab.qprogram.operations import (
    Acquire,
    Measure,
    Operation,
    Play,
    ResetPhase,
    SetFrequency,
    SetGain,
    SetMarkers,
    SetOffset,
    SetPhase,
    SetTrigger,
    Sync,
    Wait,
    WaitTrigger,
)
from qililab.qprogram.qprogram import QProgram
from qililab.typings.enums import Parameter

if TYPE_CHECKING:
    from qililab.platform.components.bus import Bus


class QdacCompilationOutput:
    """Class representing the output information generated by QdacCompiler.

    Attributes:
        sequences (Sequence): A dictionary with the buses participating in the QProgram as keys and the corresponding Sequence as values.
        acquisitions (Acquisitions): A dictionary with the buses participating in the acquisitions as keys and the corresponding Acquisitions as values.
    """

    def __init__(self, qprogram: QProgram, qdac: QDevilQDac2, trigger_position: str | None):
        self.qprogram = qprogram
        self.qdac = qdac
        self.trigger_position = trigger_position


class QdacBusCompilationInfo:
    """Class representing the information stored by QdacCompiler for a bus."""

    def __init__(self) -> None:
        # Stacks to manage block hierarchy during compilation
        self.qprogram_block_stack: deque[Block] = deque()

        # Counters to help with naming and indexing
        self.loop_counter = 0

        # Trigger to mark that the dc has been set
        self.dc_set: bool = False


class QdacCompiler:
    """A class for controling QProgram to Qdac hardware."""

    def __init__(self) -> None:
        # Handlers to map each operation to a corresponding handler function
        self._handlers: dict[type, Callable] = {
            InfiniteLoop: self._handle_infinite_loop,
            Parallel: self._handle_parallel,
            Average: self._handle_average,
            ForLoop: self._handle_for_loop,
            Loop: self._handle_loop,
            SetOffset: self._handle_set_offset,
            Play: self._handle_play,
            WaitTrigger: self._handle_wait_trigger,
            SetTrigger: self._handle_set_trigger,
            Block: self._handle_block,
        }

        self._qprogram: QProgram
        self._buses: dict[str, QdacBusCompilationInfo]
        self._qdac_buses: list["Bus"]
        self._qdac_buses_alias: list[str]
        self._channels: dict[str, int]
        self._qdac: QDevilQDac2

        self._dc_dwell: int = 2
        self._dc_delay: int = 0
        self._loop_repetitions: dict[str, int] = {}
        self._infinite_loop: bool = False

        self._trigger_hashes: dict[str, str] = {}
        self._trigger_position: str | None = None

    def compile(
        self,
        qprogram: QProgram,
        qdac: QDevilQDac2,
        qdac_buses: list["Bus"],
        bus_mapping: dict[str, str] | None = None,
        calibration: Calibration | None = None,
    ) -> QdacCompilationOutput:
        """Compile QProgram to qpysequence.Sequence
        Args:
            qprogram (QProgram): The QProgram to be compiled
            bus_mapping (dict[str, str], optional): Optional mapping of bus names. Defaults to None.
            times_of_flight (dict[str, int], optional): Optional time of flight of bus. Defaults to None.
        """

        def traverse(block: Block):
            for bus in self._buses and self._qdac_buses_alias:
                self._buses[bus].qprogram_block_stack.append(block)
            for element in block.elements:
                handler = self._handlers.get(type(element))
                if not handler:
                    self._handle_unknown(element)
                else:
                    handler(element)
                if isinstance(element, Block):
                    traverse(element)
            for bus in self._buses and self._qdac_buses_alias:
                self._buses[bus].qprogram_block_stack.pop()

        self._qprogram = qprogram
        self._qdac = qdac
        self._qdac_buses = qdac_buses
        if bus_mapping is not None:
            self._qprogram = self._qprogram.with_bus_mapping(bus_mapping=bus_mapping)
        if calibration is not None:
            self._qprogram = self._qprogram.with_calibration(calibration=calibration)
        if self._qprogram.has_calibrated_waveforms_or_weights():
            raise RuntimeError(
                "Cannot compile to hardware-native instructions because QProgram contains named operations that are not mapped. Provide a calibration instance containing all necessary mappings."
            )

        self._populate_qdac_buses()

        # Recursive traversal to convert QProgram blocks to Sequence
        traverse(self._qprogram._body)
        return QdacCompilationOutput(qprogram=self._qprogram, qdac=self._qdac, trigger_position=self._trigger_position)

    def _populate_qdac_buses(self):
        """Map each bus in the QProgram to a BusCompilationInfo instance.
        Returns:
            A dictionary where the keys are bus names and the values are BusCompilationInfo objects.
        """

        self._qdac_buses_alias = [bus.alias for bus in self._qdac_buses]
        self._buses = {bus: QdacBusCompilationInfo() for bus in self._qprogram.buses if bus in self._qdac_buses_alias}
        self._loop_repetitions.update(dict.fromkeys(self._qdac_buses_alias, 1))

        self._channels = {bus.alias: bus.channels[0] for bus in self._qdac_buses if bus.alias in self._qdac_buses_alias}
        return

    def _handle_parallel(self, element: Parallel):
        if not element.loops:
            raise NotImplementedError("Parallel block should contain loops.")

        iterations = []
        for loop in element.loops:
            if isinstance(loop, ForLoop):
                iters = QdacCompiler._convert_for_loop_values(loop)
            elif isinstance(loop, Loop):
                iters = int(loop.values.shape[0])
            iterations.append(iters)
        for qdac_bus_alias in self._qdac_buses_alias:
            self._loop_repetitions[qdac_bus_alias] *= min(iterations) + 1

        for bus in self._buses:
            self._buses[bus].loop_counter += 1
        return True

    def _handle_average(self, element: Average):
        for qdac_bus_alias in self._qdac_buses_alias:
            self._loop_repetitions[qdac_bus_alias] *= element.shots + 1
        for bus in self._buses:
            self._buses[bus].loop_counter += 1
        return True

    def _handle_infinite_loop(self, _: InfiniteLoop):
        self._infinite_loop = True
        for bus in self._buses:
            self._buses[bus].loop_counter += 1
        return True

    def _handle_for_loop(self, element: ForLoop):
        iterations = QdacCompiler._convert_for_loop_values(element)
        for qdac_bus_alias in self._qdac_buses_alias:
            self._loop_repetitions[qdac_bus_alias] *= iterations + 1
        for bus in self._buses:
            self._buses[bus].loop_counter += 1
        return True

    def _handle_loop(self, element: Loop):
        for qdac_bus_alias in self._qdac_buses_alias:
            self._loop_repetitions[qdac_bus_alias] *= int(element.values.shape[0]) + 1
        for bus in self._buses:
            self._buses[bus].loop_counter += 1
        return True

    def _handle_set_offset(self, element: SetOffset):
        if element.bus in self._qdac_buses_alias:
            self._qdac.set_parameter(
                parameter=Parameter.VOLTAGE, value=element.offset_path0, channel_id=self._channels[element.bus]
            )
            self._buses[element.bus].dc_set = True

    def _handle_set_trigger(self, element: SetTrigger):
        if element.bus in self._qdac_buses_alias:
            # Condition to check if positions are properly set.
            if element.position not in {"end", "start"}:
                raise NotImplementedError(
                    f"position must be set as 'end' or 'start', {element.position} is not recognized"
                )

            if element.outputs:
                for output in element.outputs if isinstance(element.outputs, list) else [element.outputs]:
                    trigger = self._hash_trigger(element, output)

                    if element.position == "end":
                        self._qdac.set_end_marker_external_trigger(
                            channel_id=self._channels[element.bus],
                            out_port=output,
                            trigger=trigger,
                            width_s=element.duration,
                        )
                    else:
                        self._qdac.set_start_marker_external_trigger(
                            channel_id=self._channels[element.bus],
                            out_port=output,
                            trigger=trigger,
                            width_s=element.duration,
                        )
                if not self._trigger_position:
                    self._trigger_position = "front"
            else:
                trigger = self._hash_trigger(element, None)
                if element.position == "end":
                    self._qdac.set_end_marker_internal_trigger(channel_id=self._channels[element.bus], trigger=trigger)
                else:
                    self._qdac.set_start_marker_internal_trigger(
                        channel_id=self._channels[element.bus], trigger=trigger
                    )

    def _handle_wait_trigger(self, element: WaitTrigger):
        if element.bus in self._qdac_buses_alias:
            if element.port:
                self._qdac.set_in_external_trigger(channel_id=self._channels[element.bus], in_port=element.port)
                if not self._trigger_position:
                    self._trigger_position = "back"
            else:
                self._qdac.set_in_internal_trigger(
                    channel_id=self._channels[element.bus],
                    trigger=next(trigger for _, trigger in self._trigger_hashes.items()),
                )

    def _handle_play(self, element: Play):
        if element.bus in self._qdac_buses_alias:
            convert = QdacCompiler._convert_value(element)
            waveform, _ = element.get_waveforms()
            waveform_variables = element.get_waveform_variables()
            if waveform_variables:
                logger.error("Variables in waveforms are not supported in Qdac.")
                return
            if not element.dwell:
                element.dwell = self._dc_dwell
            if not element.delay:
                element.delay = self._dc_delay
            if not element.repetitions:
                element.repetitions = self._loop_repetitions[element.bus]
            if self._infinite_loop:
                element.repetitions = -1

            self._qdac.upload_voltage_list(
                waveform=waveform,
                channel_id=self._channels[element.bus],
                dwell_us=convert(element.dwell),
                sync_delay_s=element.delay,
                repetitions=element.repetitions,
            )

            self._loop_repetitions[element.bus] = 1

    def _handle_block(self, element: Block):
        pass

    def _handle_unknown(self, element: Any):
        if type(element) in [SetFrequency, SetPhase, ResetPhase, SetGain, SetMarkers, Wait, Measure, Acquire]:
            if element.bus in self._qdac_buses_alias:
                raise NotImplementedError(f"{element.__class__} is not supported in QDACII.")

        if type(element) is Sync:
            if element.buses and any(bus in self._qdac_buses_alias for bus in element.buses):
                raise NotImplementedError(f"{element.__class__} is not supported in QDACII.")

    def _hash_trigger(self, element: SetTrigger, output: int | None):
        hash = f"trigger_{element.bus}_{output}_{element.position}"
        self._trigger_hashes[element.bus] = hash
        return hash

    @staticmethod
    def _calculate_iterations(start: int | float, stop: int | float, step: int | float):
        if step == 0:
            raise ValueError("Step value cannot be zero")

        # Calculate the raw number of iterations
        raw_iterations = (stop - start + step) / step

        # If the raw number of iterations is very close to an integer, round it to that integer
        # This accounts for potential floating-point inaccuracies
        if abs(raw_iterations - round(raw_iterations)) < 1e-9:
            return round(raw_iterations)

        # Otherwise, if we're incrementing, take the ceiling, and if we're decrementing, take the floor
        return math.floor(raw_iterations) if step > 0 else math.ceil(raw_iterations)

    @staticmethod
    def _convert_for_loop_values(for_loop: ForLoop):
        iterations = QdacCompiler._calculate_iterations(start=for_loop.start, stop=for_loop.stop, step=for_loop.step)
        return iterations

    @staticmethod
    def _convert_value(operation: Operation) -> Callable[[Any], Any]:
        conversion_map: dict[type[Operation], Callable[[Any], Any]] = {
            Play: lambda x: float(x * 1e-6),
        }
        return conversion_map.get(type(operation), lambda x: x)
