! Compile with:
! gfortran kdv_for_students.f90 -o kdv -O2 -Wall
program kdv
  implicit none

  integer, parameter          :: nx          = 400      ! Number of cells
  double precision, parameter :: pi          = acos(-1.0d0) ! pi = 3.1415...
  double precision, parameter :: dx          = 2.0d0/nx ! Grid spacing
  double precision, parameter :: t_end       = 4.0/pi   ! End time
  double precision, parameter :: delta       = 0.022d0
  double precision, parameter :: dt_output   = t_end * 4.0d-3
  double precision, parameter :: umax_for_dt = 10.0d0
  double precision, parameter :: dt_constant = 3.0d0 * sqrt(3.0d0)

  integer            :: my_unit
  integer            :: i, n, n_steps, n_output
  double precision   :: u(-1:nx+2, 3)
  double precision   :: x(nx)
  double precision   :: dt
  character(len=200) :: filename

  ! Set up initial condition
  do i = 1, nx
     x(i)    = (i-1) * dx
     u(i, 1) = cos(pi * x(i))
  end do

  ! Copy periodic boundaries
  call fill_ghost_cells(u(:, 1))

  dt = min(dx**3 / (dt_constant * delta**2), dx/umax_for_dt)

  ! Copy state to t+1 (we could also do a forward Euler step)
  u(:, 2) = u(:, 1)

  ! Update in time
  n_steps  = ceiling(t_end / dt)
  n_output = 0

  do n = 1, n_steps
     do i = 1, nx
        ! Here you have to implement the scheme from the assignment
        ! u(i, 3) corresponds to u_i at time t+1
        ! u(i, 2) corresponds to u_i at time t
        ! u(i, 1) corresponds to u_i at time t-1
     end do

     call fill_ghost_cells(u(:,3 ))

     ! Save new data at t-1 and t-2
     u(:, 1) = u(:, 2)
     u(:, 2) = u(:, 3)

     if (n * dt > n_output * dt_output) then
        n_output = n_output + 1
        write(filename, "(A,I0.4,A)") "kdv_", n_output, ".txt"
        open(newunit=my_unit, file=trim(filename), action="write")
        do i = 1, nx
           write(my_unit, *) x(i), u(i, 3)
        end do
        close(my_unit)
     end if
  end do

contains

  subroutine fill_ghost_cells(uvec)
    double precision, intent(inout) :: uvec(-1:nx+2)

    ! Copy periodic boundaries
    uvec(0)    = uvec(nx)
    uvec(-1)   = uvec(nx-1)
    uvec(nx+1) = uvec(1)
    uvec(nx+2) = uvec(2)
  end subroutine fill_ghost_cells

end program kdv
