#! /usr/bin/env python

"""
Module with simplex (Nelder-Mead) optimization for defining the flux and 
position of a companion using the Negative Fake Companion.
"""
from __future__ import print_function

import numpy as np
import matplotlib.pyplot as plt
from scipy.optimize import minimize
from .simplex_fmerit import chisquare
from ..var import frame_center
from ..conf import time_ini, timing, sep

__all__ = ['firstguess_from_coord',
           'firstguess_simplex',
           'firstguess']



def firstguess_from_coord(planet, center, cube, angs, PLSC, psf, 
                          fwhm, annulus_width, aperture_radius, ncomp, 
                          cube_ref=None, svd_mode='lapack', scaling=None, 
                          fmerit='sum', collapse='median', f_range=None,
                          display=False, verbose=True, save=False, **kwargs):
    """
    Determine a first guess for the flux of a companion at a given position 
    in the cube by doing a simple grid search evaluating the reduced chi2.
    
    Parameters
    ----------
    planet: numpy.array
        The (x,y) position of the planet in the pca processed cube.
    center: numpy.array
        The (x,y) position of the cube center.
    cube: numpy.array
        The cube of fits images expressed as a numpy.array. 
    angs: numpy.array
        The parallactic angle fits image expressed as a numpy.array.         
    PLSC: float
        The platescale, in arcsec per pixel.
    psf: numpy.array
        The scaled psf expressed as a numpy.array. 
    fwhm : float
        The FHWM in pixels.           
    annulus_width: int, optional
        The width in terms of the FWHM of the annulus on which the PCA is done.       
    aperture_radius: int, optional
        The radius of the circular aperture in terms of the FWHM.
    ncomp: int
        The number of principal components. 
    cube_ref : array_like, 3d, optional
        Reference library cube. For Reference Star Differential Imaging.
    svd_mode : {'lapack', 'randsvd', 'eigen', 'arpack'}, str optional
        Switch for different ways of computing the SVD and selected PCs.
    scaling : {'temp-mean', 'temp-standard'} or None, optional
        With None, no scaling is performed on the input data before SVD. With 
        "temp-mean" then temporal px-wise mean subtraction is done and with 
        "temp-standard" temporal mean centering plus scaling to unit variance 
        is done. 
    fmerit : {'sum', 'stddev'}, string optional
        Chooses the figure of merit to be used. stddev works better for close in
        companions sitting on top of speckle noise.
    collapse : {'median', 'mean', 'sum', 'trimmean', None}, str or None, optional
        Sets the way of collapsing the frames for producing a final image. If
        None then the cube of residuals is used when measuring the function of
        merit (instead of a single final frame).
    f_range: numpy.array, optional
        The range of flux tested values. If None, 20 values between 0 and 5000
        are tested.
    display: boolean, optional
        If True, the figure chi2 vs. flux is displayed.
    verbose: boolean
        If True, display intermediate info in the shell.
    save: boolean, optional
        If True, the figure chi2 vs. flux is saved.
    kwargs: dict, optional
        Additional parameters are passed to the matplotlib plot method.
        
    Returns
    -------
    out : numpy.array
        The radial coordinates and the flux of the companion.
                
    """  
    xy = planet-center
    r0= np.sqrt(xy[0]**2+xy[1]**2)
    theta0 = np.mod(np.arctan2(xy[1],xy[0])/np.pi*180,360) 

    if f_range is not None:    
        n = f_range.shape[0]
    else:
        n = 20
        f_range =  np.linspace(0,5000,n)
    
    chi2r = []
    if verbose:
        print('Step | flux    | chi2r')
        
    counter = 0
    for j, f_guess in enumerate(f_range):
        chi2r.append(chisquare((r0,theta0,f_guess), cube, angs, PLSC, psf, 
                                fwhm, annulus_width, aperture_radius,(r0,theta0),
                                ncomp, cube_ref, svd_mode, scaling, fmerit,
                                collapse))
        if chi2r[j] > chi2r[j-1]:  counter+=1 
        if counter == 4:  break
        if verbose:
            print('{}/{}   {:.3f}   {:.3f}'.format(j+1,n,f_guess,chi2r[j]))
         
    
    chi2r = np.array(chi2r)
    f0 = f_range[chi2r.argmin()]  

    if display:
        plt.figure(figsize=kwargs.pop('figsize',(8,4)))
        plt.title(kwargs.pop('title',''))
        plt.xlim(f_range[0], f_range[:chi2r.shape[0]].max())
        plt.ylim(chi2r.min()*0.9, chi2r.max()*1.1)
        plt.plot(f_range[:chi2r.shape[0]],chi2r,
                 linestyle = kwargs.pop('linestyle','-'),
                 color = kwargs.pop('color','gray'),
                 marker = kwargs.pop('marker','.'),
                 markerfacecolor='r', markeredgecolor='r', **kwargs)
        plt.xlabel('flux')
        plt.ylabel(r'$\chi^2_{r}$')
        plt.grid('on')
    if save:
        plt.savefig('chi2rVSflux.pdf')
    if display:
        plt.show()

    return (r0,theta0,f0)       



def firstguess_simplex(p, cube, angs, psf, plsc, ncomp, fwhm, annulus_width, 
                       aperture_radius, cube_ref=None, svd_mode='lapack', 
                       scaling=None, fmerit='sum', collapse='median', p_ini=None,
                       options=None, verbose=False, **kwargs):
    """
    Determine the position of a companion using the negative fake companion 
    technique and a standard minimization algorithm (Default=Nelder-Mead) .
    
    Parameters
    ----------
    
    p : np.array
        Estimate of the candidate position.
    cube: numpy.array
        The cube of fits images expressed as a numpy.array. 
    angs: numpy.array
        The parallactic angle fits image expressed as a numpy.array. 
    psf: numpy.array
        The scaled psf expressed as a numpy.array.        
    plsc: float
        The platescale, in arcsec per pixel.
    ncomp: int
        The number of principal components.  
    fwhm : float
        The FHWM in pixels.   
    annulus_width: int, optional
        The width in terms of the FWHM of the annulus on which the PCA is done.       
    aperture_radius: int, optional