using FFTW
using Printf
using SpecialFunctions
using LinearAlgebra

# numerical values
hbar=6.58e-16 # eV.s
m=9.11e-31 # kg
V_=6.6 # eV
I=1e17 # W/m^2
EF=3.3 # eV
omega_=1.5498*4 # eV

# dimensionless quantities
V=1
k=sqrt(EF/V_)
epsilon=sqrt(hbar^3*I/(8*m*V_*omega_^2))
omega=omega_/V_

#epsilon=8*epsilon
#omega=4*omega

nu=60
ell=20

# for Crank-Nicolson
xmin_cn=-100
xmax_cn=100
nx_cn=4000
tmax=10
# mu=dt/dx^2
mu=0.1
nt=floor(Int,tmax/(mu*((xmax_cn-xmin_cn)/nx_cn)^2))
# only save up to memt time steps
memt=10000


# asymptotic value of \varphi_t(x)
function phi_asym(x,T,gkz,k,epsilon,omega,V,nu,ell)
  # compute residues
  res=residue(x,T,gkz,k,epsilon,omega,V,nu)
  times=compute_times(omega,nu,ell)
  N=(2*nu+1)*(2*ell+1)
  fun=zeros(Complex{Float64},N)
  for m in 1:N
    for n in -nu:nu
      fun[m]+=-1im*res[n+nu+1]*exp(-1im*(k*k+n*omega)*times[m])
    end
  end
  return fun
end
# asymptotic value of \partial\varphi_t(x)
function dphi_asym(x,T,gkz,k,epsilon,omega,V,nu,ell)
  # compute residues
  res=dresidue(x,T,gkz,k,epsilon,omega,V,nu)
  times=compute_times(omega,nu,ell)
  N=(2*nu+1)*(2*ell+1)
  fun=zeros(Complex{Float64},N)
  for m in 1:N
    for n in -nu:nu
      fun[m]+=-1im*res[n+nu+1]*exp(-1im*(k*k+n*omega)*times[m])
    end
  end
  return fun
end

# return times
function compute_times(omega,nu,ell)
  N=(2*nu+1)*(2*ell+1)
  times=zeros(N)
  for j in 1:N
    times[j]=pi/(omega*(nu+0.5))*(j-1)
  end
  return times
end

# \kappa_m^{(\sigma)}
function K(m,sigma,V,k,epsilon,omega)
  return pow(V+2*epsilon*epsilon-k*k-sigma-m*omega,0.5,pi/2)
end
# mathfrak q_k
function Q(k,V)
  return sqrt(V-k*k)
end

# T_0
function T0(k,V)
  return 2im*k/(1im*k-sqrt(V-k*k))
end

# compute g
function g(kappa,epsilon,omega,N)
  B=zeros(Complex{Float64},N)
  # prepare vector of B's
  for i in 1:N
    t=2*pi/(omega*N)*(i-1)
    B[i]=exp(1im*epsilon*epsilon/omega*sin(2*omega*t)-kappa*4*epsilon/omega*cos(omega*t))
  end
  # take the fft
  fft!(B)
  # correct by a factor 1/N
  for i in 1:N
    B[i]=B[i]/N
  end
  return(B)
end

# index of m-n
function mmn(n,nu)
  return (n>=0 ? n+1 : n+4*nu+2)
end

# compute residues at -ik^2-in\omega
function residue(x,T,gkz,k,epsilon,omega,V,nu)
  # compute residue
  out=zeros(Complex{Float64},2*nu+1)
  if x>=0
    for n in -nu:nu
      for m in -nu:nu
	out[n+nu+1]+=1im*gkz[m+nu+1][mmn(n-m,nu)]*exp(-K(m,0,V,k,epsilon,omega)*x)*T[m+nu+1]
      end
    end
  else
    for n in -nu:nu
      out[n+nu+1]=(n==0 ? 1im*(exp(1im*k*x)-exp(-1im*k*x)) : 0)
      for m in -nu:nu
	out[n+nu+1]+=1im*T[m+nu+1]*gkz[m+nu+1][mmn(n-m,nu)]*exp(-1im*x*pow(k*k+n*omega,0.5,-pi/2))
      end
    end
  end

  return out
end
# compute \partial residues at -ik^2-in\omega
function dresidue(x,T,gkz,k,epsilon,omega,V,nu)
  # compute residue
  out=zeros(Complex{Float64},2*nu+1)
  if x>=0
    for n in -nu:nu
      for m in -nu:nu
	out[n+nu+1]+=-K(m,0,V,k,epsilon,omega)*1im*gkz[m+nu+1][mmn(n-m,nu)]*exp(-K(m,0,V,k,epsilon,omega)*x)*T[m+nu+1]
      end
    end
  else
    for n in -nu:nu
      out[n+nu+1]=(n==0 ? -k*(exp(1im*k*x)+exp(-1im*k*x)) : 0)
      for m in -nu:nu
	out[n+nu+1]+=pow(k*k+n*omega,0.5,-pi/2)*T[m+nu+1]*gkz[m+nu+1][mmn(n-m,nu)]*exp(-1im*x*pow(k*k+n*omega,0.5,-pi/2))
      end
    end
  end

  return out
end


# matching condition for residue
function residue_matching(k,epsilon,omega,V,nu)
  # preliminary: g^{(\kappa_m^{(0)})}
  # g^{(\kappa_m^{(0)})}_n=gkz[m+nu+1][n+1] if n>=0 and gkz[m+nu+1][n+4*nu+2] if n<0
  gkz=Array{Array{Complex{Float64}},1}(undef,2*nu+1)
  for m in -nu:nu
    gkz[m+nu+1]=g(K(m,0,V,k,epsilon,omega),epsilon,omega,4*nu+1)
  end

  # solve matching condition
  G=zeros(Complex{Float64},2*nu+1,2*nu+1)
  v=zeros(Complex{Float64},2*nu+1)
  for n in -nu:nu
    for m in -nu:nu
      G[n+nu+1,m+nu+1]=((K(m,0,V,k,epsilon,omega)-1im*pow(k*k+n*omega,0.5,-pi/2))*gkz[m+nu+1][mmn(n-m,nu)]-epsilon*(gkz[m+nu+1][mmn(n-m+1,nu)]-gkz[m+nu+1][mmn(n-m-1,nu)]))
    end
  end
  v[nu+1]=-2im*k
  
  return (G\v,gkz)
end


# fractional power with an arbitrary branch cut
function pow(x,a,cut)
  if(angle(x)/cut<=1)
    return(abs(x)^a*exp(1im*angle(x)*a))
  else
    return(abs(x)^a*exp(1im*(angle(x)-sign(cut)*2*pi)*a))
  end
end

# current
function J(phi,dphi,x,epsilon,omega,t)
  if x>=0
    return(2*imag(conj(phi)*(dphi+2im*epsilon*sin(omega*t)*phi)))
  else
    return(2*imag(conj(phi)*dphi))
  end
end

# compute \psi_t(x) using Crank-Nicolson
function psi_cn(V,omega,epsilon,k,xs,ts,memt)
  nx=length(xs)
  nt=length(ts)

  # length of output vector
  len=min(nt,memt)
  # how often to write to output
  freq=ceil(Int,nt/memt)

  # initialize
  psis=zeros(Complex{Float64},nx,len)
  times=zeros(Float64,len)

  # previous time step
  psi_prev=zeros(Complex{Float64},nx)

  # init
  for i in 1:nx
    times[1]=0
    psis[i,1]=(xs[i]<0 ? exp(1im*k*xs[i])+(T0(k,V)-1)*exp(-1im*k*xs[i]) : T0(k,V)*exp(-Q(k,V)*xs[i]))
    psi_prev[i]=psis[i,1]
  end
  # the boundary condition is that psi at the boundary is constant (nothing needs to be done for this)

  # matrix structure of the Crank-Nicolson algorithm
  # diagonals of M
  M0=zeros(Complex{Float64},nx-2)
  Mp=zeros(Complex{Float64},nx-3)
  Mm=zeros(Complex{Float64},nx-3)
  v=zeros(Complex{Float64},nx-2)

  for j in 1:nt-1
    # print progress
    progress(j,nt-1,1000)

    # t_j
    t=ts[j]
    # t_{j+1}
    tp=ts[j+1]
    dt=tp-t
    for mu in 2:nx-1
      # x_mu
      x=xs[mu]
      dx2=(xs[mu+1]-xs[mu])*(xs[mu]-xs[mu-1])
      # tridiagonal matrix
      M0[mu-1]=1im/dt-1/dx2-(x<0 ? 0 : (V-2*epsilon*omega*cos(omega*tp)*x)/2)
      if mu<nx-1
	Mp[mu-1]=1/(2*dx2)
      end
      if mu>2
	Mm[mu-2]=1/(2*dx2)
      end
      # right side
      v[mu-1]=1im*psi_prev[mu]/dt-(psi_prev[mu+1]+psi_prev[mu-1]-2*psi_prev[mu])/(2*dx2)+(x<0 ? 0 : (V-2*epsilon*omega*cos(omega*t)*x)*psi_prev[mu]/2)
      # correct for boundary conditions
      # assumes the boundary condition is constant in time!
      if mu==2
	v[mu-1]-=psi_prev[1]/(2*dx2)
      end
      if mu==nx-1
	v[mu-1]-=psi_prev[nx]/(2*dx2)
      end
    end

    M=Tridiagonal(Mm,M0,Mp)

    # copy to psi_prev
    psi_prev[2:nx-1]=M\v

    # write to output
    if j%freq==0
      psis[:,Int(j/freq)+1]=psi_prev
      times[Int(j/freq)+1]=tp
    end
  end

  return (times,psis)
end

# returns the times in b that are closest to those in a
# assumes the vectors are ordered
function nearest_times(a,b)
  out=zeros(Int,length(a))
  pointer=1
  for i in 1:length(a)
    if pointer==length(b)
      out[i]=length(b)
    end
    for j in pointer:length(b)-1
      if b[j+1]>a[i]
	out[i]=(abs(b[j+1]-a[i])<abs(b[j]-a[i]) ? j+1 : j)
	pointer=out[i]
	break
      end
      # all remaining points are beyond the range
      if j==length(b)-1
	out[i]=length(b)
	pointer=out[i]
      end
    end
  end
  return out
end

# space derivative using Crank Nicolson
function deriv_cn(phi,i,xs)
  if i==1
    return (phi[i+1]-phi[i])/(xs[i+1]-xs[i])
  end
  if i==length(phi)
    return (phi[i]-phi[i-1])/(xs[i]-xs[i-1])
  end
  return (phi[i+1]-phi[i])/(xs[i+1]-xs[i])/2 + (phi[i]-phi[i-1])/(xs[i]-xs[i-1])/2
end

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


# print progress
function progress(j,tot,freq)
  if (j-1)%ceil(Int,tot/freq)==0
    if j>1
      @printf(stderr,"\r")
    end
    @printf(stderr,"%d/%d",j,tot)
  end
  if j==tot
    @printf(stderr,"%d/%d\n",j,tot)
  end
end

# print animation data using Crank Nicolson and compare to Laplace
function print_anim_cn_cmp()
  @printf(stderr,"epsilon=% .8e omega=% .8e k=% .8e V=% .8e z=% .8e T=% .8e gamma=% .8e\n",epsilon,omega,k,V,4*sqrt(nu)*epsilon/sqrt(omega),2*pi/(omega_/hbar)*1e15,sqrt((V-k^2)/(4*epsilon^2)))
  xmin=-10
  xmax=10
  nx=200

  # array of times for the asymptote
  times=compute_times(omega,nu,ell)

  # positions at which to compute phi
  xs=Array(0:nx_cn)*(xmax_cn-xmin_cn)/nx_cn.+xmin_cn
  # times at which to compute phi
  ts=Array(0:nt-1)*tmax/nt

  # compute matching condition
  (T,gkz)=residue_matching(k,epsilon,omega,V,nu)

  # compute wave function
  (t_psi,psi)=psi_cn(V,omega,epsilon,k,xs,ts,memt)

  # nearest times using Crank Nicolson
  indices_t=nearest_times(times,t_psi)
  indices_x=nearest_times(Array(0:nx-1)*(xmax-xmin)/nx.+xmin,xs)

  ps_cn=Array{Array{Complex{Float64},1}}(undef,nx)
  ps_asym=Array{Array{Complex{Float64},1}}(undef,nx)
  dps_cn=Array{Array{Complex{Float64},1}}(undef,nx)
  dps_asym=Array{Array{Complex{Float64},1}}(undef,nx)
  for i in 1:nx
    progress(i,nx,1000)
    x=xmin+(xmax-xmin)*(i-1)/nx
    ps_cn[i]=Array{Complex{Float64}}(undef,length(times))
    dps_cn[i]=Array{Complex{Float64}}(undef,length(times))
    for j in 1:length(times)
      # phase change in magnetic gauge
      phase=(x<0 ? 1 : exp(-2im*epsilon*sin(omega*times[j])*x))
      ps_cn[i][j]=phase*psi[indices_x[i],indices_t[j]]
      dps_cn[i][j]=phase*deriv_cn(psi[:,indices_t[j]],indices_x[i],xs)-(xs[indices_x[i]]<0 ? 0 : 2im*epsilon*sin(omega*times[j])*psi[indices_x[i],indices_t[j]])
    end
    ps_asym[i]=phi_asym(x,T,gkz,k,epsilon,omega,V,nu,ell)
    dps_asym[i]=dphi_asym(x,T,gkz,k,epsilon,omega,V,nu,ell)
  end

  # print values at each time
  for j in 1:length(times)
    if times[j]>tmax
      break
    end
    for i in 1:nx
      x=(xmin+(xmax-xmin)*i/nx)
      # remove point at 0
      if i!=Int(nx/2)
	@printf("% .8e % .8e % .8e % .8e % .8e % 8e\n",times[j]*hbar/V_*1e15,x*sqrt(hbar^2*1.6e-19/(2*m*V_))*1e9,abs(ps_asym[i][j])^2,abs(ps_cn[i][j])^2,J(ps_asym[i][j],dps_asym[i][j],x,epsilon,omega,times[j]),J(ps_cn[i][j],dps_cn[i][j],x,epsilon,omega,times[j]))
      end
    end
    print('\n')
  end
end


print_anim_cn_cmp()