from copy import deepcopy

import numpy as np

from photutils import detect_threshold
from photutils import deblend_sources
from photutils import detect_sources
from photutils.segmentation import SourceCatalog

from astropy.convolution import Gaussian2DKernel
from astropy.stats import gaussian_fwhm_to_sigma
from astropy.nddata import CCDData, Cutout2D

from matplotlib import pyplot as plt

__all__ = [
    'plot_segments', 'plot_segment_residual', 'get_source_position', 'get_source_elong',
    'get_source_ellip', 'get_source_theta', 'make_kernel', 'segm_mask', 'masked_segm_image',
    'make_segments', 'deblend_segments'
]


def plot_segments(segm, image=None, vmin=None, vmax=None, alpha=0.5, title=None):
    """
    Plot segmented areas over an image (2D array, if provided)
    """

    cmap = segm.make_cmap(seed=np.random.randint(1000000))

    if image is not None:
        plt.imshow(image, vmin=vmin, vmax=vmax, cmap="gist_gray")

    plt.imshow(segm, cmap=cmap, alpha=alpha)

    if title is not None:
        plt.title(title)

    plt.xlabel("Pixels")
    plt.ylabel("Pixels")


def plot_segment_residual(segm, image, vmin=None, vmax=None):
    """
    Plot segment subtracted image (residual)
    """
    temp = image.copy()
    temp[np.where(segm.data != 0)] = 0.0
    plt.imshow(temp, vmin=vmin, vmax=vmax)


def get_source_position(source):
    """Return max x, y value of a SourceCatalog or catalog row"""
    if isinstance(source, SourceCatalog):
        x, y = source.maxval_xindex, source.maxval_yindex
    else:
        x, y = source['maxval_xindex'], source['maxval_yindex']
    return x, y


def get_source_elong(source):
    """ Return SourceCatalog elongation"""
    return source.elongation.value if isinstance(source, SourceCatalog) else source['elongation']


def get_source_ellip(source):
    """ Return SourceCatalog ellipticity"""
    return source.ellipticity.value if isinstance(source, SourceCatalog) else source['ellipticity']


def get_source_theta(source):
    """ Return SourceCatalog orientation in rad"""
    return source.orientation.to('rad').value if isinstance(source, SourceCatalog) else np.deg2rad(source['orientation'])


def make_kernel(fwhm, kernel_size):
    sigma = fwhm * gaussian_fwhm_to_sigma
    kernel = Gaussian2DKernel(sigma, x_size=kernel_size, y_size=kernel_size)
    kernel.normalize()
    return kernel


def segm_mask(source, segm, mask_background=False):
    """
    Given a segmentation and a target with an label, returns a mask
    with all other sources masked in the original image.

    Parameters
    ----------

    source : int or photutils.segmentation.properties.SourceCatalog
        The catalog object for the target or label of target in the segmentation object.

    segm : photutils.segmentation.core.SegmentationImage
        Segmentation image.

    mask_background : bool
        Option to also mask out all the un-segmented background pixels.

    Returns
    -------

    mask : bool array
    """

    if isinstance(source, SourceCatalog):
        source = source.label

    mask = (segm.data == source)
    if not mask_background:
        mask = ((segm.data == 0) | mask)
    return mask


def masked_segm_image(source, image, segm, fill=None, mask_background=False):
    """
    Returns a masked image of the original image by masking out other sources

    Parameters
    ----------

    source : int or photutils.segmentation.properties.SourceCatalog
        The catalog object for the target or label of target in the segmentation object.

    image : CCDData or array
        Image to mask.

    segm : photutils.segmentation.core.SegmentationImage
        Segmentation image.

    fill : float
        Fill in the masked pixels with this value.

    mask_background : bool
        Option to also mask out all the un-segmented background pixels.

    Returns
    -------

    masked_image : CCDData or array
    """

    fill = np.nan if fill is None else fill
    mask = segm_mask(source, segm, mask_background)

    masked_image = deepcopy(image)

    if isinstance(masked_image, np.ndarray):
        masked_image[np.invert(mask)] = fill
    else:
        masked_image.data[np.invert(mask)] = fill

    return masked_image


def make_segments(image, npixels=None, threshold=3., fwhm=8., kernel_size=4):
    """
    Segment an image.

    Parameters
    ----------
    image : array like
        Input image

    npixels : int
        The number of connected pixels, each greater than ``threshold``,
        that an object must have to be detected.  ``npixels`` must be a
        positive integer.

    threshold : float or array-like
        The data value or pixel-wise data values to be used for the
        detection threshold. A 2D ``threshold`` must have the same shape
        as ``data``. See `~photutils.segmentation.detect_threshold` for
        one way to create a ``threshold`` image.

    fwhm : float
        FWHM of smoothing gaussian kernel.

    kernel_size : int
        Size of smoothing kernel.

    Returns
    -------

    segment_image : `~photutils.segmentation.SegmentationImage` or `None`
        A 2D segmentation image, with the same shape as ``data``, where
        sources are marked by different positive integer values.  A
        value of zero is reserved for the background.  If no sources
        are found then `None` is returned.
    """

    if npixels is None:
        npixels = fwhm ** 2

    kernel = make_kernel(fwhm, kernel_size) if kernel_size else None

    return detect_sources(image, threshold, npixels=npixels, filter_kernel=kernel)


def deblend_segments(image, segm, npixels=None, fwhm=8., kernel_size=4, nlevels=30, contrast=1/1000):
    """
    Deblend overlapping sources labeled in a segmentation image.

    Parameters
    ----------
    image : array like
        Input image.

    segm : `~photutils.segmentation.SegmentationImage` or `None`
        A 2D segmentation image, with the same shape as ``data``, where
        sources are marked by different positive integer values.  A
        value of zero is reserved for the background.  If no sources
        are found then `None` is returned.

    npixels : int
        The number of connected pixels, each greater than ``threshold``,
        that an object must have to be detected.  ``npixels`` must be a
        positive integer.

    fwhm : float
        FWHM of smoothing gaussian kernel.

    kernel_size : int
        Size of smoothing kernel.

    nlevels : int, optional
        The number of multi-thresholding levels to use.  Each source
        will be re-thresholded at ``nlevels`` levels spaced
        exponentially or linearly (see the ``mode`` keyword) between its
        minimum and maximum values within the source segment.

    contrast : float, optional
        The fraction of the total (blended) source flux that a local
        peak must have (at any one of the multi-thresholds) to be
        considered as a separate object.  ``contrast`` must be between 0
        and 1, inclusive.  If ``contrast = 0`` then every local peak
        will be made a separate object (maximum deblending).  If
        ``contrast = 1`` then no deblending will occur.  The default is
        0.001, which will deblend sources with a 7.5 magnitude
        difference.

    Returns
    -------
    segment_image : `~photutils.segmentation.SegmentationImage`
        A segmentation image, with the same shape as ``data``, where
        sources are marked by different positive integer values.  A
        value of zero is reserved for the background.
    """

    if npixels is None:
        npixels = fwhm ** 2

    kernel = make_kernel(fwhm, kernel_size) if kernel_size else None

    segm_deblend = deblend_sources(image, segm,
                                   npixels=npixels, filter_kernel=kernel,
                                   nlevels=nlevels, contrast=contrast)

    return segm_deblend


def make_catalog(image, threshold, wcs=None, deblend=True,
                 npixels=None, fwhm=8., kernel_size=4,
                 nlevels=30, contrast=1/1000,
                 plot=True, vmax=None, vmin=None):
    """
    This function constructs a catalog using `PhotUtils`. The `petrofit.segmentation.make_segments` and
    `petrofit.segmentation.deblend_segments` functions are used to construct segmentation maps and
    the resulting segmentation is turned into a vatalog (`photutils.segmentation.catalog.SourceCatalog`)

    Parameters
    ----------
    image : array like
        Input image.

    threshold : float or array-like
        The data value or pixel-wise data values to be used for the
        detection threshold. A 2D ``threshold`` must have the same shape
        as ``data``. See `~photutils.segmentation.detect_threshold` for
        one way to create a ``threshold`` image.

    wcs : astropy.wcs.WCS
        Astropy WCS for catalog to convert source locations to RA-Dec.

    deblend : bool
        Flag to enable deblending of sources in the segmentation map.

    npixels : int
        The number of connected pixels, each greater than ``threshold``,
        that an object must have to be detected.  ``npixels`` must be a
        positive integer.

    fwhm : float
        FWHM of smoothing gaussian kernel.

    kernel_size : int
        Size of smoothing kernel.

    nlevels : int, optional
        The number of multi-thresholding levels to use.  Each source
        will be re-thresholded at ``nlevels`` levels spaced
        exponentially or linearly (see the ``mode`` keyword) between its
        minimum and maximum values within the source segment.

    contrast : float, optional
        The fraction of the total (blended) source flux that a local
        peak must have (at any one of the multi-thresholds) to be
        considered as a separate object.  ``contrast`` must be between 0
        and 1, inclusive.  If ``contrast = 0`` then every local peak
        will be made a separate object (maximum deblending).  If
        ``contrast = 1`` then no deblending will occur.  The default is
        0.001, which will deblend sources with a 7.5 magnitude
        difference.

    plot : bool
        Flag to toggle plotting segmentation maps.

    vmin, vmax : float
        vmax and vmin values for plot.

    Returns
    -------
    cat, segm, segm_deblend
        cat : `photutils.segmentation.catalog.SourceCatalog`
            A catalog of sources

        segm : `photutils.segmentation.core.SegmentationImage`
            Segmentation map.

        segm_deblend : `photutils.segmentation.core.SegmentationImage`
            Deblended segmentation map.
    """

    if isinstance(image, CCDData) or isinstance(image, Cutout2D):
        image = image.data

    # Make segmentation map
    segm = make_segments(image,
                         threshold=threshold,
                         kernel_size=kernel_size,
                         fwhm=fwhm, npixels=npixels)

    if plot and segm:
        # Make plots
        if deblend:
            fig, ax = plt.subplots(1, 2)
            plt.sca(ax[0])
        plot_segments(segm, image=image, vmax=vmax, vmin=vmin, title='Segmentation Map')

    # Deblend segmentation map
    segm_deblend = None
    if deblend:
        segm_deblend = deblend_segments(image, segm,
                                        contrast=contrast,
                                        nlevels=nlevels,
                                        kernel_size=kernel_size,
                                        fwhm=fwhm, npixels=npixels)

        if plot and segm_deblend:
            # Make plots
            plt.sca(ax[1])
            plot_segments(segm_deblend, image=image, vmax=vmax, vmin=vmin, title='Deblended Segmentation Map')

    # Make catalog
    cat = SourceCatalog(image, segm_deblend if deblend else segm, wcs=wcs)

    return cat, segm, segm_deblend