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

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

import copy
import numpy as np


def conjugate(self, inplace=False):
    """Return Grid object corresponding to conjugated function

    The result has spin weight equal to the negative of this object's spin weight.

    """
    s = self.view(np.ndarray)
    c = s if inplace else np.zeros_like(s)
    np.conjugate(s, out=c)
    if inplace:
        self._metadata['spin_weight'] = -self.spin_weight
        return self
    metadata = copy.copy(self._metadata)
    metadata['spin_weight'] = -self.spin_weight
    return type(self)(c, **metadata)


@property
def bar(self):
    return self.conjugate()
bar.__doc__ = conjugate.__doc__


@property
def real(self):
    """Return real-valued part of function values on the grid

    This only makes sense when the function has spin weight 0; anything else will raise an error.

    """
    if self.spin_weight != 0:
        raise ValueError("The real part of a function with non-zero spin weight is meaningless")
    s = self.view(np.ndarray)
    c = s.real
    return type(self)(c, **self._metadata)


@property
def imag(self):
    """Return imaginary-valued part of function values on the grid

    This only makes sense when the function has spin weight 0; anything else will raise an error.

    """
    if self.spin_weight != 0:
        raise ValueError("The imaginary part of a function with non-zero spin weight is meaningless")
    s = self.view(np.ndarray)
    c = s.imag
    return type(self)(c, **self._metadata)


def absolute(self):
    """Return absolute value of function values on the grid"""
    metadata = copy.copy(self._metadata)
    metadata['spin_weight'] = 0
    return type(self)(np.abs(self.view(np.ndarray)), **metadata)


def add(self, other):
    if isinstance(other, type(self)):
        s = self.spin_weight
        if s != other.spin_weight:
            raise ValueError(f"Cannot add functions with different spin weights ({s} and {other.spin_weight})")
        if self.n_theta != other.n_theta or self.n_phi != other.n_phi:
            raise ValueError(f"Shape mismatch: self.shape={self.shape}; other.shape={other.shape}")
        result = self.view(np.ndarray) + other.view(np.ndarray)
        return type(self)(result, **self._metadata)
    elif self.spin_weight != 0 and np.any(other):
        raise ValueError(
            f"It is not permitted to add non-zero scalars to a {type(self).__name__} object of non-zero spin weight"
        )
    else:
        result = self.view(np.ndarray) + other
        return type(self)(result, **self._metadata)


def subtract(self, other):
    if isinstance(other, type(self)):
        if self.spin_weight != other.spin_weight:
            raise ValueError(f"Cannot subtract functions with different spin weights ({s} and {other.spin_weight})")
        if self.n_theta != other.n_theta or self.n_phi != other.n_phi:
            raise ValueError(f"Shape mismatch: self.shape={self.shape}; other.shape={other.shape}")
        result = self.view(np.ndarray) - other.view(np.ndarray)
        return type(self)(result, **self._metadata)
    elif self.spin_weight != 0 and np.any(other):
        raise ValueError(
            f"It is not permitted to subtract non-zero scalars from a {type(self).__name__} "
            "object of non-zero spin weight"
        )
    else:
        result = self.view(np.ndarray) - other
        return type(self)(result, **self._metadata)


def multiply(self, other, truncator=None):
    """Multiply by another spin-weighted function or a scalar

    For spin-weighted functions, the spin weight of their product is the sum of the spin weights of
    the input.

    Parameters
    ----------
    other : Grid, array_like, complex, or float
        Grid object representing the spin-weighted functions, or an array or float which is
        equivalent to a spin-0 function.

    """
    if isinstance(other, type(self)):
        if self.n_theta != other.n_theta or self.n_phi != other.n_phi:
            raise ValueError(f"Shape mismatch: self.shape={self.shape}; other.shape={other.shape}")
        result = self.view(np.ndarray) * other.view(np.ndarray)
        result_s = self.spin_weight + other.spin_weight
        metadata = copy.copy(self._metadata)
        metadata['spin_weight'] = result_s
        return type(self)(result, **metadata)
    else:
        if self._check_broadcasting(other):
            return type(self)(self.view(np.ndarray) * other, **self._metadata)
        elif self._check_broadcasting(other, reverse=True):
            return type(self)(other * self.view(np.ndarray), **self._metadata)
        else:
            raise ValueError(
                "\n"
                f"    Cannot broadcast input array to this {type(self).__name__} object.  Note that the input array\n"
                f"    must broadcast to all but last two dimensions of this object; it is not allowed to\n"
                f"    multiply each grid value individually.  If you really want to hack this, view this\n"
                f"    object as an ndarray, and don't complain if your results are wrong."
            )


def divide(self, other):
    """Divide by another spin-weighted function or a scalar

    For spin-weighted functions, the spin weight of their ratio is the difference of the spin
    weights of the input.

    Parameters
    ----------
    other : Grid, array_like, complex, or float
        Grid object representing the spin-weighted functions, or an array or float which is
        equivalent to a spin-0 function.

    """
    if isinstance(other, type(self)):
        if self.n_theta != other.n_theta or self.n_phi != other.n_phi:
            raise ValueError(f"Shape mismatch: self.shape={self.shape}; other.shape={other.shape}")
        result = self.view(np.ndarray) / other.view(np.ndarray)
        result_s = self.spin_weight - other.spin_weight
        metadata = copy.copy(self._metadata)
        metadata['spin_weight'] = result_s
        return type(self)(result, **metadata)
    else:
        if self._check_broadcasting(other):
            return type(self)(self.view(np.ndarray) / other, **self._metadata)
        else:
            raise ValueError(
                "\n"
                f"    Cannot broadcast input array to this {type(self).__name__} object.  Note that the input array\n"
                f"    must broadcast to all but last two dimensions of this object; it is not allowed to\n"
                f"    multiply each grid value individually.  If you really want to hack this, view this\n"
                f"    object as an ndarray, and don't complain if your results are wrong."
            )
