"""Space-type objects representing physical distances between components."""
import logging
import types
import numpy as np
from ..env import warn
from ..parameter import float_parameter
from ..exceptions import ContextualTypeError, ContextualValueError
from ..components.general import (
    Connector,
    InteractionType,
    borrows_nodes,
    DOFDefinition,
)
from ..components.workspace import ConnectionSetting
from ..components.node import NodeDirection, NodeType, Port
from ..tracing import abcd
LOGGER = logging.getLogger(__name__)
[docs]@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")
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`, optional
        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.
    """
[docs]    def __init__(
        self,
        name,
        portA=None,
        portB=None,
        L=0.0,
        nr=1.0,
        user_gouy_x=None,
        user_gouy_y=None,
    ):
        given_name = name
        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,))
        if portA is not None and not isinstance(portB, Port):
            raise ContextualTypeError("portB", portB, allowed_types=(Port,))
        if portA is not None and portA.type != NodeType.OPTICAL:
            raise ContextualValueError({"portA": portA}, "must be an optical port")
        if portB is not None and portB.type != NodeType.OPTICAL:
            raise ContextualValueError({"portB": portB}, "must be an optical port")
        if portA is not None and portB is not None:
            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:
            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"
                )
        super().__init__(name)
        self.__portA = portA
        self.__portB = portB
        self._add_to_model_namespace = True
        if given_name is None:
            self._namespace = (".spaces",)
        else:
            # Also put into main namespace if it has a specific name
            self._namespace = (".", ".spaces")
        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)
        self.__changing_check = set((self.L, self.nr))
        # 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 = DOFDefinition(f"{self.name}.dofs.h", None, self.h.i, 1) 
    @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
[docs]    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
[docs]    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")
[docs]    def ABCD(
        self,
        from_node,
        to_node,
        direction="x",
        symbolic=False,
        copy=True,
        retboth=False,
    ):
        r"""Returns the ABCD matrix of the space for the specified coupling.
        .. _fig_abcd_space_transmission:
        .. figure:: ../images/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)