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

from ObservationSim.PSF.PSFModel import PSFModel

class PSFGauss(PSFModel):
	def __init__(self, chip, fwhm=0.187, sigSpin=0., psfRa=0.15):
		self.pix_size = chip.pix_scale
		self.chip = chip
		self.fwhm = fwhm
		self.sigSpin = sigSpin
		self.sigGauss = psfRa # 80% light radius
		self.psf = galsim.Gaussian(flux=1.0,fwhm=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
		"""
		gaussFun = lambda sigma, r: 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)
		frac = flux.tolist()
		return frac

	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