'''
	FASTPT is a numerical algorithm to calculate
	1-loop contributions to the matter power spectrum or other
	similar type integrals.
	The method is presented in the paper 2016arXiv160304826M.
	Please cite this paper if you are using FASTPT

	J. E. McEwen (c) 2016
	mcewen.24@osu.edu

	The FASTPT class is the workhose of the FASTPT algorithm.
	This class calculates integrals of the form
	\int \frac{d^3q}{(2 \pi)^3} K(q,k-q) P(q) P(|k-q|)
'''
from __future__ import division
from __future__ import print_function


import numpy as np
from numpy.fft import fft, ifft , rfft, irfft , fftfreq
from numpy import exp, log, log10, cos, sin, pi, cosh, sinh , sqrt
from scipy.special import gamma
import sys
from time import time
from .fastpt_extr import p_window, c_window, pad_left, pad_right
from .matter_power_spt import P_13_reg
from scipy.signal import convolve , fftconvolve
from scipy.interpolate import interp1d
from .gamma_funcs import g_m_vals, gamsn
from .P_extend import k_extend
from scipy.integrate import quad

log2=log(2.)
class FASTPT:

	def __init__(self,k,nu,param_mat=None,low_extrap=None,high_extrap=None,n_pad=None,verbose=False):


		# check for log spacing
		dk=np.diff(np.log(k))
		dk_test=np.ones_like(dk)*dk[0]

		log_sample_test='ERROR! FASTPT will not work if your k vector is not sampled evenly in log space!'

		np.testing.assert_array_almost_equal(dk, dk_test, decimal=4, err_msg=log_sample_test, verbose=False)

		# size of input array must be an even number
		if (k.size % 2 != 0):
			raise ValueError('Input array must contain an even number of elements.')

		self.extrap=False
		if (low_extrap is not None or high_extrap is not None):
			self.EK=k_extend(k,low_extrap,high_extrap)
			k=self.EK.extrap_k()
			self.extrap=True

		self.low_extrap=low_extrap
		self.high_extrap=high_extrap


		self.k_old=k

		delta_L=(log(np.max(k))-log(np.min(k)))/(k.size-1)

		if(n_pad !=None):
			self.id_pad=np.arange(k.size)+n_pad
			d_logk=delta_L
			k_pad=np.log(k[0])-np.arange(1,n_pad+1)*d_logk
			k_pad=np.exp(k_pad)
			k_left=k_pad[::-1]

			k_pad=np.log(k[-1])+np.arange(1,n_pad+1)*d_logk
			k_right=np.exp(k_pad)
			k=np.hstack((k_left,k,k_right))

			# check to make sure that the n padding sufficient to keep the
			# FASTPT k_min less than 1/2 of the input k_min (added a plus 1 to be safe)
			n_pad_check=int(np.log(2)/delta_L) +1
			if (n_pad < n_pad_check):
				print('Warning, you should consider increasing your zero padding to at least ', n_pad_check, ' .')
				print('So, that you ensure that k > 2k_min.')
				print(' k min in the FASTPT universe is ', k[0], ' while k min input is ', self.k_old[0])


		if(n_pad == None):
			print('Your results are only good for k > 2k_min')


		# default parameters for standard P_22_reg
		if param_mat is None:
			param_mat=np.array([[0,0,0,0],[0,0,2,0],[0,0,4,0],[2,-2,2,0],\
							[1,-1,1,0],[1,-1,3,0],[2,-2,0,1] ])

		self.k=k
		self.k_size=k.size
		self.nu=nu
		self.p_mat=param_mat
		self.p_size=param_mat.shape[0]
		self.verbose=verbose
		self.n_pad=n_pad

		alpha=self.p_mat[:,0]
		beta=self.p_mat[:,1]
		l_Bessel=self.p_mat[:,2]
		type=self.p_mat[:,3]


		self.N=k.size

		# define eta_m and eta_n=eta_m
		omega=2*pi/(float(self.N)*delta_L)
		self.m=np.arange(-self.N//2,self.N//2+1)
		self.eta_m=omega*self.m


		# define l and tau_l
		self.n_l=self.m.size + self.m.size - 1
		self.l=np.arange(-self.n_l//2+1,self.n_l//2+1)
		self.tau_l=omega*self.l

		#Q_m=np.zeros((param_mat.shape[0],self.N+1), dtype=complex)
		self.pf=np.zeros((param_mat.shape[0]))
		self.two_part_l=np.zeros((param_mat.shape[0],self.l.size), dtype=complex)
		#Q_n=Q_m
		self.g_m=np.zeros((param_mat.shape[0],self.N+1), dtype=complex)
		self.g_n=np.zeros((param_mat.shape[0],self.N+1), dtype=complex)
		self.h_l=np.zeros((param_mat.shape[0],self.l.size),dtype=complex)

		self.p=-5-2*self.nu-alpha-beta

		for i in range(param_mat.shape[0]):

			sigma=l_Bessel[i]+1/2.

			# Define Q_m and Q_n and p
			# use eta_m for Q_n, the value is the same
			Q_m=3/2.+ nu + alpha[i] + 1j*self.eta_m
			Q_n=3/2.+ nu + beta[i] + 1j*self.eta_m
			p=-5-2*nu-alpha[i]-beta[i]

			self.g_m[i,:]=g_m_vals(sigma,Q_m)

			if (type[i]==1):

				# this is the special case, Corresponding to the regularized version of
				# J_{2,-2,0,reg}(k)
				# get values for g_n
				# again use eta_m.
				s=2+nu + beta[i]
				Q_n=s+ 1j*self.eta_m

				self.g_n[i,:]=gamsn(Q_n)

				#two_part_m=2**Q_m
				self.g_m[i,:]=self.g_m[i,:]*2.**Q_m

				# prefactor
				self.pf[i]=(-1)**l_Bessel[i]/pi**3*np.sqrt(pi/2.)

				self.two_part_l[i,:]=np.ones(self.l.size)

			else:
				self.g_n[i,:]=g_m_vals(sigma,Q_n)
				# pre factor
				self.pf[i]=(-1)**l_Bessel[i]/pi**2*2.**(2+2*nu+alpha[i]+beta[i])

				self.two_part_l[i,:]=exp(1j*self.tau_l*log2)

			# calculate h_l
			#arg=(p+1-1j*tau_l)
			self.h_l[i,:]=gamsn(self.p[i]+1-1j*self.tau_l)

	def J_k(self,P,P_window=None,C_window=None):

		if(self.low_extrap is not None):
			P=self.EK.extrap_P_low(P)

		if(self.high_extrap is not None):
			P=self.EK.extrap_P_high(P)


		P_b=P*self.k_old**(-self.nu)

		if P_window is not None:
		# window the input power spectrum, so that at high and low k
		# the signal smoothly tapers to zero. This make the input
		# more like a periodic signal

			if (self.verbose):
				print('windowing biased power spectrum')
			W=p_window(self.k_old,P_window[0],P_window[1])
			P_b=P_b*W

		if (self.n_pad !=0):
			P_b=np.pad(P_b, pad_width=(self.n_pad,self.n_pad), mode='constant', constant_values=0)

		c_m_positive=rfft(P_b)
		# End point should be divided by two. However, we always filter the Fourier coefficients (the last element is set to
		# zero), so this really has no effect.
		c_m_positive[-1]=c_m_positive[-1]/2.
		c_m_negative=np.conjugate(c_m_positive[1:])

		c_m=np.hstack((c_m_negative[::-1], c_m_positive))/float(self.N)

		if (C_window != None):
			# window the Fourier coefficients.
			# This will damping the highest frequencies

			if (self.verbose):
				print('windowing the Fourier coefficients')
			c_m=c_m*c_window(self.m,int(C_window*self.N//2.))


		A_out=np.zeros((self.p_size,self.k_size))
		for i in range(self.p_size):


			# convolve f_c and g_c
			#C_l=np.convolve(c_m*self.g_m[i,:],c_m*self.g_n[i,:])
			C_l=fftconvolve(c_m*self.g_m[i,:],c_m*self.g_n[i,:])

			# multiply all l terms together
			C_l=C_l*self.h_l[i,:]*self.two_part_l[i]

			# set up to feed ifft an array ordered with l=0,1,...,-1,...,N/2-1
			c_plus=C_l[self.l>=0]
			c_minus=C_l[self.l< 0]

			C_l=np.hstack((c_plus[:-1],c_minus))
			A_k=ifft(C_l)*C_l.size # multiply by size to get rid of the normalization in ifft

			A_out[i,:]=np.real(A_k[::2])*self.pf[i]*self.k**(-self.p[i]-2)
			# note that you have to take every other element
			# in A_k, due to the extended array created from the
			# discrete convolution

		P_out=irfft(c_m[self.m>=0])*self.k**self.nu*float(self.N)
		if (self.n_pad !=0):
			# get rid of the elements created from padding
			P_out=P_out[self.id_pad]
			A_out=A_out[:,self.id_pad]

		return P_out, A_out

	def P22(self,P,P_window=None,C_window=None):

		Power, mat=self.J_k(P,P_window=P_window,C_window=C_window)
		A=1219/1470.*mat[0,:]
		B=671/1029.*mat[1,:]
		C=32/1715.*mat[2,:]
		D=1/3.*mat[3,:]
		E=62/35.*mat[4,:]
		F=8/35.*mat[5,:]
		reg=1/3.*mat[6,:]

		return Power, 2*(A+B+C+D+E+F) + reg

	def one_loop(self,P,P_window=None,C_window=None):

	    Ps,P22=self.P22(P,P_window,C_window)
	    P13=P_13_reg(self.k_old,Ps)
	    if (self.extrap):
	        _,P=self.EK.PK_original(P22+P13)
	        return P

	    return P22+P13

	def P_bias(self,P,P_window=None,C_window=None):
		# Quadraric bias Legendre components
		# See eg section B of Baldauf+ 2012 (arxiv: 1201.4827)
		# Note pre-factor convention is not standardized
		# Returns relevant correlations (including Wick contraction factors),
		# but WITHOUT bias values and other pre-factors.
		# Uses standard "full initialization" of J terms

		Power, mat=self.J_k(P,P_window=P_window,C_window=C_window)
		sig4=np.trapz(self.k_old**2*Power**2,x=self.k_old)/(2.*pi**2)
		#sig2=np.trapz(self.k_old**2*Power,x=self.k_old)/(2.*pi**2)

		Pd1d2=2.*(17./21*mat[0,:]+mat[4,:]+4./21*mat[1,:])
		Pd2d2=2.*(mat[0,:])
		Pd1s2=2.*(8./315*mat[0,:]+4./15*mat[4,:]+254./441*mat[1,:]+2./5*mat[5,:]+16./245*mat[2,:])
		Pd2s2=2.*(2./3*mat[1,:])
		Ps2s2=2.*(4./45*mat[0,:]+8./63*mat[1,:]+8./35*mat[2,:])

		if (self.extrap):
			_, Power=self.EK.PK_original(Power)
			_, Pd1d2=self.EK.PK_original(Pd1d2)
			_, Pd2d2=self.EK.PK_original(Pd2d2)
			_, Pd1s2=self.EK.PK_original(Pd1s2)
			_, Pd2s2=self.EK.PK_original(Pd2s2)
			_, Ps2s2=self.EK.PK_original(Ps2s2)

		return Power, Pd1d2, Pd2d2, Pd1s2, Pd2s2, Ps2s2, sig4


if __name__ == "__main__":
	# An example script to run FASTPT and get the plot for
	# P_22 + P_13

	# load the data file
	#d=np.loadtxt('Pk_Planck15.dat')
	d=np.loadtxt('Pk_test.dat')
	# declare k and the power spectrum
	k=d[:,0]; P=d[:,1]

	# set the parameters for the power spectrum window and
	# Fourier coefficient window
	#P_window=np.array([.2,.2])
	C_window=.75

	# bias parameter and padding length
	nu=-2; n_pad=1000

	from time import time

	# initialize the FASTPT class
	# including extrapolation to higher and lower k

	fastpt=FASTPT(k,nu,low_extrap=-5,high_extrap=3,n_pad=n_pad)
	#fastpt=FASTPT(k,nu,n_pad=n_pad)


	t1=time()
	# without P_windowing (better if you are using zero padding)
	P_spt=fastpt.one_loop(P,C_window=C_window)
	#Ps,_,_,_,_,_,_=fastpt.P_bias(P,C_window=C_window)

	t2=time()
	print('time', "%10.3f" %(t2-t1),'s')


	# make a plot
	import matplotlib.pyplot as plt

	ax=plt.subplot(111)
	ax.set_xscale('log')
	ax.set_yscale('log')
	ax.set_ylabel(r'$P(k)$', size=30)
	ax.set_xlabel(r'$k$', size=30)

	ax.plot(k,P,label='linear')
	ax.plot(k,P_spt, label=r'$P_{22}(k) + P_{13}(k)$' )

	plt.legend(loc=3)
	plt.grid()
	plt.show()
