# MIT License
#
# Copyright (c) 2025 Quandela
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.

"""
Pure descriptive components that describe WHAT we want, not HOW to implement it.
Components are platform-agnostic and focus on intent rather than implementation.
"""

from dataclasses import dataclass
from enum import Enum
from typing import Any


class ParameterRole(Enum):
    """Clear role definition for parameters."""

    FIXED = "fixed"  # Fixed value, not a parameter
    INPUT = "input"  # Input parameter (fed from data)
    TRAINABLE = "trainable"  # Trainable parameter (nn.Parameter)


@dataclass
class Rotation:
    """
    Rotation gate description.

    The actual parameter name will be generated by the backend based on role.
    """

    target: int
    role: ParameterRole = ParameterRole.FIXED
    value: float = 0.0  # Used when role=FIXED
    axis: str = "z"

    # Optional: let user specify custom name if they want
    custom_name: str | None = None

    def get_params(self) -> dict[str, Any]:
        """Return declared parameter placeholders for the rotation.

        Non-fixed rotations expose either their custom name or the automatically
        generated identifier so that downstream tooling can bind data or trainable
        tensors to the gate.

        Returns:
            Dict[str, Any]: Mapping from parameter name to placeholder value.
        """
        if self.custom_name and self.role != ParameterRole.FIXED:
            return {
                self.custom_name: self.value
                if self.role == ParameterRole.FIXED
                else None
            }
        return {}


@dataclass
class BeamSplitter:
    """Beam splitter description."""

    targets: tuple
    theta_role: ParameterRole = ParameterRole.FIXED
    theta_value: float = 0.7854  # pi/4 default
    phi_role: ParameterRole = ParameterRole.FIXED
    phi_value: float = 0.0

    # Optional custom names
    theta_name: str | None = None
    phi_name: str | None = None

    def get_params(self) -> dict[str, Any]:
        """Describe which phase shifter angles should be exposed as parameters.

        Returns:
            Dict[str, Any]: Parameter placeholders keyed by their symbolic names.
        """
        params = dict[str, Any]()
        if self.theta_name and self.theta_role != ParameterRole.FIXED:
            params[self.theta_name] = None
        if self.phi_name and self.phi_role != ParameterRole.FIXED:
            params[self.phi_name] = None
        return params


@dataclass
class EntanglingBlock:
    """Entangling block description."""

    targets: str | list[int] = "all"
    pattern: str = "nearest_neighbor"
    depth: int = 1
    trainable: bool = True
    name_prefix: str | None = None

    def get_params(self) -> dict[str, Any]:
        """Entangling blocks themselves carry no direct parameters.

        Returns:
            Dict[str, Any]: Always empty because entangling blocks are metadata-only.
        """
        return {}


@dataclass
class GenericInterferometer:
    """Generic interferometer block spanning multiple modes."""

    start_mode: int
    span: int
    trainable: bool = True
    name_prefix: str | None = None
    model: str = "mzi"
    trainable_inner: bool | None = None
    trainable_outer: bool | None = None

    def __post_init__(self):
        """Validate and normalise the interferometer model."""
        if not isinstance(self.model, str):
            raise TypeError("GenericInterferometer model must be provided as a string")
        normalized = (self.model or "mzi").lower()
        if normalized not in {"mzi", "bell"}:
            raise ValueError(
                "GenericInterferometer model must be either 'mzi' or 'bell'"
            )
        self.model = normalized
        if self.trainable_inner is None:
            self.trainable_inner = self.trainable
        if self.trainable_outer is None:
            self.trainable_outer = self.trainable
        if not self.trainable:
            self.trainable_inner = False
            self.trainable_outer = False
        # Normalise the aggregate flag so downstream logic can rely on it
        self.trainable = bool(self.trainable_inner or self.trainable_outer)

    def get_params(self) -> dict[str, Any]:
        """Return placeholder names for every internal interferometer parameter.

        Returns:
            Dict[str, Any]: Mapping of generated parameter names to ``None`` placeholders.
        """
        if self.span < 2 or not self.trainable:
            return {}

        prefix = self.name_prefix or "gi"
        count = self.span * (self.span - 1) // 2
        params: dict[str, Any] = {}
        for idx in range(count):
            if self.trainable_inner:
                params[f"{prefix}_li{idx}"] = None
            if self.trainable_outer:
                params[f"{prefix}_lo{idx}"] = None
        return params


# For backward compatibility, keep old interface but map to new
class Component:
    """Base class for backward compatibility."""

    pass
