program landau_damping

    !------ based on code written by Michael Barnes ------!

    implicit none

    !-------------- input parameters ---------------------!

    ! Te/Ti
    real, parameter :: alpha = 1.

    ! time grid parameters
    integer, parameter :: nstep = 40000
    real, parameter :: dt = 0.0001
    integer, parameter :: nwrite = 1000

    ! k-space parameters
    ! k = 2*pi*kint
    integer, parameter :: kint = 1

    ! Hermite (m-space) parameters
    ! number of Hermite polynomials
    integer, parameter :: nm = 128

    ! coefficient for artificial hyperviscosity of form -nu_m * m**4
    real, parameter :: nu_m = 0.000000001
    !real, parameter :: nu_m = 0.01
    
    !------------------------------------------------------!
    ! keeps track of time variable
    real :: time

    ! the constant pi
    real :: pi

    ! the imaginary unit
    complex :: zi = (0.,1.)

    ! wavenumber of mode in Fourier approach
    real :: k

    ! for convenience (k*dt/sqrt(2))
    real :: kdt
    
    ! gk(m)
    complex, dimension (:), allocatable :: gkm
    real, dimension (:), allocatable :: gkmr, gkmi

    ! arrays needed to fill tridiagnoal advance matrix in spectral approach
    real, dimension (:), allocatable :: aa, bb, cc

    integer :: phik2_unit = 104, gkm_unit = 105

    ! define pi for later use
    pi = 2.*acos(0.)

    k = 2.*pi*kint
    kdt = k*dt/sqrt(2.0)
    
    call init_io

    ! allocates and initializes g(k,m)
    call init_gkm
    ! implicit in time, spectral in x and v solver
    call spectral_solve

    call finish_io

    deallocate (gkm, gkmr, gkmi)
    deallocate (aa, bb, cc)

  contains

    subroutine init_io

      implicit none

      open (phik2_unit, file='landau_damping.phik2', status='replace')
      open (gkm_unit, file='landau_damping.gkm', status='replace')

    end subroutine init_io

    subroutine finish_io

      implicit none

      close (phik2_unit)
      close (gkm_unit)

    end subroutine finish_io

    subroutine init_gkm

      implicit none

      if (.not.allocated(gkm)) then
        allocate (gkm(0:nm-1))
        allocate (gkmr(nm))
        allocate (gkmi(nm))
      end if

      gkm(0) = 1.0
      gkm(1:) = 0.0

    end subroutine init_gkm

    subroutine spectral_solve

      implicit none

      integer :: istep

      time = 0.0
      call write_data (0)

      call init_matrix

      do istep = 1, nstep
        call update_gkm
        time = time + dt
        call write_data (istep)
      end do

    end subroutine spectral_solve

    subroutine init_matrix

      implicit none

      integer :: i, sgn
      Real :: ir
      
      if (.not. allocated(aa)) then
        allocate (aa(nm)) ; aa = 0.
        allocate (bb(nm)) ; bb = 1.
        allocate (cc(nm)) ; cc = 0.
      end if

      ! m = 2...(nm-1)
      do i = 3, nm-1
        ir = REAL(i)
        sgn = (-1)**(i-1)
        aa(i) = sgn*kdt*sqrt(ir-1.0)
        cc(i) = sgn*kdt*sqrt(ir)
      end do

      ! m = 0
      cc(1) = kdt

      ! m = 1
      aa(2) = -kdt*(1.0 + alpha)
      cc(2) = -k*dt

      ! m = nm-1
      ir = REAL(nm)
      aa(nm) = (-1)**(nm-1)*kdt*sqrt(ir-1.0)

      do i = 3, nm
         ir = REAL(i-1)
         bb(i) = bb(i) + nu_m*ir**4.0
         !bb(i) = bb(i) + nu_m*mreal
      end do

    end subroutine init_matrix

    subroutine update_gkm

      implicit none

      integer :: i

      do i = 0, nm-1
        if (mod(i,2)==0) then
          gkmr(i+1) = real(gkm(i))
          gkmi(i+1) = aimag(gkm(i))
        else
          gkmr(i+1) = -aimag(gkm(i))
          gkmi(i+1) = real(gkm(i))
        end if
      end do
      call tridiag (aa, bb, cc, gkmr)
      call tridiag (aa, bb, cc, gkmi)
      do i = 0, nm-1
        if (mod(i,2)==0) then
          gkm(i) = gkmr(i+1) + zi*gkmi(i+1)
        else
          gkm(i) = gkmi(i+1) - zi*gkmr(i+1)
        end if
      end do

    end subroutine update_gkm

    subroutine write_data (istep)

      implicit none

      integer, intent (in) :: istep

      integer :: i

      if (mod(istep,nwrite)==0) then
        do i = 0, nm-1
          write (gkm_unit,*) time, i, real(gkm(i)), aimag(gkm(i))
        end do
        write (gkm_unit,*)
      end if

      write (phik2_unit,*) time, real(conjg(gkm(0))*gkm(0))*alpha**2

      write (*,*) 'time = ', time, '|phik|**2 = ',real(conjg(gkm(0))*gkm(0))*alpha**2

    end subroutine write_data

    ! solves system Ax = b for x (which is returned as sol)
    ! inputs are aa, bb, and cc (the elements to the left, center, and right
    ! of diagonal in tridiagonal matrix A)
    ! and sol=b as the rhs of the linear system Ax = b
    subroutine tridiag (aa, bb, cc, sol)

      implicit none

      real, dimension (:), intent (in) :: aa, bb, cc
      real, dimension (:), intent (in out) :: sol

      integer :: ix, npts
      real :: bet

      real, dimension (:), allocatable :: gam

      npts = size(aa)
      allocate (gam(npts))

      bet = bb(1)
      sol(1) = sol(1)/bet

      do ix = 2, npts
        gam(ix) = cc(ix-1)/bet
        bet = bb(ix) - aa(ix)*gam(ix)
        if (bet == 0.0) write (*,*) 'tridiagonal solve failed'
        sol(ix) = (sol(ix)-aa(ix)*sol(ix-1))/bet
      end do

      do ix = npts-1, 1, -1
        sol(ix) = sol(ix) - gam(ix+1)*sol(ix+1)
      end do

      deallocate (gam)

    end subroutine tridiag

  end program landau_damping