import autograd.numpy as np
import numpy.ma as ma
import proxmin.operators
from abc import abstractmethod

from .bbox import Box, overlapped_slices
from .constraint import (
    ConstraintChain,
    L0Constraint,
    PositivityConstraint,
    MonotonicityConstraint,
    MonotonicMaskConstraint,
    SymmetryConstraint,
    CenterOnConstraint,
    NormalizationConstraint,
)
from .frame import Frame
from .model import Model, UpdateException
from .parameter import Parameter, prepare_param, relative_step
from .psf import PSF
from .wavelet import Starlet, starlet_reconstruction, get_multiresolution_support
from . import fft
from . import initialization as init


class Morphology(Model):
    """Morphology base class

    The class describes the 2D image of the spatial dependence of
    `~scarlet.FactorizedComponent`.

    Parameters
    ----------
    frame: `~scarlet.Frame`
        Characterization of the model
    parameters: list of `~scarlet.Parameter`
    bbox: `~scarlet.Box`
        2D bounding box of this model
    """

    def __init__(self, frame, *parameters, bbox=None):
        assert isinstance(frame, Frame)
        self.frame = frame

        if bbox is None:
            bbox = frame.bbox
        assert isinstance(bbox, Box)
        self.bbox = bbox

        super().__init__(*parameters)

    def shrink_box(self, image, thresh=0):
        # peel the onion
        size = max(image.shape)
        dist = 0
        while (
            np.all(image[dist, :] <= thresh)
            and np.all(image[-dist - 1, :] <= thresh)
            and np.all(image[:, dist] <= thresh)
            and np.all(image[:, -dist - 1] <= thresh)
        ):
            dist += 1
        newsize = init.get_minimal_boxsize(size - 2 * dist)
        if newsize < size:
            dist = (size - newsize) // 2
            # adjust bbox
            self.bbox.origin = tuple(o + dist for o in self.bbox.origin)
            self.bbox.shape = (newsize, newsize)


class ImageMorphology(Morphology):
    """Morphology from an image

    The class uses an arbitrary image as non-parametric model. To allow for subpixel
    offsets, a Fourier-based shifting transformation is available.

    Parameters
    ----------
    frame: `~scarlet.Frame`
        Characterization of the model
    image: 2D array or `~scarlet.Parameter`
        Image parameter
    bbox: `~scarlet.Box`
        2D bounding box for focation of the image in `frame`
    shift: None or `~scarlet.Parameter`
        2D shift parameter (in units of image pixels)
    resizing: bool
        Whether to resize the box dynamically
    """

    def __init__(
        self, frame, image, bbox=None, shifting=False, shift=None, resizing=True
    ):
        if isinstance(image, Parameter):
            assert image.name == "image"
        else:
            constraint = PositivityConstraint()
            image = Parameter(
                image, name="image", step=relative_step, constraint=constraint
            )

        if bbox is None:
            assert frame.bbox[1:].shape == image.shape
            bbox = Box(image.shape)
        else:
            assert bbox.shape == image.shape

        self.resizing = resizing
        self.shifting = shifting

        # create the shift parameter to allow for dynamic resizing
        if shift is None:
            shift = Parameter(np.zeros(2), name="shift", step=1e-2, fixed=self.shifting)
        else:
            assert shift.shape == (2,)
            if isinstance(shift, Parameter):
                assert shift.name == "shift"
            else:
                shift = Parameter(shift, name="shift", step=1e-2)

        parameters = (image, shift)
        super().__init__(frame, *parameters, bbox=bbox)

    def get_model(self, *parameters):
        image = self.get_parameter(0, *parameters)
        shift = self.get_parameter(1, *parameters)

        if self.shifting:
            image = fft.shift(image, shift, return_Fourier=False)
        return image

    def update(self):
        image = self._parameters[0]

        if not self.resizing or image.fixed:
            return

        # shrink the box?
        bbox = self.bbox.copy()
        self.shrink_box(image)
        if bbox != self.bbox:
            slice, _ = overlapped_slices(bbox, self.bbox)
            image = Parameter(
                image[slice],
                name=image.name,
                prior=image.prior,
                constraint=image.constraint,
                step=image.step / 2,
                fixed=image.fixed,
                m=image.m[slice] if image.m is not None else None,
                v=image.v[slice] if image.v is not None else None,
                vhat=image.vhat[slice] if image.vhat is not None else None,
            )

            # set new parameters
            self._parameters = (image,) + self._parameters[1:]

            raise UpdateException

        # grow the box?
        # because the PSF moves power across the box, the gradients at the edge
        # accummulate flux from beyond the box
        elif image.m is not None:
            # next adam gradient update
            gu = -image.m / np.sqrt(np.sqrt(ma.masked_equal(image.v, 0))) * image.step
            gu_pull = gu * (image > 0)  # check if model has flux at the edge at all
            edge_pull = np.array(
                (
                    gu_pull[:, 0].mean(),
                    gu_pull[:, -1].mean(),
                    gu_pull[0, :].mean(),
                    gu_pull[-1, :].mean(),
                )
            )

            # 0.1 compared to 1 at center
            if np.any(edge_pull > 0.1):
                # find next larger boxsize
                size = max(bbox.shape)
                newsize = init.get_minimal_boxsize(size + 1)
                pad_width = (newsize - size) // 2

                # Create new parameter for extended image
                image = Parameter(
                    np.pad(image, pad_width, mode="linear_ramp"),
                    name=image.name,
                    prior=image.prior,
                    constraint=image.constraint,
                    step=image.step / 2,
                    fixed=image.fixed,
                    m=np.pad(image.m, pad_width, mode="constant")
                    if image.m is not None
                    else None,
                    v=np.pad(image.v, pad_width, mode="constant")
                    if image.v is not None
                    else None,
                    vhat=np.pad(image.vhat, pad_width, mode="constant")
                    if image.vhat is not None
                    else None,
                )
                # set new parameters
                self._parameters = (image,) + self._parameters[1:]

                # adjust bbox
                self.bbox.origin = tuple(o - pad_width for o in self.bbox.origin)
                self.bbox.shape = (newsize, newsize)
                raise UpdateException


class ProfileMorphology(Morphology):
    """Morphology from a radial profile

    Parameters
    ----------
    frame: `~scarlet.Frame`
        Characterization of the model
    func:
        Radial profile function, signature func(R2, *parameters),
        where R2 is the squared radius.
    parameters: list of `~scarlet.Parameter`
        Parameters for the radial profile. Needs to include at least the following:
        * "radius": float
        * "center": 2D float
        * "ellipticity": 2D float
    boxsize: int
        Size of bounding box over which to evaluate the function, in frame pixels
    resizing: bool
        Whether to resize the box dynamically
    """

    def __init__(self, frame, func, *parameters, boxsize=None, resize=True):

        # define radial profile function
        self.f = func

        # retain center attribute
        center = self.get_parameter("center", *parameters)
        self.center = center

        # get bounding box
        bbox = self.get_box(*parameters, boxsize=boxsize)
        self.resizing = resize

        # x/y evaluation locations
        self._Y = np.arange(bbox.shape[-2], dtype="float") + bbox.origin[-2]
        self._X = np.arange(bbox.shape[-1], dtype="float") + bbox.origin[-1]

        # make sure radius stays positive
        radius = self.get_parameter("radius", *parameters)
        radius.constraint = self._radius_prox

        # make sure ellipticity stays in unit circle
        eps = self.get_parameter("ellipticity", *parameters)
        eps.constraint = self._eps_prox

        super().__init__(frame, *parameters, bbox=bbox)

    def get_model(self, *parameters):
        center = self.get_parameter("center", *parameters)
        _Y = self._Y - center[-2]
        _X = self._X - center[-1]

        e = self.get_parameter("ellipticity", *parameters)
        if np.all(e == 0) and not parameters:  # need e for gradients even if 0
            R2 = _Y[:, None] ** 2 + _X[None, :] ** 2
        else:
            e1, e2 = e[0], e[1]
            __X = ((1 - e1) * _X[None, :] - e2 * _Y[:, None]) / np.sqrt(
                1 - (e1 ** 2 + e2 ** 2)
            )
            __Y = (-e2 * _X[None, :] + (1 + e1) * _Y[:, None]) / np.sqrt(
                1 - (e1 ** 2 + e2 ** 2)
            )
            R2 = __Y ** 2 + __X ** 2

        Rp = self.get_parameter("radius", *parameters)
        R2 /= Rp ** 2

        morph = self.f(R2, *parameters)

        return morph

    @property
    @abstractmethod
    def integral(self):
        pass

    def update(self):
        if not self.resizing:
            return

        bbox = self.get_box()
        if bbox != self.bbox:
            # adjust bbox
            self.bbox.origin = bbox.origin
            self.bbox.shape = bbox.shape
            # x/y evaluation locations
            self._Y = np.arange(bbox.shape[-2], dtype="float") + bbox.origin[-2]
            self._X = np.arange(bbox.shape[-1], dtype="float") + bbox.origin[-1]
            raise UpdateException

    def get_box(self, *parameters, boxsize=None):
        if boxsize is None:
            Rp = self.get_parameter("radius", *parameters)
            size = 10 * Rp
            boxsize = init.get_minimal_boxsize(size)
        shape = (boxsize, boxsize)

        center = self.get_parameter("center", *parameters)
        assert center is not None and len(center) >= 2
        # define bbox around center
        origin = (
            int(round(center[-2])) - (boxsize // 2),
            int(round(center[-1])) - (boxsize // 2),
        )
        bbox = Box(shape, origin=origin)
        return bbox

    def _radius_prox(self, x, step):
        return np.maximum(x, 1e-2)

    def _eps_prox(self, x, step):
        norm2 = (x ** 2).sum()
        if norm2 > 1:
            x /= np.sqrt(norm2) * 1.1  # need to get inside of unit circle
        return x


class GaussianMorphology(ProfileMorphology):
    """Morphology from a Gaussian radial profile

    Parameters
    ----------
    frame: `~scarlet.Frame`
        Characterization of the model
    center: array or `~scarlet.Parameter`
        2D center parameter (in units of frame pixels)
        If it is to be optimized, it needs to be provided as a full defined `Parameter`.
    sigma: float or `~scarlet.Parameter`
        Standard deviation of the Gaussian in frame pixels
        If it is to be optimized, it needs to be provided as a full defined `Parameter`.
    ellipticity: array or ~scarlet.Parameter`
        Two-component ellipticity (e1,e2).
        If it is to be optimized, it needs to be provided as a full defined `Parameter`.
    boxsize: int
        Size of bounding box over which to evaluate the function, in frame pixels
    """

    def __init__(self, frame, center, sigma, ellipticity=(0, 0), boxsize=None):

        assert len(center) == 2
        self.center = prepare_param(center, name="center")
        radius = prepare_param(sigma, name="radius")
        assert len(ellipticity) == 2
        ellipticity = prepare_param(ellipticity, name="ellipticity")
        parameters = (self.center, radius, ellipticity)

        if boxsize is None:
            boxsize = int(np.ceil(10 * sigma))

        super().__init__(frame, self._f, *parameters, boxsize=boxsize)

    def _f(self, R2, *parameters):
        return np.exp(-R2 / 2)

    @property
    def integral(self):
        radius = self.get_parameter("radius")
        return 2 * np.pi * radius ** 2


# we need gamma and kv from scipy.special
from autograd.scipy.special import gamma
from scipy.optimize import fsolve

# but kv is not ported to autograd ...
import scipy.special
from autograd.extend import primitive, defvjp

kv = primitive(scipy.special.kv)
defvjp(kv, None, lambda ans, n, x: lambda g: g * (-kv(n - 1, x) - kv(n + 1, x)) / 2.0)


class SpergelMorphology(ProfileMorphology):
    """Morphology from a Spergel (2010) radial profile

    Parameters
    ----------
    frame: `~scarlet.Frame`
        Characterization of the model
    center: array or `~scarlet.Parameter`
        2D center parameter (in units of frame pixels)
        If it is to be optimized, it needs to be provided as a full defined `Parameter`.
    nu: float or `~scarlet.Parameter`
        Bessel function order.
        If it is to be optimized, it needs to be provided as a full defined `Parameter`.
    rhalf: float or `~scarlet.Parameter`
        Half-light radius in frame pixels.
        If it is to be optimized, it needs to be provided as a full defined `Parameter`.
    ellipticity: array or None
        Two-component ellipticity (e1,e2)
        If it is to be optimized, it needs to be provided as a full defined `Parameter`.
    integrate: bool
        Whether pixel integration is performed
    boxsize: int
        Size of bounding box over which to evaluate the function, in frame pixels
    """

    def __init__(self, frame, center, nu, rhalf, ellipticity=(0, 0), boxsize=None):

        assert len(center) == 2
        self.center = prepare_param(center, name="center")

        self._minimum_nu = -0.85
        self._maximum_nu = 4.00
        nu = prepare_param(nu, name="nu")
        assert nu[0] >= self._minimum_nu and nu[0] <= self._maximum_nu
        # make sure nu stays within limits
        nu.constraint = self._nu_prox

        radius = prepare_param(rhalf, name="radius")

        assert len(ellipticity) == 2
        ellipticity = prepare_param(ellipticity, name="ellipticity")
        parameters = (self.center, nu, radius, ellipticity)

        if boxsize is None:
            boxsize = int(np.ceil(10 * rhalf))

        # compute the cnu function
        # see Table 1 in Spergel (2010) and the sentence below Eqn (8).
        # _nu = np.linspace(self._minimum_nu,
        #                   self._maximum_nu,
        #                   1000)
        #
        # def func(x, nu):
        #     return (1 + nu) * self._f_nu(x, nu + 1) - 0.25
        #
        # _cnu = [fsolve(lambda x: func(x, v), x0=0.2)[0] for v in _nu]
        # self._z = np.polyfit(_nu, _cnu, 4)  # fit a 4th order polynomial
        # instead of solving this every time, we simply store the result
        self._z = np.array(
            [-0.00788962, 0.0735303, -0.27770785, 0.99483285, 1.25227402]
        )

        super().__init__(frame, self._f, *parameters, boxsize=boxsize)

    def _f(self, R2, *parameters):
        nu = self.get_parameter("nu", *parameters)
        cnu = self._cnu(nu)

        x = np.sqrt(R2 + 1e-4) * cnu
        return self._f_nu(x, nu)

    @property
    def integral(self):
        radius = self.get_parameter("radius")
        nu = self.get_parameter("nu")
        cnu = self._cnu(nu)
        return 2 * np.pi * radius ** 2 / cnu ** 2

    def _f_nu(self, x, nu):
        # Eqn 3 in Spergel (2010).
        # kv is the modified Bessel function of the second kind.
        # gamma is the gamma function.
        return (x / 2) ** nu * kv(nu, x) / gamma(nu + 1)

    def _cnu(self, nu):
        z = self._z
        return z[0] * nu ** 4 + z[1] * nu ** 3 + z[2] * nu ** 2 + z[3] * nu + z[4]

    def _nu_prox(self, x, step):
        return np.maximum(np.minimum(self._maximum_nu, x), self._minimum_nu)


class PointSourceMorphology(Morphology):
    """Morphology from a PSF

    The class uses `frame.psf` as model, evaluated at `center`

    Parameters
    ----------
    frame: `~scarlet.Frame`
        Characterization of the model
    center: array or `~scarlet.Parameter`
        2D center parameter (in units of frame pixels)
    """

    def __init__(self, frame, center):

        assert frame.psf is not None and isinstance(frame.psf, PSF)
        self.psf = frame.psf

        # define bbox
        pixel_center = tuple(np.round(center).astype("int"))
        shift = (0, *pixel_center)
        bbox = self.psf.bbox + shift

        # parameters is simply 2D center
        self.center = prepare_param(center, name="center")
        super().__init__(frame, self.center, bbox=bbox)

    def get_model(self, *parameters):
        center = self.get_parameter(0, *parameters)
        box_center = np.mean(self.bbox.bounds[1:], axis=1)
        offset = center - box_center
        return self.psf.get_model(offset=offset)  # no "internal" PSF parameters here

    @property
    def integral(self):
        # silly for Gaussian PSFs, but since the PSFs are not implemented
        # as ProfileMorphology, we can't analytically integrate
        return self.psf.get_model().sum()


class StarletMorphology(Morphology):
    """Morphology from a starlet representation of an image

    The class uses the starlet parameterization as an overcomplete, non-parametric model.

    Parameters
    ----------
    frame: `~scarlet.Frame`
        Characterization of the model
    image: 2D array
        Initial image to construct starlet transform
    bbox: `~scarlet.Box`
        2D bounding box for focation of the image in `frame`
    monotonic: bool
        Whether to constrain every starlet scale to be monotonic; otherwise they are
        hard-thresholded by `threshold`.
    threshold: float
        Lower bound on threshold for all but the last starlet scale
    """

    def __init__(self, frame, image, bbox=None, monotonic=False, threshold=0):

        if bbox is None:
            assert frame.bbox[1:].shape == image.shape
            bbox = Box(image.shape)

        self.monotonic = monotonic

        # Starlet transform of morphologies (n1,n2) with 3 dimensions: (scales+1,n1,n2)
        self.transform = Starlet.from_image(image)
        # The starlet transform is the model
        coeffs = self.transform.coefficients

        if not self.monotonic:
            # wavelet-scale norm
            starlet_norm = self.transform.norm
            # One threshold per wavelet scale: thresh*norm
            thresh_array = np.zeros(coeffs.shape) + threshold
            thresh_array *= starlet_norm[:, None, None]
            # We don't threshold the last scale
            thresh_array[-1] = 0
            constraint = ConstraintChain(
                PositivityConstraint(0), L0Constraint(thresh_array)
            )
        else:
            center = tuple(s // 2 for s in bbox.shape)
            constraint = MonotonicMaskConstraint(center, center_radius=1)

        coeffs = Parameter(coeffs, name="coeffs", step=1e-2, constraint=constraint)
        super().__init__(frame, coeffs, bbox=bbox)

    def get_model(self, *parameters):
        # Takes the inverse transform of parameters as starlet coefficients
        coeffs = self.get_parameter(0, *parameters)
        return starlet_reconstruction(coeffs)

    def update(self):
        coeffs = self.get_parameter(0)
        if coeffs.fixed:
            return

        # shrink the box?
        image = self.get_model()
        bbox = self.bbox.copy()
        self.shrink_box(image, thresh=1e-8)
        if bbox != self.bbox:
            slice, _ = overlapped_slices(bbox, self.bbox)
            center = tuple(s // 2 for s in self.bbox.shape)
            if self.monotonic:
                # non-monotonic can keep its constraint as it's independent of size
                constraint = MonotonicMaskConstraint(center, center_radius=1)
            coeffs = Parameter(
                coeffs[:, slice[0], slice[1]],
                name=coeffs.name,
                prior=coeffs.prior,
                constraint=coeffs.constraint,
                step=coeffs.step,
                fixed=coeffs.fixed,
                m=coeffs.m[:, slice[0], slice[1]] if coeffs.m is not None else None,
                v=coeffs.v[:, slice[0], slice[1]] if coeffs.v is not None else None,
                vhat=coeffs.vhat[:, slice[0], slice[1]]
                if coeffs.vhat is not None
                else None,
            )

            # set new parameters
            self._parameters = (coeffs,) + self._parameters[1:]

            raise UpdateException


class ExtendedSourceMorphology(ImageMorphology):
    def __init__(
        self,
        frame,
        center,
        image,
        bbox=None,
        monotonic="angle",
        symmetric=False,
        min_grad=0,
        shifting=False,
        resizing=True,
    ):
        """Non-parametric image morphology designed for galaxies as extended sources.

        Parameters
        ----------
        frame: `~scarlet.Frame`
            The frame of the full model
        center: tuple
            Center of the source
        image: `numpy.ndarray`
            Image of the source.
        bbox: `~scarlet.Box`
            2D bounding box for focation of the image in `frame`
        monotonic: ['flat', 'angle', 'nearest'] or None
            Which version of monotonic decrease in flux from the center to enforce
        symmetric: `bool`
            Whether or not to enforce symmetry.
        min_grad: float in [0,1)
            Minimal radial decline for monotonicity (in units of reference pixel value)
        shifting: `bool`
            Whether or not a subpixel shift is added as optimization parameter
        resize: bool
            Whether to resize the box dynamically
        """

        constraints = []
        # backwards compatibility: monotonic was boolean
        if monotonic is True:
            monotonic = "angle"
        elif monotonic is False:
            monotonic = None
        if monotonic is not None:
            # most astronomical sources are monotonically decreasing
            # from their center
            constraints.append(
                MonotonicityConstraint(neighbor_weight=monotonic, min_gradient=min_grad)
            )

        if symmetric:
            # have 2-fold rotation symmetry around their center ...
            constraints.append(SymmetryConstraint())

        constraints += [
            # ... and are positive emitters
            PositivityConstraint(),
            # prevent a weak source from disappearing entirely
            CenterOnConstraint(),
            # break degeneracies between sed and morphology
            NormalizationConstraint("max"),
        ]
        morph_constraint = ConstraintChain(*constraints)
        image = Parameter(image, name="image", step=1e-2, constraint=morph_constraint)

        self.pixel_center = np.round(center).astype("int")
        if shifting:
            shift = Parameter(center - self.pixel_center, name="shift", step=1e-1)
        else:
            shift = None
        self.shift = shift

        super().__init__(
            frame, image, bbox=bbox, shifting=shifting, shift=shift, resizing=resizing
        )

    @property
    def center(self):
        if self.shift is not None:
            return self.pixel_center + self.shift
        else:
            return self.pixel_center
