# Copyright 2024 IQM
#
# 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.
"""Tools for building instruction schedules.

Note: This is a minimal stub for Kipu Quantum Hub. Only the elements needed for circuit
serialization and validation are included. The full ScheduleBuilder class and
pulse-level scheduling code have been removed as they require proprietary IQM
dependencies (iqm-pulse) that are not publicly available.
"""

from __future__ import annotations

import copy
import logging
from dataclasses import dataclass, field, replace
from typing import Any

# Type alias for qubit names - v33 SDK uses tuples
Locus = tuple[str, ...]


class QuantumOpTable(dict):
    """Dictionary mapping operation names to QuantumOp definitions."""
    pass


@dataclass
class QuantumOp:
    """Describes a quantum operation supported by the QPU."""
    name: str
    arity: int = 0
    params: dict[str, tuple[type, ...]] = field(default_factory=dict)
    optional_params: dict[str, tuple[type, ...]] = field(default_factory=dict)
    implementations: dict[str, Any] = field(default_factory=dict)
    symmetric: bool = False
    factorizable: bool = False
    defaults_for_locus: dict = field(default_factory=dict)

    def get_default_implementation_for_locus(self, locus: list[str]) -> str | None:
        """Get the default implementation name for a specific locus.

        Returns the implementation name from defaults_for_locus if it exists,
        otherwise returns None (no specific implementation required).
        """
        locus_key = tuple(locus) if isinstance(locus, list) else locus
        return self.defaults_for_locus.get(locus_key)


class GateImplementation:
    """Stub implementation that represents a gate that doesn't need calibration."""
    def __init__(self, *args, **kwargs):
        pass

    def needs_calibration(self) -> bool:
        """Return False for stub implementations - operations don't need calibration data."""
        return False


# Stub implementation for operations that bypass architecture validation
_NO_CALIBRATION_IMPL = GateImplementation()

# Define canonical quantum operations library (based on IQM SDK v33 default_gates.py)
# Note: Single-qubit gates (prx, cc_prx) marked factorizable=True so validation
# checks only that qubit exists, not that it's in a specific locus list.
# Operations with implementations={None: _NO_CALIBRATION_IMPL} bypass architecture validation.
_quantum_ops_library: dict[str, QuantumOp] = {
    "barrier": QuantumOp(
        name="barrier", arity=0, symmetric=True, factorizable=True,
        implementations={None: _NO_CALIBRATION_IMPL}  # Bypass architecture validation
    ),
    "delay": QuantumOp(
        name="delay", arity=0, params={"duration": (float,)}, symmetric=True, factorizable=True,
        implementations={None: _NO_CALIBRATION_IMPL}  # Bypass architecture validation
    ),
    "measure": QuantumOp(
        name="measure",
        arity=0,
        params={"key": (str,)},
        optional_params={"feedback_key": (str,)},
        factorizable=True,
    ),
    "prx": QuantumOp(name="prx", arity=1, params={"angle": (float, int), "phase": (float, int)}, factorizable=True),
    "cz": QuantumOp(name="cz", arity=2, symmetric=True),
    "move": QuantumOp(name="move", arity=2),
    "cc_prx": QuantumOp(
        name="cc_prx",
        arity=1,
        params={
            "angle": (float, int),
            "phase": (float, int),
            "feedback_qubit": (str,),
            "feedback_key": (str,),
        },
        factorizable=True,
    ),
    "reset": QuantumOp(name="reset", arity=0, symmetric=True, factorizable=True),
}

_deprecated_ops: set[str] = set()
_deprecated_implementations: dict[str, set[str]] = {}

logger = logging.getLogger(__name__)


@dataclass
class CircuitOperation:
    r"""Specific quantum operation applied on a specific part of the QPU, e.g. in a quantum circuit.

    We currently support the following native operations for circuit execution:

    ================ =========== ======================================= ===========
    name             # of qubits args                                    description
    ================ =========== ======================================= ===========
    measure          >= 1        ``key: str``, ``feedback_key: str``     Measurement in the Z basis.
    prx              1           ``angle: float``, ``phase: float``      Phased x-rotation gate.
    cc_prx           1           ``angle: float``, ``phase: float``,
                                 ``feedback_qubit: str``,
                                 ``feedback_key: str``                   Classically controlled PRX gate.
    reset            >= 1                                                Reset the qubit(s) to :math:`|0\rangle`.
    cz               2                                                   Controlled-Z gate.
    move             2                                                   Move a qubit state between a qubit and a
                                                                         computational resonator, as long as
                                                                         at least one of the components is
                                                                         in the :math:`|0\rangle` state.
    barrier          >= 1                                                Execution barrier.
    delay            >= 1        ``duration: float``                     Force a delay between circuit operations.
    ================ =========== ======================================= ===========

    For each CircuitOperation you may also optionally specify :attr:`implementation`,
    which contains the name of an implementation of the operation to use.
    Support for multiple implementations is currently experimental and in normal use the
    field should be omitted, this selects the default implementation for the operation for that locus.

    """

    name: str
    """name of the quantum operation"""
    locus: Locus
    """names of the information-bearing QPU components (qubits, computational resonators...) the operation acts on"""
    args: dict[str, Any] = field(default_factory=dict)
    """arguments for the operation"""
    implementation: str | None = None
    """name of the implementation"""

    def validate(self, op_table: QuantumOpTable) -> None:
        """Validate the operation against a table of operation definitions.

        Args:
            op_table: table containing allowed quantum operations

        Raises:
            ValueError: operation is not valid

        """
        op_type = op_table.get(self.name)
        if op_type is None:
            message = ", ".join(op_table)
            raise ValueError(f"Unknown operation '{self.name}'. Supported operations are '{message}'.")
        self._validate_implementation(op_type)
        self._validate_locus(op_type)
        self._validate_args(op_type)

    def _validate_implementation(self, op_type: QuantumOp) -> None:
        if self.implementation is not None:
            if not self.implementation:
                raise ValueError("Implementation of the instruction should be None, or a non-empty string")
            if self.implementation not in op_type.implementations:
                raise ValueError(f"Unknown implementation '{self.implementation}' for quantum operation '{self.name}'.")

    def _validate_locus(self, op_type: QuantumOp) -> None:
        arity = op_type.arity
        if (0 < arity) and (arity != len(self.locus)):
            raise ValueError(
                f"The '{self.name}' operation acts on {arity} qubit(s), but {len(self.locus)} were given: {self.locus}."
            )
        if len(self.locus) != len(set(self.locus)):
            raise ValueError(f"Repeated locus components: {self.locus}.")

    def _validate_args(self, op_type: QuantumOp) -> None:
        # Check argument names
        submitted_arg_names = set(self.args)
        allowed_arg_types = op_type.params | op_type.optional_params
        if not set(op_type.params) <= submitted_arg_names:
            raise ValueError(
                f"The operation '{self.name}' requires "
                f"the argument(s) {tuple(op_type.params)}, "
                f"but {tuple(submitted_arg_names)} were given."
            )

        if not submitted_arg_names <= set(allowed_arg_types):
            allowed_arg_names = tuple(allowed_arg_types)
            message = f"the arguments {allowed_arg_names}" if allowed_arg_names else "no arguments"
            raise ValueError(f"The operation '{self.name}' allows {message}, but {submitted_arg_names} were given.")
        # Check argument types
        for arg_name, arg_value in self.args.items():
            allowed_types = allowed_arg_types[arg_name]
            if not isinstance(arg_value, allowed_types):
                raise ValueError(
                    f"The argument '{arg_name}' should be of one of the following supported types"
                    f" {allowed_types}, but ({type(arg_value)}) was given."
                )


def build_quantum_ops(ops: dict[str, dict[str, Any]]) -> QuantumOpTable:
    """Builds the table of known quantum operations.

    Hardcoded canonical ops table is extended by the ones in ``ops``.
    In case of name collisions, the content of ``ops`` takes priority over the defaults,
    with the following caveats:

    * canonical implementation names cannot be redefined, and
    * for canonical operations you can only change :attr:`implementations` and :attr:`defaults_for_locus`.

    Args:
        ops: Contents of the ``gate_definitions`` section defining
            the quantum operations in the configuration YAML file.
            NOTE: Modified by the function.

    Returns:
        Mapping from quantum operation names to their definitions.

    """
    # build the table of native ops
    op_table = {}
    for op_name, op in _quantum_ops_library.items():
        # filter out deprecated ops, unless requested by the user in ``ops``
        if op_name in _deprecated_ops and op_name not in ops:
            continue

        new_op = copy.deepcopy(op)  # to prevent modifications to the hardcoded table
        if (deprecated_impls := _deprecated_implementations.get(op_name)) is not None:
            # filter out deprecated implementations
            non_deprecated = {
                impl_name: impl for impl_name, impl in op.implementations.items() if impl_name not in deprecated_impls
            }
            new_op = replace(op, implementations=non_deprecated)
        op_table[op_name] = new_op

    # add ops defined by the user (simplified - no implementation class resolution)
    for op_name, op_definition in ops.items():
        # validate defaults_for_locus
        defaults_for_locus: dict[Locus, str] = op_definition.pop("defaults_for_locus", {})
        implementations = op_definition.pop("implementations", {})

        if (old_op := _quantum_ops_library.get(op_name)) is not None:
            # modify a canonical operation - only implementations and defaults_for_locus can be changed
            if op_definition:
                logger.warning(
                    f"'{op_name}' is a canonical operation, which means the fields {set(op_definition)} "
                    "provided by the user may not be changed."
                )
            op = replace(old_op, implementations=implementations, defaults_for_locus=defaults_for_locus)
        else:
            # entirely new quantum operation defined by the user
            op = QuantumOp(
                name=op_name,
                implementations=implementations,
                defaults_for_locus=defaults_for_locus,
                **op_definition,
            )
        op_table[op_name] = op

    return op_table
