"""Space-type objects representing physical distances between components."""

import logging
import types
import finesse
import numpy as np

from finesse.symbols import Variable, Matrix
from finesse.env import warn
from finesse.parameter import float_parameter
from finesse.exceptions import ContextualTypeError, ContextualValueError
from finesse.components.general import (
    Connector,
    InteractionType,
    borrows_nodes,
    LocalDegreeOfFreedom,
)
from finesse.components.readout import _Readout
from finesse.components.workspace import ConnectionSetting
from finesse.components.node import NodeType, Port, NodeDirection
from finesse.tracing import abcd

LOGGER = logging.getLogger(__name__)


@borrows_nodes()
@float_parameter(
    "L",
    "Length",
    validate="_check_L",
    units="m",
    is_geometric=True,
)
@float_parameter(
    "nr",
    "Refractive index",
    validate="_check_nr",
    is_geometric=True,
    changeable_during_simulation=False,
)
@float_parameter("user_gouy_x", "Gouy phase (x)", units="Degrees")
@float_parameter("user_gouy_y", "Gouy phase (y)", units="Degrees")
# IMPORTANT: renaming this class impacts the katscript spec and should be avoided!
class Space(Connector):
    """Represents a space between two components in the interferometer configuration,
    with a given length and index of refraction.

    There can be many spaces in a model which are not of much interest and do not
    need to be referred to. For example, the ``link`` command will create spaces
    between components without giving an explicit name, just specifying a distance
    between them.

    All Space elements are added to the Model's ``.spaces`` namespace so they can
    all easily be iterated over. You can also find spaces which had no name
    specified and thus have an autogenerated name. User specified Space's with
    names will be put in the main Model namespace so that they can be accessed
    like any other element in the model.

    Parameters
    ----------
    name : str, optional
        Name of newly created space. If not specified, a name is automatically
        generated.

    portA, portB : :class:`.Port`
        Ports to connect.

    L : float, optional
        Geometric length of newly created :class:`.Space` instance; defaults to 0.

    nr : float, optional
        Index of refraction of newly created :class:`.Space` instance; defaults to 1.0.

    user_gouy_x, user_gouy_y : float, optional
        User-defined gouy phase to override the calculated value.
    """

    def __init__(
        self,
        name,
        portA,
        portB,
        L=0.0,
        nr=1.0,
        user_gouy_x=None,
        user_gouy_y=None,
    ):
        # Override None name with auto https://gitlab.com/ifosim/finesse/finesse3/-/merge_requests/251
        if (portA is None) != (portB is None):
            warn(
                "Can't construct a space with only one port connected; ignoring ports."
            )
            portA = None
            portB = None

        if portA is not None and not isinstance(portA, Port):
            raise ContextualTypeError("portA", portA, allowed_types=(Port,), name=name)
        if portA is not None and not isinstance(portB, Port):
            raise ContextualTypeError("portB", portB, allowed_types=(Port,), name=name)

        if portA is not None and portA.type != NodeType.OPTICAL:
            raise ContextualValueError(
                {"portA": portA},
                f"must be an optical port for {name}.portA",
            )
        if portB is not None and portB.type != NodeType.OPTICAL:
            raise ContextualValueError(
                {"portB": portB},
                f"must be an optical port for {name}.portB",
            )

        if portA is not None and portB is not None:
            if portA is portB:
                raise ValueError("Port A and B are the same port.")
            for n1 in portA.nodes:
                for n2 in portB.nodes:
                    if n1 is n2:
                        raise ValueError(f"Port A and B contain the same node {n1}")

            for port in (portA, portB):
                if isinstance(port.component, _Readout):
                    # check if readout has copied over nodes from another port
                    port.component.check_connection()

            if portA.component._model is not portB.component._model:
                raise ValueError("Port A and B are not part of the same model")

        if name is None:
            auto_generated = True
            if portA is not None and portB is not None:
                compA = portA.component.name
                compB = portB.component.name
                name = f"{compA}_{portA.name}__{compB}_{portB.name}"
            else:
                raise ValueError(
                    "Cannot create an unconnected space without providing a name"
                )
            self._auto_generated = auto_generated
            self._auto_generated_name = name
        else:
            self._auto_generated = False
            self._auto_generated_name = None

        if self._auto_generated:
            super().__init__(self._auto_generated_name)
            self._namespace = (".spaces",)
        else:
            super().__init__(name)
            # Also put into main namespace if it has a specific name
            self._namespace = (".", ".spaces")

        self.__portA = portA
        self.__portB = portB
        self._add_to_model_namespace = True

        self.L = L
        self.nr = nr
        self.user_gouy_x = user_gouy_x
        self.user_gouy_y = user_gouy_y

        self._add_port("p1", NodeType.OPTICAL)
        self._add_port("p2", NodeType.OPTICAL)

        # Phase modulation input
        self._add_port("phs", NodeType.ELECTRICAL)
        self.phs._add_node("i", NodeDirection.INPUT)
        # Amplitude modulation input
        self._add_port("amp", NodeType.ELECTRICAL)
        self.amp._add_node("i", NodeDirection.INPUT)
        # strain input
        self._add_port("h", NodeType.ELECTRICAL)
        self.h._add_node("i", NodeDirection.INPUT)

        if portA is not None and portB is not None:
            self.connect(portA, portB)

        # Define typical degrees of freedom for this component
        self.dofs = types.SimpleNamespace()
        # Strain doesn't have a DC term really in Finesse
        # changing a space length won't generate a signal
        self.dofs.h = LocalDegreeOfFreedom(f"{self.name}.dofs.h", None, self.h.i, 1)

    def optical_equations(self):
        """Calculates the optical equations for the space component.

        Returns
        -------
            dict: A dictionary containing the optical equations for the space component.
                The keys match the optical couplings defined by the Space, and the values
                represent the corresponding complex coupling equations.
        """
        with finesse.symbols.simplification():
            settings = self._model._settings
            _f_ = Variable("_f_")
            _c0_ = finesse.symbols.CONSTANTS["c0"]
            pi = finesse.symbols.CONSTANTS["pi"]
            gouy = Matrix(f"{self.name}.gouy")

            pre_factor = 2 * pi * self.nr.ref * self.L.ref / _c0_
            phi = pre_factor * _f_

            if settings.is_modal:
                return {
                    f"{self.name}.P1i_P2o": np.exp(-1j * phi) * gouy,
                    f"{self.name}.P2i_P1o": np.exp(-1j * phi) * gouy,
                }
            else:
                # Otherwise we just get the phase from the propagation along the
                # space if the frequency is different to that of the model reference
                # wavelength
                return {
                    f"{self.name}.P1i_P2o": np.exp(-1j * phi),
                    f"{self.name}.P2i_P1o": np.exp(-1j * phi),
                }

    @property
    def gouy_x(self):
        # return user value if it is set
        # otherwise try to return calculated value or default

        if self.user_gouy_x.value is not None:
            return self.user_gouy_x.value
        else:
            try:
                return self.gouy(self.portA.o.qx, self.portB.i.qx)
            except RuntimeError:
                # No simulation, use the default value
                return 0.0

    @property
    def gouy_y(self):
        # return user value if it is set
        # otherwise try to return calculated value or default
        if self.user_gouy_y.value is not None:
            return self.user_gouy_y.value
        else:
            try:
                return self.gouy(self.portA.o.qy, self.portB.i.qy)
            except RuntimeError:
                # No simulation, use the default value
                return 0.0

    def gouy(self, q1, q2):
        """Computes the Gouy phase in degrees from beam parameters.

        Parameters
        ----------
        q1 : complex, :class:`.BeamParam`
            Starting beam parameter

        q2 : complex, :class:`.BeamParam`
            Ending beam parameter

        Returns
        -------
        out : float
            Gouy phase (in degrees)
        """
        return np.degrees(
            abs(np.arctan2(q1.real, q1.imag) - np.arctan2(q2.real, q2.imag))
        )

    @property
    def portA(self):
        return self.__portA

    @property
    def portB(self):
        return self.__portB

    def connect(self, portA, portB):
        """Sets the ports of this `Space`.

        Parameters
        ----------
        portA : :class:`.Port`, optional
            Port to connect

        portB : :class:`.Port`, optional
            Port to connect
        """
        if portA.is_connected:
            raise Exception(f"Port {portA} has already been connected to")
        if portB.is_connected:
            raise Exception(f"Port {portB} has already been connected to")

        # From the Space's perspective the input and output
        # nodes are swapped around for its ports
        self.p1._add_node("i", None, portA.o)
        self.p1._add_node("o", None, portA.i)
        self.p2._add_node("i", None, portB.o)
        self.p2._add_node("o", None, portB.i)

        self._register_node_coupling(
            "P1i_P2o",
            self.p1.i,
            self.p2.o,
            interaction_type=InteractionType.TRANSMISSION,
        )
        self._register_node_coupling(
            "P2i_P1o",
            self.p2.i,
            self.p1.o,
            interaction_type=InteractionType.TRANSMISSION,
        )
        self._register_node_coupling("SIGPHS_P1o", self.phs.i, self.p1.o)
        self._register_node_coupling("SIGPHS_P2o", self.phs.i, self.p2.o)
        self._register_node_coupling("SIGAMP_P1o", self.amp.i, self.p1.o)
        self._register_node_coupling("SIGAMP_P2o", self.amp.i, self.p2.o)
        self._register_node_coupling("H_P1o", self.h.i, self.p1.o)
        self._register_node_coupling("H_P2o", self.h.i, self.p2.o)

    def _get_workspace(self, sim):
        from finesse.components.modal.space import (
            space_carrier_fill,
            space_signal_fill,
            space_set_gouy,
            SpaceWorkspace,
        )

        _, is_changing = self._eval_parameters()

        carrier_refill = sim.carrier.any_frequencies_changing and (
            float(self.L.value) > 0 or self.L.is_changing
        )
        carrier_refill |= len(is_changing)
        carrier_refill |= self in sim.trace_forest
        carrier_refill |= self.user_gouy_x.is_changing
        carrier_refill |= self.user_gouy_y.is_changing

        ws = SpaceWorkspace(self, sim)

        # use user gouy if it is defined
        ws.use_user_gouy_x = self.user_gouy_x.value is not None
        ws.use_user_gouy_y = self.user_gouy_y.value is not None

        ws.set_gouy_function(space_set_gouy)
        # Set the fill function for this simulation
        ws.carrier.add_fill_function(space_carrier_fill, carrier_refill)
        ws.carrier.connection_settings["P1i_P2o"] = ConnectionSetting.DIAGONAL
        ws.carrier.connection_settings["P2i_P1o"] = ConnectionSetting.DIAGONAL

        if sim.signal:
            signal_refill = sim.signal.any_frequencies_changing and (
                # only non-zero spaces which aren't changing need to be added
                # if signal frequency is changing
                float(self.L.value) > 0
                or self.L.is_changing
            )
            signal_refill |= self in sim.trace_forest
            signal_refill |= self.user_gouy_x.is_changing
            signal_refill |= self.user_gouy_y.is_changing
            ws.signal.add_fill_function(space_signal_fill, signal_refill)

            ws.signal.connection_settings["P1i_P2o"] = ConnectionSetting.DIAGONAL
            ws.signal.connection_settings["P2i_P1o"] = ConnectionSetting.DIAGONAL
            ws.signal.connection_settings["SIGPHS_P1o"] = ConnectionSetting.DIAGONAL
            ws.signal.connection_settings["SIGPHS_P2o"] = ConnectionSetting.DIAGONAL
            ws.signal.connection_settings["SIGAMP_P1o"] = ConnectionSetting.DIAGONAL
            ws.signal.connection_settings["SIGAMP_P2o"] = ConnectionSetting.DIAGONAL
            ws.signal.connection_settings["H_P1o"] = ConnectionSetting.DIAGONAL
            ws.signal.connection_settings["H_P2o"] = ConnectionSetting.DIAGONAL

        # Initialise the ABCD matrix memory-views
        if sim.is_modal:
            ws.abcd = self.ABCD(self.p1.i, self.p2.o, "x", copy=False)

        return ws

    def _check_L(self, value):
        if value < 0:
            raise ValueError("Length of a space must not be negative.")

        return value

    def _check_nr(self, value):
        if value < 1:
            raise ValueError("Index of refraction must be >= 1")

        return value

    def _resymbolise_ABCDs(self):
        L = self.L.ref
        nr = self.nr.ref
        M_sym = abcd.space(L, nr)

        # Matrices same in both propagations and both planes so
        # only need one register call for all couplings here
        self.register_abcd_matrix(
            M_sym,
            (self.p1.i, self.p2.o),
            (self.p2.i, self.p1.o),
        )

    @property
    def abcd(self):
        """Numeric ABCD matrix.

        Equivalent to any of ``space.ABCD(1, 2, "x")``, ``space.ABCD(2, 1, "x")``,
        ``space.ABCD(1, 2, "y")``, ``space.ABCD(2, 1, "y")``.

        :`getter`: Returns a copy of the (numeric) ABCD matrix (read-only).
        """
        return self.ABCD(1, 2, "x")

    def ABCD(
        self,
        from_node,
        to_node,
        direction="x",
        symbolic=False,
        copy=True,
        retboth=False,
        allow_reverse=False,
    ):
        r"""Returns the ABCD matrix of the space for the specified coupling.

        .. image:: /images/space/abcd_spacet.*
            :align: center

        This is given by,

        .. math::
            M = \begin{pmatrix}
                    1 & \frac{L}{n_r} \\
                    0 & 1
                \end{pmatrix},

        where :math:`L` is the length of the space and :math:`n_r` is
        the index of refraction.

        See :meth:`.Connector.ABCD` for descriptions of parameters, return values and possible
        exceptions.
        """
        return super().ABCD(
            from_node, to_node, direction, symbolic, copy, retboth, allow_reverse
        )

    # replacing spaces makes no sense, you can only replace with another space
    def replace(self, *args, **kwargs) -> str:
        raise NotImplementedError(
            "Replacing a space component is currently not implemented. Consider "
            "breaking up the space into two spaces and a Nothing component to insert "
            "new components in the middle of a space"
        )
