# Copyright 2023 Agnostiq Inc.
#
# This file is part of Covalent.
#
# Licensed under the Apache License 2.0 (the "License"). A copy of the
# License may be obtained with this software package or at
#
#     https://www.apache.org/licenses/LICENSE-2.0
#
# Use of this file is prohibited except in compliance with the License.
# 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.

"""QNode and QElectron information containers."""

from typing import Any, Dict, Optional, Sequence, Union

from pydantic import BaseModel  # pylint: disable=no-name-in-module


class QNodeSpecs(BaseModel):
    """
    A container for the specifications of a QNode generated by `qml.specs`.
    """

    gate_sizes: Dict[int, int]
    gate_types: Dict[str, int]
    num_operations: int
    num_observables: int
    num_diagonalizing_gates: int
    num_used_wires: int
    depth: int
    num_trainable_params: int = None
    num_device_wires: int
    device_name: str
    diff_method: Optional[str] = None
    expansion_strategy: str
    gradient_options: Dict[str, int]
    interface: Optional[str] = None
    gradient_fn: Any = None  # can be string or `qml.gradients.gradient_transform`
    num_gradient_executions: Any = 0
    num_parameter_shift_executions: int = None


class QElectronInfo(BaseModel):
    """
    A container for global and QNode device settings.
    """

    name: str
    description: Optional[str] = None
    device_name: str  # name of the original device, e.g. "default.qubit"
    device_import_path: str  # used to inherit type converters and other methods
    device_shots: Union[
        None, int, Sequence[int], Sequence[Union[int, Sequence[int]]]
    ] = None  # optional default for execution devices
    device_shots_type: Any = None
    device_wires: int  # this can not be reliably inferred from tapes alone
    pennylane_active_return: bool  # client-side status of `pennylane.active_return()`
