! Parameters shared by all subroutines
module params

  integer, parameter :: a        = 5
  integer, parameter :: ndivs    = 200
  real   , parameter :: m        = 1.0
  real   , parameter :: omega    = 10.0
  real   , parameter :: pi       = acos(-1.0)
  real   , parameter :: sigma    = 0.2
  real   , parameter :: x0       = 1.0
  real   , parameter :: x1       = 1.0
  real   , parameter :: dV       = 5.0
  real   , parameter :: delta_x  = real(a) / real(ndivs)

  integer, parameter :: nsteps   = 5
  real   , parameter :: delta_t  = 0.001

  real   , parameter :: hbar     = 1.0
  complex, parameter :: zi       = (0.0, 1.0)
  real   , parameter :: hm       = hbar * hbar / (m * delta_x * delta_x)
  complex, parameter :: idt      = zi * delta_t / (2.0 * hbar)

end module params


! Defines a harmonic potential
subroutine def_pot(x, v, n) 
  use params
  implicit none

  integer, intent(in) :: n
  real, dimension(n), intent(out) :: x, v

  integer :: i
  
  do i = 1, n
    x(i) = -real(a)/2.0 + real(i-1)*delta_x
    v(i) = m * omega ** 2.0 * x(i) ** 2.0 / 2.0
  enddo
end subroutine def_pot


! Defines a double-well potential
subroutine def_dw_pot(x, v, n) 
  use params
  implicit none

  integer, intent(in) :: n
  real, dimension(n), intent(out) :: x, v

  integer :: i
  
  do i = 1, n
    x(i) = -real(a)/2.0 + real(i-1)*delta_x
    v(i) = dV * ((x(i) / x1)** 2.0 - 1.0)**2.0 
  enddo
end subroutine def_dw_pot


! Defines the initial function psi
subroutine def_psi(x, psi, n)
  use params
  implicit none

  real   , dimension(n), intent(in)  :: x
  complex, dimension(n), intent(out) :: psi
  integer, intent(in) :: n

  integer :: i

  do i = 1, n
    psi(i) = 1.0 / (sigma * sqrt(2.0 * pi)) * exp(- (x(i) - x0)**2.0/(2 * sigma ** 2.0))
  enddo
end subroutine def_psi


! Computes the norm of a wavefunction
real function wavenorm(psi, n)
  use params
  implicit none

  complex, dimension(n), intent(in) :: psi
  integer, intent(in) :: n

  real :: norm
  integer :: i

  norm = 0.0
  do i = 1, n
    norm = norm + delta_x * cabs(psi(i))**2.0
  enddo
  wavenorm = norm
end function wavenorm


! Compute the energy of a state
real function energy(psi, v, n)
  use params
  implicit none

  complex, dimension(n), intent(in) :: psi
  real   , dimension(n), intent(in) :: v
  integer, intent(in) :: n

  real :: en
  integer :: i

  en = 0.0
  en = en + conjg(psi(1)) * ((- hm / 2.0) * psi(2)   + (hm + v(1)) * psi(1)) * delta_x
  en = en + conjg(psi(n)) * ((- hm / 2.0) * psi(n-1) + (hm + v(n)) * psi(n)) * delta_x
  do i = 2, (n-1)
    en = en + conjg(psi(i)) * ((- hm / 2.0) * (psi(i-1) + psi(i+1)) + (hm + v(i)) * psi(i)) * delta_x
  enddo
  energy = en
end function energy


! Normalizes a wavefunction
subroutine normalize(psi, n)
  use params
  implicit none

  complex, dimension(n), intent(inout) :: psi
  integer, intent(in) :: n
  real, external :: wavenorm

  real :: norm
  integer :: i

  norm = wavenorm(psi, n)

  do i = 1, n
    psi(i) = psi(i) / sqrt(norm)
  enddo

end subroutine normalize


! Computes the matrix A (tridiagonal) and the vector B
subroutine compute_AB(n, v, psi, c, d, e, b)
  use params
  implicit none

  integer,                 intent(in)  :: n
  real   , dimension(n)  , intent(in)  :: v
  complex, dimension(n-1), intent(out) :: c, e
  complex, dimension(n)  , intent(out) :: psi, d, b

  integer :: i

  do i = 1, n-1
    c(i) = - idt * hm / 2.0
    e(i) = - idt * hm / 2.0
  enddo

  do i = 1, n
    d(i) = 1.0 + idt * (hm + v(i))
  enddo

  do i = 2, (n-1)
    b(i) = (idt * hm / 2.0) * (psi(i-1) + psi(i+1)) + (1.0 - idt * (hm + v(i))) * psi(i)
  enddo
  b(1) = (idt * hm / 2.0) * psi(2)   + (1.0 - idt * (hm + v(i))) * psi(1)
  b(n) = (idt * hm / 2.0) * psi(n-1) + (1.0 - idt * (hm + v(i))) * psi(n)

end subroutine compute_AB


! Computes the position of a particle
real function pposition(n, psi)
  use params
  implicit none

  integer, intent(in) :: n
  complex, dimension(n), intent(in) :: psi
  
  real :: x, res
  integer :: i

  res = 0.0
  do i = 1, n
    x = real(i-1) * delta_x - real(a)/2.0
    res = res + conjg(psi(i)) * x * psi(i) * delta_x
  enddo
  pposition = res

end function pposition


! Main program
program main
  use params
  implicit none

  real, dimension(ndivs) :: x, v
  complex, dimension(ndivs) :: psi

  complex, dimension(ndivs-1) :: c, e
  complex, dimension(ndivs)   :: d, b

  real, external :: wavenorm, energy, pposition

  integer :: steps, i, info
  real    :: t, xpos, realpsi, imgpsi
  logical :: ok

  call def_pot(x, v, ndivs)

  inquire(file='results/tds_anim.res',exist=ok)

  if (.not. ok) then 
    call def_psi(x, psi, ndivs)
    call normalize(psi, ndivs)
  else 
    open(1, file='results/tds_anim.res')
    do i = 1, ndivs
      read(1,*) t, realpsi, imgpsi 
      psi(i) = complex(realpsi, imgpsi)
    enddo
    close(1)
  endif

  open(1, file='results/tds_anim.res')
  do steps = 1, nsteps
    t = real(steps) * delta_t
    call compute_AB(ndivs, v, psi, c, d, e, b)
    call cgtsv(ndivs, 1, c, d, e, b, ndivs, info)
    if(.not. (info == 0)) then
      write(*,*) "CGTSV failure"
      exit
    endif
    psi = b
  enddo
  do i = 1, ndivs
    xpos = -real(a)/2.0 + real(i-1)*delta_x
    write(1,*) xpos, real(real(psi(i))), real(aimag(psi(i)))
  enddo
  close(1)
  
end program main


