import pyfits
import numpy as np
import math
from scipy import optimize
from astropy.stats import mad_std, sigma_clip
from photutils import CircularAperture
from photutils import aperture_photometry

################################################################
##global variables
PSF_model = []

##################################################################
##2D moffat phot
def D2_moffat(A, F, x0, y0):
    B = PSF_model[0]
    C = PSF_model[1]
    D = PSF_model[2]
    E = PSF_model[3]
    return lambda y,x:A*(1 + ((x-x0)*B)**2.+((y-y0)*C)**2.+((x-x0)*(y-y0)*(D**2.)))**(-E)+F                   

def D2_moffat_phot(ROI, A, x_coo, y_coo, FWHM):
    x0=x_coo - math.floor(x_coo)+ROI.shape[1]/2.
    y0=y_coo - math.floor(y_coo)+ROI.shape[0]/2.
    
    ##recompute center of star
    params = (A, PSF_model[4], x0, y0)
    errorfunction = lambda p: np.ravel(D2_moffat(*p)(*np.indices(ROI.shape)) - ROI)
    p, success = optimize.leastsq(errorfunction, params, maxfev=1000, ftol=0.05)

    ##make mask
    grid = np.indices(ROI.shape)
    mask = np.sqrt((grid[1]-p[2])*(grid[1]-p[2]) + (grid[0]-p[3])*(grid[0]-p[3]))
##    plt.imshow(mask)
##    plt.show()
    
    ##calc sigma clipped background    
    ROI[mask < FWHM*3.0] = np.nan
    Sky = sigma_clip(ROI, sigma=2, iters=5, stdfunc=mad_std).filled(np.nan)
    NSky = np.count_nonzero(~np.isnan(Sky))
    median_Sky = np.nanmedian(Sky)
    sigma_Sky = np.nanstd(Sky)


    return (p[2]+ math.floor(x_coo)-(ROI.shape[1]/2.), p[3]+ math.floor(y_coo)-(ROI.shape[0]/2.), \
            median_Sky, NSky, sigma_Sky)

####################################################################
##search max element in 2d array
def get_max(ROI, x_coo, y_coo, R_ap):
    Max = np.amax(ROI)
    i,j = np.unravel_index(ROI.argmax(), ROI.shape) #get maximum pixel
    return (x_coo+j-R_ap, y_coo+i-R_ap, Max)

####################################################################
def PSF_phot(Data, XY, PSF, RSearch, RAper, Gain, Rnoise):
    global PSF_model
    PSF_model=PSF
    
    ##outputs empty lists
    coo=[]
    sky = []
    nsky = []
    ssky = []

    for ii in range(0, len(XY)):
        x_coo = XY[ii, 0] 
        y_coo = XY[ii, 1] 
        _coo = (0., 0.)
        _sky = np.nan        ##sky level
        _nsky = np.nan       ##number of pixel for sky
        _ssky = np.nan       ##sigma of sky

        ##check maximum level for linearity
        if (y_coo>0 and y_coo<Data.shape[0] and x_coo>0 and x_coo<Data.shape[1]):
            x_coo, y_coo, Max = get_max(Data[int(y_coo-RSearch):int(y_coo+RSearch), int(x_coo-RSearch):int(x_coo+RSearch)], x_coo, y_coo, RSearch)          #get XY coo of Max pixel in aperture
            ##copy subarray
            R3 = np.ceil(RAper*6.) ##sky annulus outer radii
            ROI =  np.copy(Data[int(y_coo-R3):int(y_coo+R3), int(x_coo-R3):int(x_coo+R3)])
            try:
                param = D2_moffat_phot(ROI, Max, x_coo, y_coo, RAper)
                _coo = (param[0], param[1])
                _sky = param[2]
                _nsky = param[3]
                _ssky = param[4]

            except:
                print ('#', ii, 'photomety failed')
                pass
                
        ##append 0 if photometry failed    
        coo.append(_coo)
        sky.append(_sky)
        nsky.append(_nsky)
        ssky.append(_ssky)
        
    coo = np.array(coo)
    coo = np.around(coo,3)
    sky = np.array(sky)
    sky = np.around(sky,3)
    nsky = np.array(nsky)
    ssky = np.around(ssky,3)
    shift = XY-coo
    
    aper = CircularAperture(coo, r=RAper)
    phot_table = aperture_photometry(Data, aper, method='exact')
    bkg_sum = sky * aper.area()
    flux = phot_table['aperture_sum'] - bkg_sum
    flux = np.array(flux)

    mag = 20. - 2.5*np.log10(flux)
    err = 1.0857*np.sqrt(flux*Gain+aper.area()*(ssky*ssky*Gain))/(flux*Gain)
    sn = 1.0857/err
    
    return coo, shift, flux, sn, sky, mag, err

################################################################################