## Copyright 2021-2023 Ian Jauslin
## 
## Licensed under the Apache License, Version 2.0 (the "License");
## you may not use this file except in compliance with the License.
## You may obtain a copy of the License at
## 
##     http://www.apache.org/licenses/LICENSE-2.0
## 
## Unless required by applicable law or agreed to in writing, software
## distributed under the License is distributed on an "AS IS" BASIS,
## WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
## See the License for the specific language governing permissions and
## limitations under the License.

# Chebyshev expansion
@everywhere function chebyshev(
  a::Array{Float64,1},
  taus::Array{Float64,1},
  weights::Tuple{Array{Float64,1},Array{Float64,1}},
  P::Int64,
  N::Int64,
  J::Int64,
  nu::Int64
)
  out=zeros(Float64,J*(P+1))
  for zeta in 0:J-1
    for n in 0:P
      for j in 1:N
	out[zeta*(P+1)+n+1]+=(2-(n==0 ? 1 : 0))/2*weights[2][j]*cos(n*pi*(1+weights[1][j])/2)*a[zeta*N+j]/(1-(taus[zeta+2]-taus[zeta+1])/2*sin(pi*weights[1][j]/2)+(taus[zeta+2]+taus[zeta+1])/2)^nu
      end
    end
  end

  return out
end

# evaluate function from Chebyshev expansion
@everywhere function chebyshev_eval(
  Fa::Array{Float64,1},
  x::Float64,
  taus::Array{Float64,1},
  chebyshev::Array{Polynomial,1},
  P::Int64,
  J::Int64,
  nu::Int64
)
  # change variable
  tau=(1-x)/(1+x)

  out=0.
  for zeta in 0:J-1
    # check that the spline is right
    if tau<taus[zeta+2] && tau>=taus[zeta+1]
      for n in 0:P
	out+=Fa[zeta*(P+1)+n+1]*chebyshev[n+1]((2*tau-(taus[zeta+1]+taus[zeta+2]))/(taus[zeta+2]-taus[zeta+1]))
      end
      break
    end
  end

  return (1+tau)^nu*out
end


# convolution 
# input the Chebyshev expansion of a and b, as well as the A matrix
@everywhere function conv_chebyshev(
  Fa::Array{Float64,1},
  Fb::Array{Float64,1},
  A::Array{Array{Float64,2},1}
)
  out=zeros(Float64,length(A))
  for i in 1:length(A)
    out[i]=dot(Fa,A[i]*Fb)
  end
  return out
end

# <ab>
@everywhere function avg_chebyshev(
  Fa::Array{Float64,1},
  Fb::Array{Float64,1},
  A0::Float64
)
  return dot(Fa,A0*Fb)
end

# 1_n * a
@everywhere function conv_one_chebyshev(
  n::Int64,
  zetapp::Int64,
  Fa::Array{Float64,1},
  A::Array{Array{Float64,2},1},
  taus::Array{Float64,1},
  weights::Tuple{Array{Float64,1},Array{Float64,1}},
  P::Int64,
  N::Int64,
  J::Int64,
  nu1::Int64
)
  out=zeros(Float64,N*J)
  for m in 1:N*J
    for l in 0:P
      for p in 1:J*(P+1)
	out[m]+=(2-(l==0 ? 1 : 0))/2*weights[2][n]*cos(l*pi*(1+weights[1][n])/2)/(1-(taus[zetapp+2]-taus[zetapp+1])/2*sin(pi*weights[1][n]/2)+(taus[zetapp+2]+taus[zetapp+1])/2)^nu1*A[m][zetapp*(P+1)+l+1,p]*Fa[p]
      end
    end
  end
  return out
end
# a * v
@everywhere function conv_v_chebyshev(
  a::Array{Float64,1},
  Upsilon::Array{Array{Float64,1},1},
  k::Array{Float64,1},
  taus::Array{Float64,1},
  weights::Tuple{Array{Float64,1},Array{Float64,1}},
  N::Int64,
  J::Int64
)
  out=zeros(Float64,J*N)
  for i in 1:J*N
    for zetap in 0:J-1
      for j in 1:N
	xj=weights[1][j]
	out[i]+=(taus[zetap+2]-taus[zetap+1])/(32*pi)*weights[2][j]*cos(pi*xj/2)*(1+k[zetap*N+j])^2*k[zetap*N+j]*a[zetap*N+j]*Upsilon[zetap*N+j][i]
      end
    end
  end
  return out
end
@everywhere function conv_one_v_chebyshev(
  n::Int64,
  zetap::Int64,
  Upsilon::Array{Array{Float64,1},1},
  k::Array{Float64,1},
  taus::Array{Float64,1},
  weights::Tuple{Array{Float64,1},Array{Float64,1}},
  N::Int64,
  J::Int64
)
  out=zeros(Float64,J*N)
  xj=weights[1][n]
  for i in 1:J*N
    out[i]=(taus[zetap+2]-taus[zetap+1])/(32*pi)*weights[2][n]*cos(pi*xj/2)*(1+k[zetap*N+n])^2*k[zetap*N+n]*Upsilon[zetap*N+n][i]
  end
  return out
end

# <av>
@everywhere function avg_v_chebyshev(a,
  Upsilon0::Array{Float64,1},
  k::Array{Float64,1},
  taus::Array{Float64,1},
  weights::Tuple{Array{Float64,1},Array{Float64,1}},
  N::Int64,
  J::Int64
)
  out=0.
  for zetap in 0:J-1
    for j in 1:N
      xj=weights[1][j]
      out+=(taus[zetap+2]-taus[zetap+1])/(32*pi)*weights[2][j]*cos(pi*xj/2)*(1+k[zetap*N+j])^2*k[zetap*N+j]*a[zetap*N+j]*Upsilon0[zetap*N+j]
    end
  end
  return out
end
# <1_nv>
@everywhere function avg_one_v_chebyshev(
  n::Int64,
  zetap::Int64,
  Upsilon0::Array{Float64,1},
  k::Array{Float64,1},
  taus::Array{Float64,1},
  weights::Tuple{Array{Float64,1},Array{Float64,1}},
  N::Int64
)
  xj=weights[1][n]
  return (taus[zetap+2]-taus[zetap+1])/(32*pi)*weights[2][n]*cos(pi*xj/2)*(1+k[zetap*N+n])^2*k[zetap*N+n]*Upsilon0[zetap*N+n]
end

# compute \int dq dxi u1(k-xi)u2(q)u3(xi)u4(k-q)u5(xi-q)
@everywhere function double_conv_S_chebyshev(
  FU1::Array{Float64,1},
  FU2::Array{Float64,1},
  FU3::Array{Float64,1},
  FU4::Array{Float64,1},
  FU5::Array{Float64,1},
  Abar::Array{Float64,5}
)
  out=zeros(Float64,length(Abar))
  for i in 1:length(Abar)
    for j1 in 1:length(FU1)
      for j2 in 1:length(FU2)
	for j3 in 1:length(FU3)
	  for j4 in 1:length(FU4)
	    for j5 in 1:length(FU5)
	    out[i]+=Abar[i][j1,j2,j3,j4,j5]*FU1[j1]*FU2[j2]*FU3[j3]*FU4[j4]*FU5[j5]
	    end
	  end
	end
      end
    end
  end
  return out
end


# compute A
@everywhere function Amat(
  k::Array{Float64,1},
  weights::Tuple{Array{Float64,1},Array{Float64,1}},
  taus::Array{Float64,1},
  T::Array{Polynomial,1},
  P::Int64,
  N::Int64,
  J::Int64,
  nua::Int64,
  nub::Int64
)
  out=Array{Array{Float64,2},1}(undef,J*N)
  for i in 1:J*N
    out[i]=zeros(Float64,J*(P+1),J*(P+1))
    for zeta in 0:J-1
      for n in 0:P
	for zetap in 0:J-1
	  for m in 0:P
	    out[i][zeta*(P+1)+n+1,zetap*(P+1)+m+1]=1/(pi^2*k[i])*integrate_legendre(tau->(1-tau)/(1+tau)^(3-nua)*T[n+1]((2*tau-(taus[zeta+1]+taus[zeta+2]))/(taus[zeta+2]-taus[zeta+1]))*(alpham(k[i],tau)>taus[zetap+2] || alphap(k[i],tau)<taus[zetap+1] ? 0. : integrate_legendre(sigma->(1-sigma)/(1+sigma)^(3-nub)*T[m+1]((2*sigma-(taus[zetap+1]+taus[zetap+2]))/(taus[zetap+2]-taus[zetap+1])),max(taus[zetap+1],alpham(k[i],tau)),min(taus[zetap+2],alphap(k[i],tau)),weights)),taus[zeta+1],taus[zeta+2],weights)
	  end
	end
      end
    end
  end

  return out
end

# compute Upsilon
@everywhere function Upsilonmat(
  k::Array{Float64,1},
  v::Function,
  weights::Tuple{Array{Float64,1},Array{Float64,1}}
)
  out=Array{Array{Float64,1},1}(undef,length(k))
  for i in 1:length(k)
    out[i]=Array{Float64,1}(undef,length(k))
    for j in 1:length(k)
      out[i][j]=integrate_legendre(s->s*v(s)/k[j],abs(k[j]-k[i]),k[j]+k[i],weights)
    end
  end
  return out
end
@everywhere function Upsilon0mat(
  k::Array{Float64,1},
  v::Function,
  weights::Tuple{Array{Float64,1},Array{Float64,1}}
)
  out=Array{Float64,1}(undef,length(k))
  for j in 1:length(k)
    out[j]=2*k[j]*v(k[j])
  end
  return out
end

# alpha_-
@everywhere function alpham(
  k::Float64,
  t::Float64
)
  return (1-k-(1-t)/(1+t))/(1+k+(1-t)/(1+t))
end
# alpha_+
@everywhere function alphap(
  k::Float64,
  t::Float64
)
  return (1-abs(k-(1-t)/(1+t)))/(1+abs(k-(1-t)/(1+t)))
end


# compute \bar A
@everywhere function barAmat(
  k::Float64,
  weights::Tuple{Array{Float64,1},Array{Float64,1}},
  taus::Array{Float64,1},
  T::Array{Polynomial,1},
  P::Int64,
  N::Int64,
  J::Int64,
  nu1::Int64,
  nu2::Int64,
  nu3::Int64,
  nu4::Int64,
  nu5::Int64
)
  out=zeros(Float64,J*(P+1),J*(P+1),J*(P+1),J*(P+1),J*(P+1))
  for zeta1 in 0:J-1
    for n1 in 0:P
      for zeta2 in 0:J-1
	for n2 in 0:P
	  for zeta3 in 0:J-1
	    for n3 in 0:P
	      for zeta4 in 0:J-1
		for n4 in 0:P
		  for zeta5 in 0:J-1
		    for n5 in 0:P
		      out[zeta1*(P+1)+n1+1,
			  zeta2*(P+1)+n2+1,
			  zeta3*(P+1)+n3+1,
			  zeta4*(P+1)+n4+1,
			  zeta5*(P+1)+n5+1]=1/((2*pi)^5*k^2)*integrate_legendre(tau->barAmat_int1(tau,k,taus,T,weights,nu1,nu2,nu3,nu4,nu5,zeta1,zeta2,zeta3,zeta4,zeta5,n1,n2,n3,n4,n5),taus[zeta1+1],taus[zeta1+2],weights)
		    end
		  end
		end
	      end
	    end
	  end
	end
      end
    end
  end

  return out
end
@everywhere function barAmat_int1(tau,
  k::Float64,
  taus::Array{Float64,1},
  T::Array{Polynomial,1},
  weights::Tuple{Array{Float64,1},Array{Float64,1}},
  nu1::Int64,
  nu2::Int64,
  nu3::Int64,
  nu4::Int64,
  nu5::Int64,
  zeta1::Int64,
  zeta2::Int64,
  zeta3::Int64,
  zeta4::Int64,
  zeta5::Int64,
  n1::Int64,
  n2::Int64,
  n3::Int64,
  n4::Int64,
  n5::Int64
)
  if(alpham(k,tau)<taus[zeta2+2] && alphap(k,tau)>taus[zeta2+1])
    return 2*(1-tau)/(1+tau)^(3-nu1)*T[n1+1]((2*tau-(taus[zeta1+1]+taus[zeta1+2]))/(taus[zeta1+2]-taus[zeta1+1]))*integrate_legendre(sigma->barAmat_int2(tau,sigma,k,taus,T,weights,nu2,nu3,nu4,nu5,zeta2,zeta3,zeta4,zeta5,n2,n3,n4,n5),max(taus[zeta2+1],alpham(k,tau)),min(taus[zeta2+2],alphap(k,tau)),weights)
  else
    return 0.
  end
end
@everywhere function barAmat_int2(tau,
  sigma::Float64,
  k::Float64,
  taus::Array{Float64,1},
  T::Array{Polynomial,1},
  weights::Tuple{Array{Float64,1},Array{Float64,1}},
  nu2::Int64,
  nu3::Int64,
  nu4::Int64,
  nu5::Int64,
  zeta2::Int64,
  zeta3::Int64,
  zeta4::Int64,
  zeta5::Int64,
  n2::Int64,
  n3::Int64,
  n4::Int64,
  n5::Int64
)
  return 2*(1-sigma)/(1+sigma)^(3-nu2)*T[n2+1]((2*sigma-(taus[zeta2+1]+taus[zeta2+2]))/(taus[zeta2+2]-taus[zeta2+1]))*integrate_legendre(taup->barAmat_int3(tau,sigma,taup,k,taus,T,weights,nu3,nu4,nu5,zeta3,zeta4,zeta5,n3,n4,n5),taus[zeta3+1],taus[zeta3+2],weights)
end
@everywhere function barAmat_int3(tau,
  sigma::Float64,
  taup::Float64,
  k::Float64,
  taus::Array{Float64,1},
  T::Array{Polynomial,1},
  weights::Tuple{Array{Float64,1},Array{Float64,1}},
  nu3::Int64,
  nu4::Int64,
  nu5::Int64,
  zeta3::Int64,
  zeta4::Int64,
  zeta5::Int64,
  n3::Int64,
  n4::Int64,
  n5::Int64
)
  if(alpham(k,taup)<taus[zeta4+2] && alphap(k,taup)>taus[zeta4+1])
    return 2*(1-taup)/(1+taup)^(3-nu3)*T[n3+1]((2*taup-(taus[zeta3+1]+taus[zeta3+2]))/(taus[zeta3+2]-taus[zeta3+1]))*integrate_legendre(sigmap->barAmat_int4(tau,sigma,taup,sigmap,k,taus,T,weights,nu4,nu5,zeta4,zeta5,n4,n5),max(taus[zeta4+1],alpham(k,taup)),min(taus[zeta4+2],alphap(k,taup)),weights)
  else
    return 0.
  end
end
@everywhere function barAmat_int4(tau,
  sigma::Float64,
  taup::Float64,
  sigmap::Float64,
  k::Float64,
  taus::Array{Float64,1},
  T::Array{Polynomial,1},
  weights::Tuple{Array{Float64,1},Array{Float64,1}},
  nu4::Int64,
  nu5::Int64,
  zeta4::Int64,
  zeta5::Int64,
  n4::Int64,
  n5::Int64
)
  return 2*(1-sigmap)/(1+sigmap)^(3-nu4)*T[n4+1]((2*sigma-(taus[zeta4+1]+taus[zeta4+2]))/(taus[zeta4+2]-taus[zeta4+1]))*integrate_legendre(theta->barAmat_int5(tau,sigma,taup,sigmap,theta,k,taus,T,weights,nu5,zeta5,n5),0.,2*pi,weights)
end
@everywhere function barAmat_int5(tau,
  sigma::Float64,
  taup::Float64,
  sigmap::Float64,
  theta::Float64,
  k::Float64,
  taus::Array{Float64,1},
  T::Array{Polynomial,1},
  weights::Tuple{Array{Float64,1},Array{Float64,1}},
  nu5::Int64,
  zeta5::Int64,
  n5::Int64
)
  R=barAmat_R((1-sigma)/(1+sigma),(1-tau)/(1+tau),(1-sigmap)/(1+sigmap),(1-taup)/(1+taup),theta,k)
  if((1-R)/(1+R)<taus[zeta5+2] && (1-R)/(1+R)>taus[zeta5+1])
    return (2/(2+R))^nu5*T[n5+1]((2*(1-R)/(1+R)-(taus[zeta5+1]+taus[zeta5+2]))/(taus[zeta5+2]-taus[zeta5+1]))
  else
    return 0.
  end
end
# R(s,t,s',t,theta,k)
@everywhere function barAmat_R(
  s::Float64,
  t::Float64,
  sp::Float64,
  tp::Float64,
  theta::Float64,
  k::Float64
)
  return sqrt(k^2*(s^2+t^2+sp^2+tp^2)-k^4-(s^2-t^2)*(sp^2-tp^2)-sqrt((4*k^2*s^2-(k^2+s^2-t^2)^2)*(4*k^2*sp^2-(k^2+sp^2-tp^2)^2))*cos(theta))/(sqrt(2.)*k)
end

# compute Chebyshev polynomials
@everywhere function chebyshev_polynomials(
  P::Int64
)
  T=Array{Polynomial,1}(undef,P+1)
  T[1]=Polynomial([1])
  T[2]=Polynomial([0,1])
  for n in 1:P-1
    # T_n
    T[n+2]=2*T[2]*T[n+1]-T[n]
  end

  return T
end

# compute \int f*u dk/(2*pi)^3
@everywhere function integrate_f_chebyshev(
  f::Function,
  u::Array{Float64,1},
  k::Array{Float64,1},
  taus::Array{Float64,1},
  weights::Tuple{Array{Float64,1},Array{Float64,1}},
  N::Int64,
  J::Int64
)
  out=0.
  for zeta in 0:J-1
    for i in 1:N
      out+=(taus[zeta+2]-taus[zeta+1])/(16*pi)*weights[2][i]*cos(pi*weights[1][i]/2)*(1+k[zeta*N+i])^2*k[zeta*N+i]^2*u[zeta*N+i]*f(k[zeta*N+i])
    end
  end
  return out
end

@everywhere function inverse_fourier_chebyshev(
  u::Array{Float64,1},
  x::Float64,
  k::Array{Float64,1},
  taus::Array{Float64,1},
  weights::Tuple{Array{Float64,1},Array{Float64,1}},
  N::Int64,
  J::Int64
)
  out=0.
  for zeta in 0:J-1
    for j in 1:N
      out+=(taus[zeta+2]-taus[zeta+1])/(16*pi*x)*weights[2][j]*cos(pi*weights[1][j]/2)*(1+k[zeta*N+j])^2*k[zeta*N+j]*u[zeta*N+j]*sin(k[zeta*N+j]*x)
    end
  end
  return out
end

# compute B (for the computation of the fourier transform of the two-point correlation)
@everywhere function Bmat(
  q::Float64,
  k::Array{Float64,1},
  weights::Tuple{Array{Float64,1},Array{Float64,1}},
  taus::Array{Float64,1},
  T::Array{Polynomial,1},
  P::Int64,
  N::Int64,
  J::Int64,
  nu::Int64
)
  out=Array{Array{Float64,1},1}(undef,J*N)
  for i in 1:J*N
    out[i]=zeros(Float64,J*(P+1))
    for zeta in 0:J-1
      for n in 0:P
	out[i][zeta*(P+1)+n+1]=1/(8*pi^3*k[i]*q)*(betam(k[i],q)>taus[zeta+2] || betap(k[i],q)<taus[zeta+1] ? 0. : integrate_legendre(sigma->(1-sigma)/(1+sigma)^(3-nu)*T[n+1]((2*sigma-(taus[zeta+1]+taus[zeta+2]))/(taus[zeta+2]-taus[zeta+1])),max(taus[zeta+1],betam(k[i],q)),min(taus[zeta+2],betap(k[i],q)),weights))
      end
    end
  end

  return out
end
# beta_-
@everywhere function betam(
  k::Float64,
  q::Float64
)
  return (1-k-q)/(1+k+q)
end
# beta_+
@everywhere function betap(
  k::Float64,
  q::Float64
)
  return (1-abs(k-q))/(1+abs(k-q))
end

# mathfrak S (for the computation of the fourier transform of the two-point correlation)
@everywhere function chebyshev_frakS(
  Ff::Array{Float64,1},
  B::Array{Array{Float64,1},1}
)
  out=zeros(Float64,length(B))
  for i in 1:length(B)
    out[i]=dot(Ff,B[i])
  end
  return out
end