import galsim
import sep
import numpy as np
from scipy.interpolate import interp1d

from observation_sim.psf.PSFModel import PSFModel


class PSFGauss(PSFModel):
    def __init__(self, chip, fwhm=0.187, sigSpin=0., psfRa=None):
        self.pix_size = chip.pix_scale
        self.chip = chip
        if psfRa is None:
            self.fwhm = fwhm
            self.sigGauss = 0.15
        else:
            self.fwhm = self.fwhmGauss(r=psfRa)
            self.sigGauss = psfRa  # 80% light radius
        self.sigSpin = sigSpin
        self.psf = galsim.Gaussian(flux=1.0, fwhm=self.fwhm)

    def perfGauss(self, r, sig):
        """
        pseudo-error function, i.e. Cumulative distribution function of Gaussian distribution

        Parameter:
        r: radius
        sig: sigma of the Gaussian distribution

        Return:
        the value of the pseudo CDF
        """
        def gaussFun(sigma, r): return 1.0/(np.sqrt(2.0*np.pi)
                                            * sigma) * np.exp(-r**2/(2.0*sigma**2))
        nxx = 1000
        rArr = np.linspace(0.0, r, nxx)
        gauss = gaussFun(sig, rArr)
        erf = 2.0*np.trapz(gauss, rArr)
        return erf

    def fracGauss(self, sig, r=0.15, pscale=None):
        """
        For a given Gaussian PSF with sigma=sig,
        derive the flux ratio ar the given radius r

        Parameters:
        sig: sigma of the Gauss PSF Function in arcsec
        r:   radius in arcsec
        pscale: pixel scale

        Return: the flux ratio
        """
        if pscale == None:
            pscale = self.pix_size
        gaussx = galsim.Gaussian(flux=1.0, sigma=sig)
        gaussImg = gaussx.drawImage(scale=pscale, method='no_pixel')
        gaussImg = gaussImg.array
        size = np.size(gaussImg, axis=0)
        cxy = 0.5*(size-1)
        flux, ferr, flag = sep.sum_circle(
            gaussImg, [cxy], [cxy], [r/pscale], subpix=0)
        return flux

    def fwhmGauss(self, r=0.15, fr=0.8, pscale=None):
        """
        Given a total flux ratio 'fr' within a fixed radius 'r',
        estimate the fwhm of the Gaussian function

        return the fwhm in arcsec
        """
        if pscale == None:
            pscale = self.pix_size
        err = 1.0e-3
        nxx = 100
        sig = np.linspace(0.5*pscale, 1.0, nxx)
        frA = np.zeros(nxx)
        for i in range(nxx):
            frA[i] = self.fracGauss(sig[i], r=r, pscale=pscale)
        index = [i for i in range(nxx-1) if (fr-frA[i])
                 * (fr-frA[i+1]) <= 0.0][0]

        while abs(frA[index]-fr) > 1.0e-3:
            sig = np.linspace(sig[index], sig[index+1], nxx)
            for i in range(nxx):
                frA[i] = self.fracGauss(sig[i], r=r, pscale=pscale)
            index = [i for i in range(
                nxx-1) if (fr-frA[i])*(fr-frA[i+1]) <= 0.0][0]

        fwhm = 2.35482*sig[index]
        return fwhm

    def get_PSF(self, pos_img, chip=None, bandpass=None, folding_threshold=5.e-3):
        dx = pos_img.x - self.chip.cen_pix_x
        dy = pos_img.y - self.chip.cen_pix_y
        return self.PSFspin(dx, dy)

    def PSFspin(self, x, y):
        """
        The PSF profile at a given image position relative to the axis center

        Parameters:
        theta : spin angles in a given exposure in unit of [arcsecond]
        dx, dy: relative position to the axis center in unit of [pixels]

        Return:
        Spinned PSF: g1, g2 and axis ratio 'a/b'
        """
        a2Rad = np.pi/(60.0*60.0*180.0)

        ff = self.sigGauss * 0.107 * (1000.0/10.0)  # in unit of [pixels]
        rc = np.sqrt(x*x + y*y)
        cpix = rc*(self.sigSpin*a2Rad)

        beta = (np.arctan2(y, x) + np.pi/2)
        ell = cpix**2/(2.0*ff**2+cpix**2)
        # ell *= 10.0
        qr = np.sqrt((1.0+ell)/(1.0-ell))

        # psfShape = galsim.Shear(e=ell, beta=beta)
        # g1, g2 = psfShape.g1, psfShape.g2
        # qr = np.sqrt((1.0+ell)/(1.0-ell))

        # return ell, beta, qr
        PSFshear = galsim.Shear(e=ell, beta=beta*galsim.radians)
        return self.psf.shear(PSFshear), PSFshear