#! /usr/bin/env python

"""
Module with fake companion injection functions.
"""

from __future__ import division
from __future__ import print_function

__author__ = 'Carlos Alberto Gomez Gonzalez'
__all__ = ['create_psf_template',
           'psf_norm',
           'cube_inject_companions',
           'frame_inject_companion',
           'inject_fcs_cube',
           'inject_fc_frame']

import numpy as np
import photutils
from ..preproc import cube_crop_frames, frame_shift, frame_crop
from ..var import frame_center, fit_2dgaussian, get_circle


# TODO: remove this in VIP v1.0.0. Created for backward compatibility
def inject_fcs_cube(array, psf_template, angle_list, flevel, plsc,
                    rad_dists, n_branches=1, theta=0, imlib='opencv',
                    verbose=True):
    return cube_inject_companions(array, psf_template, angle_list, flevel, plsc,
                                  rad_dists, n_branches, theta, imlib, verbose)


# TODO: remove this in VIP v1.0.0. Created for backward compatibility
def inject_fc_frame(array, array_fc, pos_y, pos_x, flux):
    return frame_inject_companion(array, array_fc, pos_y, pos_x, flux)


def cube_inject_companions(array, psf_template, angle_list, flevel, plsc,
                           rad_dists, n_branches=1, theta=0, imlib='opencv',
                           interpolation='lanczos4', verbose=True):
    """ Injects fake companions in branches, at given radial distances.

    Parameters
    ----------
    array : array_like
        Input frame or cube.
    psf_template : array_like
        2d array with the normalized psf template. It should have an odd shape.
        It's recommended to run the function psf_norm to get a proper PSF
        template. In the ADI+IFS case it must be a 3d array.
    flevel : float or list
        Factor for controlling the brightness of the fake companions.
    plsc : float
        Value of the plsc in pixels.
    rad_dists : list or array 1d
        Vector of radial distances of fake companions in pixels.
    n_branches : int, optional
        Number of azimutal branches.
    theta : float, optional
        Angle in degrees for rotating the position of the first branch that by
        default is located at zero degrees. Theta counts counterclockwise from
        the positive x axis.
    imlib : str, optional
        See the documentation of the ``vip_hci.preproc.frame_shift`` function.
    interpolation : str, optional
        See the documentation of the ``vip_hci.preproc.frame_shift`` function.
    verbose : {True, False}, bool optional
        If True prints out additional information.

    Returns
    -------
    array_out : array_like
        Output array with the injected fake companions.

    """

    if not (array.ndim == 3 or array.ndim == 4):
        raise ValueError('Array is not a cube, 3d or 4d array')
    if array.ndim == 4:
        if not psf_template.ndim == 3:
            raise ValueError('Psf must be a 3d array')

    # ADI case
    if array.ndim == 3:
        # TODO: extend to even-sized psf arrays
        if psf_template.shape[1]%2 == 0:
            raise ValueError("Only odd-sized PSF is accepted")
        ceny, cenx = frame_center(array[0])
        ceny = int(float(ceny))
        cenx = int(float(cenx))
        rad_dists = np.array(rad_dists)
        if not rad_dists[-1] < array[0].shape[0]/2.:
            msg = 'rad_dists last location is at the border (or outside) of the field'
            raise ValueError(msg)

        size_fc = psf_template.shape[0]
        nframes = array.shape[0]
        fc_fr = np.zeros_like(array[0])
        n_fc_rad = rad_dists.shape[0]

        w = int(np.floor(size_fc/2.))
        # fcomp in the center of a zeros frame
        fc_fr[ceny-w:ceny+w+1, cenx-w:cenx+w+1] = psf_template

        array_out = np.zeros_like(array)
        for fr in range(nframes):
            tmp = np.zeros_like(array[0])
            for branch in range(n_branches):
                ang = (branch * 2 * np.pi / n_branches) + np.deg2rad(theta)
                for i in range(n_fc_rad):
                    rad = rad_dists[i]
                    y = rad * np.sin(ang - np.deg2rad(angle_list[fr]))
                    x = rad * np.cos(ang - np.deg2rad(angle_list[fr]))
                    tmp += frame_shift(fc_fr, y, x, imlib=imlib,
                                       interpolation=interpolation) * flevel
                    array_out[fr] = array[fr] + tmp

        if verbose:
            for branch in range(n_branches):
                print('Branch '+str(branch+1)+':')
                for i in range(n_fc_rad):
                    ang = (branch * 2 * np.pi / n_branches) + np.deg2rad(theta)
                    posy = rad_dists[i] * np.sin(ang) + ceny
                    posx = rad_dists[i] * np.cos(ang) + cenx
                    rad_arcs = rad_dists[i]*plsc
                    msg ='\t(X,Y)=({:.2f}, {:.2f}) at {:.2f} arcsec ({:.2f} pxs)'
                    print(msg.format(posx, posy, rad_arcs, rad_dists[i]))

        return array_out

    # ADI+IFS case
    if array.ndim == 4 and psf_template.ndim == 3:
        # TODO: extend to even-sized psf arrays
        if psf_template.shape[2] % 2 == 0:
            raise ValueError("Only odd-sized PSF is accepted")
        ceny, cenx = frame_center(array[0,0])
        ceny = int(float(ceny))
        cenx = int(float(cenx))
        rad_dists = np.array(rad_dists)
        if not rad_dists[-1]<array[0].shape[1]/2.:
            msg = 'rad_dists last location is at the border (or outside) of the field'
            raise ValueError(msg)

        sizey = array.shape[2]
        sizex = array.shape[3]
        size_fc = psf_template.shape[2] # considering square frames
        nframes_wav = array.shape[0]
        nframes_adi = array.shape[1]
        fc_fr = np.zeros((nframes_wav, sizey, sizex), dtype=np.float64) # now a 3d cube
        n_fc_rad = rad_dists.shape[0]

        for i in range(nframes_wav):
            w = int(np.floor(size_fc/2.))
            # fcomp in the center of a zeros frame
            if (psf_template[0].shape[1]%2) == 0:
                fc_fr[i, ceny-w:ceny+w, cenx-w:cenx+w] = psf_template[i]
            else:
                fc_fr[i, ceny-w:ceny+w+1, cenx-w:cenx+w+1] = psf_template[i]

        array_out = np.zeros_like(array)

        for fr in range(nframes_adi):
            tmp = np.zeros_like(fc_fr)
            for branch in range(n_branches):
                ang = (branch * 2 * np.pi / n_branches) + np.deg2rad(theta)
                for i in range(n_fc_rad):
                    rad = rad_dists[i]
                    y = rad * np.sin(ang - np.deg2rad(angle_list[fr]))
                    x = rad * np.cos(ang - np.deg2rad(angle_list[fr]))
                    if isinstance(flevel, int) or isinstance(flevel, float):
                        tmp += _cube_shift(fc_fr, y, x, imlib=imlib,
                                           interpolation=interpolation) * flevel
                    else:
                        shift = _cube_shift(fc_fr, y, x, imlib=imlib,
                                            interpolation=interpolation)
                        tmp += [shift[i]*flevel[i] for i in range(len(flevel))]
            array_out[:,fr] = array[:,fr] + tmp

        if verbose:
            for branch in range(n_branches):
                print('Branch '+str(branch+1)+':')
                for i in range(n_fc_rad):
                    ang = (branch * 2 * np.pi / n_branches) + np.deg2rad(theta)
                    rad_arcs = rad_dists[i]*plsc
                    msg ='\t(X,Y)=({:.2f}, {:.2f}) at {:.2f} arcsec ({:.2f} pxs)'
                    posy = rad_dists[i] * np.sin(ang) + ceny
                    posx = rad_dists[i] * np.cos(ang) + cenx
                    print(msg.format(posx, posy, rad_arcs, rad_dists[i]))

    return array_out


def frame_inject_companion(array, array_fc, pos_y, pos_x, flux,
                           imlib='opencv', interpolation='lanczos4'):
    """ Injects a fake companion in a single frame (it could be a single
     multi-wavelength frame) at given coordinates.
    """
    if not (array.ndim == 2 or array.ndim == 3):
        raise TypeError('Array is not a 2d or 3d array.')
    if array.ndim == 2:
        size_fc = array_fc.shape[0]
        ceny, cenx = frame_center(array)
        ceny = int(ceny)
        cenx = int(cenx)
        fc_fr = np.zeros_like(array)
        w = int(np.floor(size_fc/2.))
        # fcomp in the center of a zeros frame
        fc_fr[ceny-w:ceny+w+1, cenx-w:cenx+w+1] = array_fc
        array_out = array + frame_shift(fc_fr, pos_y-ceny, pos_x-cenx, imlib,
                                        interpolation) * flux

    if array.ndim == 3:
        size_fc = array_fc.shape[1]
        ceny, cenx = frame_center(array[0])
        ceny = int(ceny)
        cenx = int(cenx)
        fc_fr = np.zeros_like(array)
        w = int(np.floor(size_fc/2.))
        # fcomp in the center of a zeros frame
        fc_fr[:, ceny-w:ceny+w+1, cenx-w:cenx+w+1] = array_fc
        array_out = array + _cube_shift(fc_fr, pos_y - ceny, pos_x - cenx,
                                        imlib, interpolation) * flux

    return array_out


def create_psf_template(array, size, fwhm=4, verbose=True, collapse='mean'):
    """ Creates a psf template from a cube of non-saturated off-axis frames of
    the star by taking the mean and normalizing the psf flux.

    Parameters
    ----------
    array : array_like, 3d
        Input cube.
    size : int
        Size of the squared subimage.
    fwhm: float, optional
        The size of the Full Width Half Maximum in pixel.
    verbose : {True,False}, bool optional
        Whether to print to stdout information about file opening, cropping and
        completion of the psf template.
    collapse : {'mean','median'}, string optional
        Defines the way the frames are collapsed.

    Returns
    -------
    psf_normd : array_like
        Normalized PSF.
    """
    if not (array.ndim == 3 or array.ndim == 4):
        raise TypeError('Array is not a cube, 3d or 4d array.')

    n = array.shape[0]
    psf = cube_crop_frames(array, size=size, verbose=verbose)
    if collapse == 'mean':
        psf = np.mean(psf, axis=0)
    elif collapse == 'median':
        psf = np.median(psf, axis=0)
    else:
        raise TypeError('Collapse mode not recognized.')

    psf_normd = psf_norm(psf, size=size, fwhm=fwhm)

    if verbose:
        print("Done scaled PSF template from the average of", n,"frames.")
    return psf_normd


def psf_norm(array, fwhm=4, size=None, threshold=None, mask_core=None,
             imlib='opencv', interpolation='lanczos4', full_output=False,
             verbose=False):
    """ Scales a PSF (2d or 3d array), so the 1*FWHM aperture flux equals 1.

    Parameters
    ----------
    array: array_like
        The PSF, 2d (ADI data) or 3d array (IFS data).
    fwhm: int or float or 1d array, optional
        The the Full Width Half Maximum in pixels. It can handle a different
        FWHM value for different wavelengths (IFS data).
    size : int or None, optional
        If int it will correspond to the size of the squared subimage to be
        cropped form the psf array.
    threshold : None of float, optional
        Sets to zero small values, trying to leave only the core of the PSF.
    mask_core : None of float, optional
        Sets the radius of a circular aperture for the core of the PSF,
        everything else will be set to zero.
    imlib : str, optional
        See the documentation of the ``vip_hci.preproc.frame_shift`` function.
    interpolation : str, optional
        See the documentation of the ``vip_hci.preproc.frame_shift`` function.
    full_output : bool, optional
        If True the flux in a FWHM aperture is returned along with the
        normalized PSF.
    verbose : bool, optional
        If True intermediate results are printed out.

    Returns
    -------
    psf_norm: array_like
        The normalized psf.

    If ``full_output`` is True the flux in a FWHM aperture is returned along
    with the normalized PSF.
    """
    def psf_norm_2d(array, fwhm, size, threshold, mask_core, imlib,
                    interpolation, full_output, verbose):
        """ Helper function for 2d case """
        if size is not None:
            if size < array.shape[0]:
                psfs = frame_crop(array, size, verbose=False)
            else:
                psfs = array.copy()
        else:
            psfs = array.copy()
            # TODO: verify correct handling of even/odd cases (old behavior: if
            # frame size is even we drop last row and last column)
            if psfs.shape[0]%2 == 0:
                psfs = psfs[:-1, :]
            if psfs.shape[1]%2 == 0:
                psfs = psfs[:, :-1]

        # we check if the psf is centered and fix it if needed
        cy, cx = frame_center(psfs, verbose=False)
        maxpxy = np.where(psfs == psfs.max())[0]
        maxpxx = np.where(psfs == psfs.max())[1]
        if cy != maxpxy or cx != maxpxx:
            # first we find the centroid and put it in the center of the array
            centroidy, centroidx = fit_2dgaussian(psfs, fwhmx=fwhm, fwhmy=fwhm)
            shiftx, shifty = centroidx - cx, centroidy - cy
            psfs = frame_shift(psfs, -shifty, -shiftx, imlib=imlib,
                               interpolation=interpolation)
            for _ in range(2):
                centroidy, centroidx = fit_2dgaussian(psfs, fwhmx=fwhm,
                                                      fwhmy=fwhm)
                cy, cx = frame_center(psfs, verbose=False)
                shiftx, shifty = centroidx - cx, centroidy - cy
                psfs = frame_shift(psfs, -shifty, -shiftx)

        # we check whether the flux is normalized and fix it if needed
        fwhm_aper = photutils.CircularAperture((frame_center(psfs)), fwhm/2.)
        fwhm_aper_phot = photutils.aperture_photometry(psfs, fwhm_aper,
                                                       method='exact')
        fwhm_flux = np.array(fwhm_aper_phot['aperture_sum'])
        if verbose:
            print("Flux in 1xFWHM aperture: {}".format(fwhm_flux))

        if fwhm_flux > 1.1 or fwhm_flux < 0.9:
            psf_norm_array = psfs / np.array(fwhm_aper_phot['aperture_sum'])
        else:
            psf_norm_array = psfs

        if threshold is not None:
            psf_norm_array[np.where(psf_norm_array < threshold)] = 0

        if mask_core is not None:
            psf_norm_array = get_circle(psf_norm_array, radius=mask_core)

        if full_output:
            return psf_norm_array, fwhm_flux
        else:
            return psf_norm_array
    #---------------------------------------------------------------------------

    if array.ndim == 2:
        res = psf_norm_2d(array, fwhm, size, threshold, mask_core, imlib,
                          interpolation, full_output, verbose)

    elif array.ndim == 3:
        if isinstance(fwhm, (int, float)):
            fwhm = [fwhm for _ in range(array.shape[0])]
        array_out = np.zeros((array.shape[0], size, size))
        if full_output:
            fwhm_flux = np.zeros((array.shape[0]))

        for fr in range(array.shape[0]):
            restemp = psf_norm_2d(array[fr], fwhm[fr], size, threshold,
                                  mask_core, imlib, interpolation, full_output,
                                  False)
            if full_output:
                array_out[fr] = restemp[0]
                fwhm_flux[fr] = restemp[1]
            else:
                array_out[fr] = restemp

        if full_output:
            res = (array_out, fwhm_flux)
        else:
            res = array_out

    return res


def _cube_shift(cube, y, x, imlib, interpolation):
    " Shifts the X-Y coordinates of a cube or 3D array by x and y values"
    cube_out = np.zeros_like(cube)
    for i in range(cube.shape[0]):
        cube_out[i] = frame_shift(cube[i], y, x, imlib, interpolation)
    return cube_out
