from typing import Any
import logging
import numpy as np

from .photometry import Photometry
from pyobs.images import Image

log = logging.getLogger(__name__)


class SepPhotometry(Photometry):
    """Perform photometry using SEP."""
    __module__ = 'pyobs.images.processors.photometry'

    def __init__(self, threshold: float = 1.5, minarea: int = 5, deblend_nthresh: int = 32,
                 deblend_cont: float = 0.005, clean: bool = True, clean_param: float = 1.0, **kwargs: Any):
        """Initializes a wrapper for SEP. See its documentation for details.

        Highly inspired by LCO's wrapper for SEP, see:
        https://github.com/LCOGT/banzai/blob/master/banzai/photometry.py

        Args:
            threshold: Threshold pixel value for detection.
            minarea: Minimum number of pixels required for detection.
            deblend_nthresh: Number of thresholds used for object deblending.
            deblend_cont: Minimum contrast ratio used for object deblending.
            clean: Perform cleaning?
            clean_param: Cleaning parameter (see SExtractor manual).
            *args:
            **kwargs:
        """
        Photometry.__init__(self, **kwargs)

        # store
        self.threshold = threshold
        self.minarea = minarea
        self.deblend_nthresh = deblend_nthresh
        self.deblend_cont = deblend_cont
        self.clean = clean
        self.clean_param = clean_param

    def __call__(self, image: Image) -> Image:
        """Do aperture photometry on given image.

        Args:
            image: Image to do aperture photometry on.

        Returns:
            Image with attached catalog.
        """
        import sep
        from pyobs.images.processors.detection import SepSourceDetection

        # check data
        if image.data is None:
            log.warning('No data found in image.')
            return image
        if image.catalog is None:
            log.warning('No catalog found in image.')
            return image

        # no mask?
        mask = image.mask if image.mask is not None else np.ones(image.data.shape, dtype=bool)

        # remove background
        data, bkg = SepSourceDetection.remove_background(image.data, mask)

        # fetch catalog
        sources = image.catalog.copy()

        # match SEP conventions
        x, y = sources['x'] - 1, sources['y'] - 1

        # get gain
        gain = image.header['DET-GAIN'] if 'DET-GAIN' in image.header else None

        # perform aperture photometry for diameters of 1" to 8"
        for diameter in [1, 2, 3, 4, 5, 6, 7, 8]:
            if image.pixel_scale is not None:
                flux, fluxerr, flag = sep.sum_circle(data, x, y,
                                                     diameter / 2. / image.pixel_scale,
                                                     mask=image.mask, err=image.uncertainty, gain=gain)
                sources['fluxaper{0}'.format(diameter)] = flux
                sources['fluxerr{0}'.format(diameter)] = fluxerr

            else:
                sources['fluxaper{0}'.format(diameter)] = 0
                sources['fluxerr{0}'.format(diameter)] = 0

        # average background at each source
        # since SEP sums up whole pixels, we need to do the same on an image of ones for the background_area
        bkgflux, fluxerr, flag = sep.sum_ellipse(bkg.back(), x, y,
                                                 sources['a'], sources['b'], np.pi / 2.0,
                                                 2.5 * sources['kronrad'], subpix=1)
        background_area, _, _ = sep.sum_ellipse(np.ones(shape=bkg.back().shape), x, y,
                                                sources['a'], sources['b'], np.pi / 2.0,
                                                2.5 * sources['kronrad'], subpix=1)
        sources['background'] = bkgflux
        sources['background'][background_area > 0] /= background_area[background_area > 0]

        # pick columns for catalog
        new_columns = ['fluxaper1', 'fluxerr1', 'fluxaper2', 'fluxerr2', 'fluxaper3', 'fluxerr3',
                       'fluxaper4', 'fluxerr4', 'fluxaper5', 'fluxerr5', 'fluxaper6', 'fluxerr6',
                       'fluxaper7', 'fluxerr7', 'fluxaper8', 'fluxerr8', 'background']
        cat = sources[image.catalog.colnames + new_columns]

        # copy image, set catalog and return it
        img = image.copy()
        img.catalog = cat
        return img


__all__ = ['SepPhotometry']
