! Module with some useful functions and global variables.
! This is used to automatically create a function interface.
module utils
  implicit none

  real, parameter :: w0=1.0, m=1.0, xi=0.0, vi=0.0
  real, parameter :: tau=0.5, gamma_=1.0, vf=0.1
  real, parameter :: dt=1e-3

  real, parameter :: tf=200.0

  integer, parameter :: npads=5

  real :: v0

  contains
  
  subroutine rk4(x,y,dx,n,df)
    implicit none
    integer, intent(in)                  :: n
    real   , dimension(n), intent(inout) :: y
    real   , intent(in)                  :: x, dx
    real   , dimension(n)                :: d1, d2, d3, d4, yp
    real                                 :: ddx, c = 1.0/6.0

    ddx = dx/2.0
    call df(x    , y , d1, n); yp = y+d1*ddx
    call df(x+ddx, yp, d2, n); yp = y+d2*ddx
    call df(x+ddx, yp, d3, n); yp = y+d3*dx
    call df(x+dx , yp, d4, n); y  = y+c*(d1+2*d2+2*d3+d4)*dx

  end subroutine rk4

end

! Function that computes derivatives for the harmonic oscillator
subroutine harmonic_df(t, eq, deq, n)
  use utils
  implicit none

  real :: t
  integer :: n
  real, dimension(0:1) :: eq, deq

  deq(0) = eq(1)
  deq(1) = -(w0**2.0)*eq(0)

end

! Function that computes derivatives for npads pads
subroutine pad_df(t, eq, deq, n)
  use utils
  implicit none

  real :: t
  integer :: n, i
  real, dimension(0:2*npads-1) :: eq, deq
  real :: xprev, xnext

  xprev = v0 * t

  if (npads > 1) then
    xnext = eq(2)
  else
    xnext = eq(0)
  endif

  do i = 0, npads-1
    deq(2*i)   = eq(2*i+1)
    deq(2*i+1) = w0 * w0 * (xprev - 2.0 * eq(2*i) + xnext) - eq(2*i+1) / tau 
    deq(2*i+1) = deq(2*i+1) - sign(1.0, eq(2*i+1)) * gamma_ * exp(-(abs(eq(2*i+1)) / vf))
    xprev      = eq(2*i)
    if (i < npads-2) then
      xnext = eq(2*(i+2))
    else
      xnext = eq(2*(i+1))
    endif
  enddo

end

! Main program
program main
  use utils
  implicit none

  integer :: i
  real    :: cur_t, energy
  real, dimension(0:2*npads-1) :: x
  
  external :: harmonic_df, pad_df

  v0 = 0.1

  open(1, file='results/pads5.out')

  do i = 0, npads-1
    x(2*i  ) = xi
    x(2*i+1) = vi
  enddo

  do i = 0, int(tf/dt)
    cur_t = real(i) * dt
    call rk4(cur_t, x, dt, 2*npads, pad_df)
    ! energy = w0 * w0 * m * x(0) * x(0) / 2.0 + m * x(1) * x(1) / 2.0
    write(1, *) cur_t, x(0:2*npads-1:1)!, energy
  enddo

  close(1)

end
