# Copyright (c) 2020, Michael Boyle
# See LICENSE file for details: <https://github.com/moble/spherical_functions/blob/master/LICENSE>

### NOTE: The functions in this file are intended purely for inclusion in the Modes class.  In
### particular, they assume that the first argument, `self` is an instance of Modes.  They should
### probably not be used outside of that class.


def index(self, ell, m):
    """Return index of (ell, m) mode in data

    Note that the modes are always stored in the last axis of the data, so if your Modes object
    `modes` has (or even just might have) more than one dimension, you might extract a
    particular mode weight with something like this:

        i = modes.index(ell, m)
        ell_m_modes = modes[..., i]

    The ellipsis just represents all other dimensions (even if there are none).

    """
    from .. import LM_index
    if ell < abs(self.s) or ell < abs(m):
        raise ValueError(f"Bad index (ell, m)=({ell}, {m}) for spin weight s={self.s}")
    if ell < self.ell_min or ell > self.ell_max:
        raise ValueError(f"Requested ell index {ell} outside bounds of this data ({self.ell_min, self.ell_max})")
    return LM_index(ell, m, self.ell_min)


def truncate_ell(self, new_ell_max):
    """Slice array so that new ell max is the given value"""
    if new_ell_max >= self.ell_max:
        return self
    truncated = self[..., :self.index(new_ell_max, new_ell_max)+1]
    truncated._metadata['ell_max'] = new_ell_max
    return truncated


def grid(self, n_theta=None, n_phi=None, **kwargs):
    """Return values of function on an equi-angular grid

    This method uses `spinsfast` to convert mode weights of spin-weighted function to values on a
    grid.  The grid has `n_theta` evenly spaced points along the usual polar (colatitude) angle
    theta, and `n_phi` evenly spaced points along the usual azimuthal angle phi.  This grid
    corresponds to the one produced by `spherical_functions.theta_phi`; see that function for
    specifics.

    The output array has one more dimension than this object; rather than the last axis giving the
    mode weights, the last two axes give the values on the two-dimensional grid.

    Parameters
    ==========
    n_theta: None or int [defaults to None]
        Number of points to use in theta direction.  None is equivalent to 2*self.ell_max+1, which
        is the minimum number that can capture behavior up to and including ell_max.  If you need to
        multiply the result with some `other` spin-weighted function, you should use an n_theta
        value of 2 * (self.ell_max + other.ell_max) + 1 to avoid aliasing.
    n_phi: None or int [defaults to None]
        Number of points to use in the phi direction.  Here, None is equivalent to n_phi=n_theta,
        after calculation of the default value for n_theta.  Note that the same comments apply about
        avoiding aliasing.
    **kwargs: any types
        Additional keyword arguments are passed through to the Grid constructor on output

    """
    import copy
    import numpy as np
    import spinsfast
    from .. import Grid
    n_theta = n_theta or 2*self.ell_max+1
    n_phi = n_phi or n_theta
    metadata = copy.copy(self._metadata)
    metadata.pop('ell_max', None)
    metadata.update(**kwargs)
    return Grid(spinsfast.salm2map(self.view(np.ndarray), self.s, self.ell_max, n_theta, n_phi), **metadata)


def evaluate(self, rotors, **kwargs):
    """Return values of function on input rotors"""
    import numpy as np
    import spherical_functions as sf
    SWSHs = sf.SWSH_grid(rotors, self.spin_weight, self.ell_max)
    return np.tensordot(
        self.view(np.ndarray),
        SWSHs[..., sf.LM_index(self.ell_min, -self.ell_min, 0):sf.LM_index(self.ell_max, self.ell_max, 0)+1],
        axes=([-1], [-1])
    )


def _check_broadcasting(self, array, reverse=False):
    """Test whether or not the given array can broadcast against this object"""
    import numpy as np

    if isinstance(array, type(self)):
        try:
            if reverse:
                np.broadcast(array, self)
            else:
                np.broadcast(self, array)
        except ValueError:
            return False
        else:
            return True
    else:
        if np.ndim(array) > np.ndim(self)-1:
            raise ValueError(f"Cannot broadcast array of {np.ndim(array)} dimensions against {type(self).__name__} "
                             f"object of fewer ({np.ndim(self)-1}) non-mode dimensions.\n"
                             "This is to ensure that scalars do not operate on individual "
                             "mode weights; they must operate on all simultaneously.\n"
                             "If that is the case and you still want to broadcast, add more "
                             "dimensions before this object's first dimension.")
        try:
            if reverse:
                np.broadcast(array, self[..., 0])
            else:
                np.broadcast(self[..., 0], array)
        except ValueError:
            return False
        else:
            return True
