! attention, emacs: the -*- mode: f90; compile-command: "make homework-2" -*- mode

module vector_calculus
  intrinsic :: dot_product, sqrt

contains

  function unit(r)
    real, dimension(:), intent(in) :: r
    real, dimension(size(r)) :: unit
    unit = r / len(r)
  end function unit

  function len(r)
    real, dimension(:), intent(in) :: r
    real :: len
    len = sqrt(dot_product(r,r))
  end function len

end module vector_calculus

program three_body
  use vector_calculus

  ! Simulation parameters
  real, parameter :: &
       precision_goal = 0.01, &
       dt_init = 0.01, &
       duration = 20, &
       n = 2, &
       d = 2

  ! Body masses
  real, dimension(n), parameter :: m = (/ 1.0, 1.0 /)

  ! Body parameters
  real, dimension(n,d) :: &
       x, v, &                  ! Configuration at current t
       x_n, v_n, &              ! Configuration at t+dt with step dt
       x_hq, v_hq               ! Configuration at t=dt with step dt/2



  ! Physical constants
  real, parameter :: &
       ! G = 6.67428e-11
       G = 1

  ! Variables
  real :: &
       dt = dt_init, &
       t = 0
  integer :: i,j

  ! Initial conditions
  x(1,:) = (/ 1.0, 0.0 /)
  x(2,:) = (/ 0.0, 0.0 /)
  !x(3,:) = (/ 5.0, 5.0 /)
  v(1,:) = (/ 0.0, 1.0 /)
  v(2,:) = (/ 0.0, 0.0 /)
  !v(3,:) = (/ 0.0, 0.0 /)

  ! Loop over time
  main: do

     ! Calculate next position with step dt
     call step(x,v,dt, x_n, v_n)

     ! Calculate next position with step dt/2
     call step(x, v, dt/2, x_hq, v_hq)
     call step(x_hq, v_hq, dt/2, x_hq, v_hq)

     ! Redo with dt/2 if we do not reach precision goal
     if (any(abs(x_hq - x_n > precision_goal)) .or. &
          any(abs(v_hq - v_n > precision_goal))) then
        dt = dt/2
        cycle
     end if

     ! Update variables
     t = t + dt
     x = x_n
     v = v_n

     ! Fix this guy for now
     x(2,:) = (/ 0.0, 0.0 /)

     if (t > duration) exit
  end do main



contains

  subroutine step(x, v, dt, xn, vn)
    real, dimension(:,:), intent(in) :: x, v
    real, dimension(:,:), intent(out) :: xn, vn
    real, dimension(size(x,1),size(x,2)) :: a_m, v_m, x_m, a
    real, intent(in) :: dt
    integer :: n,d


    n = size(x,1)
    d = size(x,2)

    ! For each body
    do i=1,n

       ! Calculate current acceleration
       a(i,:)=0
       do j=1,d
          if (i /= j) then
             a(i, :) = a(i, :) + G * m(j) * unit(x(j,:) - x(i,:)) / &
                  len(x(i,:) - x(j,:))**2
          end if
       end do

       ! Calculate expected new velocity and position
       v_m(i,:) = v(i,:) + 0.5 * a(i,:) * dt
       x_m(i,:) = x(i,:) + 0.5 * v(i,:) * dt
    end do


    ! For each body
    do i=i,n
       ! Calculate new acceleration
       a_m(i,:) = 0
       do j=1,d
          if (i /= j) then
             a_m(i, :) = a_m(i, :) + G * m(j) * unit(x_m(j,:) - x_m(i,:)) / &
                  len(x_m(i,:) - x_m(j,:))**2
          end if
       end do

       ! Calculate final new velocity and position
       vn(i,:) = v(i,:) + a_m(i,:) * dt
       xn(i,:) = x(i,:) + v_m(i,:) * dt
    end do
  end subroutine step

end program three_body