import numpy as np

from ...solutions import BaseSolution
from .base import Action
from ...components.node import OpticalNode


class DCFieldsSolutions(BaseSolution):
    """Contains the result of a :class:`DCFields` action. The returned fields are the
    internal description of the carrier are all nodes, modes, and frequencies. Fields
    are defined as P=|E|^2.

    Attributes
    ----------
    homs : list
        The Hermite-Gaussian higher order mode indices at each node, of each frequency.
    nodes : tuple[str]
        Name of nodes the fields were extracted at
    frequencies : array_like[dtype=float]
        Values of the optical frequencies at each node
    fields : array_like
        Field data array with shape [nodes, frequencies, HOMs.shape[0]]

    Notes
    -----
    The `fields` attribute will contain all the required information with its ordering
    defined by the `nodes`, `frequencies`, and `homs` array. For easier indexing you can
    access the solution directly with node names.

    Select all the frequencies and modes at l1.p1.i:
    >>> sol['l1.p1.i']

    Select all the first modes at all frequencies and modes at l1.p1.i:
    >>> sol['l1.p1.i', :, 0]

    Select all the frequencies and modes at two nodes:
    >>> sol[('l1.p1.i', 'l1.p1.o')]

    The same nodes can be selected directly with the OpticalNodes:
    >>> sol[(model.l1.p1.i, model.l1.p1.o)]

    Select all the HG02 modes at every node and frequency
    >>> sol[:, :, sol.homs.index([0, 2]))]

    Select all the HG02 modes at every output node and every frequency
    >>> idx = np.array([node.split(".")[-1] == "o" for node in sol.nodes])
    >>> sol[idx, :, sol.homs.index([0, 2])]

    The frequency and mode selection must be a slice or an array of integer or boolean
    values.
    """

    homs = None
    nodes = None
    frequencies = None
    fields = None

    def selector(self, key):
        def node_name(key):
            if isinstance(key, str):
                return key
            elif isinstance(key, OpticalNode):
                return key.full_name
            else:
                raise TypeError("key should be a string or optical node")

        if isinstance(key, (str, int, slice, np.ndarray, OpticalNode)):
            # specifying a single node or a slice of multiple nodes
            key = [key]
        elif isinstance(key, tuple):
            # this is either a tuple specifying multiple nodes or
            # specifying multiple dimensions
            assert len(key) > 1
            if len(key) in [2, 3]:
                if isinstance(key[1], (str, OpticalNode)):
                    # specifying multiple nodes
                    key = [key]
                else:
                    # specifying multiple dimensions
                    key = list(key)
            else:
                # specifying multiple nodes
                key = [list(key)]
        else:
            raise TypeError(
                f"Unrecognized key type {type(key)}."
                "Multiple nodes should be specified by tuples, slices, or ndarrays."
            )
        if not isinstance(key[0], (slice, int, np.ndarray)):
            nodes = np.array([node_name(n) for n in np.atleast_1d(key[0])])
            for i in nodes:
                if i not in self.nodes:
                    raise KeyError(f"`{i}` is not in the solution nodes")
            key[0] = tuple(self.nodes.index(i) for i in nodes)

        return tuple(key)

    def __getitem__(self, key):
        if isinstance(key, str) and self.name == key:
            return self
        else:
            return self.fields[self.selector(key)]


class DCFields(Action):
    """An action that saves the DC (carrier) fields at all nodes, frequencies, and
    higher order modes for the current state of the simulation.

    Parameters
    ----------
    name : str, optional
        Name of the solution generated by this action
    """

    def __init__(self, *, name="dcfields"):
        super().__init__(name)

    def _requests(self, model, memo, first=True):
        pass

    def _do(self, state):
        sol = DCFieldsSolutions(self.name)
        # do the DC solve again just in case a parameter has changed and another
        # action hasn't done so already
        state.sim.run_carrier()
        sol.homs = state.model.homs.tolist()
        sol.nodes = tuple(state.sim.carrier.nodes.keys())
        sol.frequencies = np.array(
            [f.f for f in state.sim.carrier.optical_frequencies.frequencies]
        )
        sol.fields = np.reshape(
            state.sim.carrier.M().rhs_view[0, :].copy(),
            (len(sol.nodes), len(sol.frequencies), len(sol.homs)),
            order="C",
        )
        # rescale into sqrt(W) units
        sol.fields /= np.sqrt(2)
        return sol
