MODULE Mod_Orbits

!######################################################################
!#
!# Filename: Mod_Orbits.f90
!#
!# Description: N-body module for solving 2-d, collisionless orbits.  The
!#  user defines the mass, initial position, and initial velocity of each
!#  body.  The solver is a 4th-order Runge-Kutta.  Units are AU, years, and
!#  solar mass. 
!#
!# Dependencies:
!#   NONE
!#
!# History:
!#   v1.0 - 04/2/07 - Peter Mendygral
!#
!######################################################################

IMPLICIT NONE

!### create a data structure for each of the bodies ###
!### mass is in solar unit, positions in AU, velocities ###
!### are in AU / year ###
TYPE :: Body
	REAL*8 :: mass, x, y, vx, vy
END TYPE Body

!### create a structure for the simulation ###
TYPE :: Orbits
	TYPE(Body), ALLOCATABLE :: bodies (:)
	INTEGER :: nbodies

	!### dt is initial time step, mu is center of mass ###
	REAL*8 :: dt, xcm, ycm, totmass, tolerance, error
END TYPE Orbits

!### define our constants ###
REAL*8, PARAMETER :: G = -3.94784176044d1

!######################################################################
!### SUBROUTINES
!######################################################################

CONTAINS

SUBROUTINE OrbitsCopy(orig, copy)

	TYPE(Orbits) :: orig
	TYPE(Orbits), INTENT(INOUT) :: copy

	copy%bodies    = orig%bodies
	copy%nbodies   = orig%nbodies
	copy%dt        = orig%dt
	copy%xcm       = orig%xcm
	copy%ycm       = orig%ycm
	copy%totmass   = orig%totmass
	copy%tolerance = orig%tolerance
	copy%error     = orig%error

END SUBROUTINE

!### OrbitsConstruct will construct the class for the simulation and ###
!### takes options for the number of bodies and the nominal dt ###
SUBROUTINE OrbitsConstruct (self, nbodies, dt, tolerance)

	TYPE(Orbits), INTENT(INOUT) :: self
	INTEGER :: nbodies
	REAL*8 :: dt, tolerance

	!### allocate space to the bodies array ###
	ALLOCATE(self%bodies(nbodies))

	!### set the passed parameters ###
	self%nbodies   = nbodies
	self%dt        = dt
	self%tolerance = tolerance
	self%error     = 0.d0
END SUBROUTINE

!### OrbitsInit takes arrays for initial x, initial y, initial vx, and ###
!### initial vy for each of the bodies and initializes the simulation ###
SUBROUTINE OrbitsInit (self, masses, xs, ys, vx, vy)

	TYPE(Orbits), INTENT(INOUT) :: self
	REAL*8, DIMENSION (self%nbodies) :: masses, xs, ys, vx, vy
	INTEGER :: n

	self%totmass = 0.d0
	self%xcm     = 0.d0
	self%ycm     = 0.d0

	!### loop through each of the bodies ###
	DO n = 1, self%nbodies

		!### store the mass, position and velocity ###
		self%bodies(n)%mass = masses(n)
		self%bodies(n)%x    = xs(n)
		self%bodies(n)%y    = ys(n)
		self%bodies(n)%vx   = vx(n)
		self%bodies(n)%vy   = vy(n)

		!### sum up the center of mass ###
		self%totmass = self%totmass + masses(n)
		self%xcm     = self%xcm + masses(n) * self%bodies(n)%x
		self%ycm     = self%ycm + masses(n) * self%bodies(n)%y
	END DO

	self%xcm = self%xcm / self%totmass
	self%ycm = self%ycm / self%totmass

END SUBROUTINE

!### p_fderiv is a private function that will return the first ###
!### derivative of the orbit equation ###
REAL*8 FUNCTION p_fderiv (self, x, v, t)

	TYPE(Orbits), INTENT(INOUT) :: self
	REAL*8 :: x, v, t
	p_fderiv = v
	RETURN
END function

!### p_sderiv is a private function that will return the second ###
!### derivative of the orbit equation ###
REAL*8 FUNCTION p_sderiv (self, x, y, mass, xo, yo, t)

	TYPE(Orbits), INTENT(INOUT) :: self
	REAL*8 :: x, y, mass, xo, yo, t
	p_sderiv = G * (x - xo) * mass / (((x - xo)**2 + (y - yo)**2)**(3.d0 / 2.d0))
	RETURN
END FUNCTION

!### OrbitAdaptStep will create an adaptive step to the ODE ###
!### if the error in the solution is sufficiently large a half ###
!### time step will be made.  If the error is small the time step ###
!### will be increased ###
SUBROUTINE OrbitAdaptStep (self)

	TYPE(Orbits), INTENT(INOUT) :: self
	TYPE(Orbits) :: selfcopyA, selfcopyB
	REAL*8 :: maxerrorA, maxerrorB
	LOGICAL :: converged

	!### allocate space to the temporary data structures ###
	ALLOCATE(selfcopyA%bodies(self%nbodies),selfcopyB%bodies(self%nbodies))

	!### create temporary copies of the particle data ###
	CALL OrbitsCopy(self, selfcopyA)
	CALL OrbitsCopy(self, selfcopyB)

	!### first try an update with the last given time step ###
	CALL OrbitStep(selfcopyA)

	!### try another with twice the time step ###
	selfcopyB%dt = 2.d0 * self%dt
	CALL OrbitStep(selfcopyB)

	!### calculate the maximum error in position for the distribution ###
	CALL CalcMaxError(self, selfcopyA, maxerrorA)
	CALL CalcMaxError(self, selfcopyB, maxerrorB)

	!### loop on the update until the error in each paricles position ###
	!### is within machine error ###
	converged = .false.
	DO WHILE (.NOT.converged)

		!PRINT '("DEBUG: maxerrorA = ",ES16.4," maxerrorB = ",ES16.4)', maxerrorA, maxerrorB

		!### if maxerrorB is within range we can still try doubling the time step ###
		IF (maxerrorB.LT.self%tolerance) THEN

			CALL OrbitsCopy(selfcopyB, selfcopyA)
			maxerrorA    = maxerrorB
			CALL OrbitsCopy(self, selfcopyB)
			selfcopyB%dt = 2.d0 * selfcopyA%dt

			CALL OrbitStep(selfcopyB)
			CALL CalcMaxError(self, selfcopyB, maxerrorB)

		!### if maxerrorA is out of tolerance then we'll half the time step ###
		!### and try again ###
		ELSE IF (maxerrorA.GT.self%tolerance) THEN

			CALL OrbitsCopy(selfcopyA, selfcopyB)
			maxerrorB    = maxerrorA
			CALL OrbitsCopy(self, selfcopyA)
			selfcopyA%dt = 5.d-1 * selfcopyB%dt

			CALL OrbitStep(selfcopyB)
			CALL CalcMaxError(self, selfcopyA, maxerrorA)
		
		!### if maxerrorB is out of tolerance and maxerrorA is within, ###
		!### use the A update and we're done here ###
		ELSE IF (maxerrorB.GT.self%tolerance .AND. maxerrorA.LE.self%tolerance) THEN
			
			CALL OrbitsCopy(selfcopyA, self)
			self%error = maxerrorA
			converged = .true.

		END IF

	END DO

	!### clean up memory ###
	DEALLOCATE(selfcopyA%bodies,selfcopyB%bodies)

END SUBROUTINE

!### CalcMaxError will scan the particle array and calculate the largest ###
!### change in particle position in x and y ###
SUBROUTINE CalcMaxError (self, selfnew, maxerror)

	TYPE(Orbits) :: self, selfnew
	REAL*8, INTENT(INOUT) :: maxerror
	REAL*8 :: dx, dy
	INTEGER :: n

	!### initialize maxerror as something small ###
	maxerror = 1.d-24

	!### loop through each particle and calculate the change in x and y position ###
	DO n = 1, self%nbodies

		dx = DABS(self%bodies(n)%x - selfnew%bodies(n)%x)
		dy = DABS(self%bodies(n)%y - selfnew%bodies(n)%y)
		!PRINT '("DEBUG: dx = ",ES16.4," dy = ",ES16.4)', dx, dy

		maxerror = DMAX1(dx, dy, maxerror)

	END DO
	IF((maxerror/maxerror).NE.1.) STOP

END SUBROUTINE

!### OrbitStep will solve a 2nd order ODE using a 4th order Runge Kutta ###
!### scheme.  ###
SUBROUTINE OrbitStep (self)

	TYPE(Orbits), INTENT(INOUT) :: self
	INTEGER :: n, j
	REAL*8, DIMENSION (self%nbodies) :: ky1, ky2, ky3, ky4, kx1, kx2, kx3, kx4
	REAL*8, DIMENSION (self%nbodies) :: kvy1, kvy2, kvy3, kvy4, kvx1, kvx2, kvx3, kvx4
	REAL*8 :: t, sixth, xcm, ycm

	sixth = 1.d0 / 6.d0

	!### loop k1's through all of the bodies ###
	DO n = 1, self%nbodies

		kx1(n)  = 0.d0
		kvx1(n) = 0.d0
		ky1(n)  = 0.d0
		kvy1(n) = 0.d0

		!### loop through all other bodies ###
		kx1(n)  = self%dt * p_fderiv(self, self%bodies(n)%x, self%bodies(n)%vx, t)
		ky1(n)  = self%dt * p_fderiv(self, self%bodies(n)%y, self%bodies(n)%vy, t)
		DO j = 1, self%nbodies

		IF (n.ne.j) THEN

		!### construct the k1 coefficients ###
		kvx1(n) = kvx1(n) + self%dt * p_sderiv(self, self%bodies(n)%x, self%bodies(n)%y, self%bodies(j)%mass, self%bodies(j)%x, self%bodies(j)%y, t)
		kvy1(n) = kvy1(n) + self%dt * p_sderiv(self, self%bodies(n)%y, self%bodies(n)%x, self%bodies(j)%mass, self%bodies(j)%x, self%bodies(j)%y, t)

		END IF
		END DO

	END DO

	!### loop k2's through all of the bodies ###
	DO n = 1, self%nbodies

		kx2(n)  = 0.d0
		kvx2(n) = 0.d0
		ky2(n)  = 0.d0
		kvy2(n) = 0.d0

		!### loop through all other bodies ###
		kx2(n)  = self%dt * p_fderiv(self, self%bodies(n)%x + 0.5d0 * kx1(n), self%bodies(n)%vx + 0.5d0 * kvx1(n), t)
		ky2(n)  = self%dt * p_fderiv(self, self%bodies(n)%y + 0.5d0 * ky1(n), self%bodies(n)%vy + 0.5d0 * kvy1(n), t)
		DO j = 1, self%nbodies

		IF (n.ne.j) THEN

		!### construct the k2 coefficients ###
		kvx2(n) = kvx2(n) + self%dt * p_sderiv(self, self%bodies(n)%x + 0.5d0 * kx1(n), self%bodies(n)%y + 0.5d0 * ky1(n), self%bodies(j)%mass, self%bodies(j)%x + 0.5d0 * kx1(j), self%bodies(j)%y + 0.5d0 * ky1(j), t)
		kvy2(n) = kvy2(n) + self%dt * p_sderiv(self, self%bodies(n)%y + 0.5d0 * ky1(n), self%bodies(n)%x + 0.5d0 * kx1(n), self%bodies(j)%mass, self%bodies(j)%x + 0.5d0 * kx1(j), self%bodies(j)%y + 0.5d0 * ky1(j), t)

		END IF
		END DO

	END DO

	!### loop k3's through all of the bodies ###
	DO n = 1, self%nbodies

		kx3(n)  = 0.d0
		kvx3(n) = 0.d0
		ky3(n)  = 0.d0
		kvy3(n) = 0.d0

		!### loop through all other bodies ###
		kx3(n)  = self%dt * p_fderiv(self, self%bodies(n)%x + 0.5d0 * kx2(n), self%bodies(n)%vx + 0.5d0 * kvx2(n), t)
		ky3(n)  = self%dt * p_fderiv(self, self%bodies(n)%y + 0.5d0 * ky2(n), self%bodies(n)%vy + 0.5d0 * kvy2(n), t)
		DO j = 1, self%nbodies

		IF (n.ne.j) THEN

		!### construct the k3 coefficients ###
		kvx3(n) = kvx3(n) + self%dt * p_sderiv(self, self%bodies(n)%x + 0.5d0 * kx2(n), self%bodies(n)%y + 0.5d0 * ky2(n), self%bodies(j)%mass, self%bodies(j)%x + 0.5d0 * kx2(j), self%bodies(j)%y + 0.5d0 * ky2(j), t)
		kvy3(n) = kvy3(n) + self%dt * p_sderiv(self, self%bodies(n)%y + 0.5d0 * ky2(n), self%bodies(n)%x + 0.5d0 * kx2(n), self%bodies(j)%mass, self%bodies(j)%x + 0.5d0 * kx2(j), self%bodies(j)%y + 0.5d0 * ky2(j), t)

		END IF
		END DO

	END DO

	!### loop k4's through all of the bodies ###
	DO n = 1, self%nbodies

		kx4(n)  = 0.d0
		kvx4(n) = 0.d0
		ky4(n)  = 0.d0
		kvy4(n) = 0.d0

		!### loop through all other bodies ###
		kx4(n)  = self%dt * p_fderiv(self, self%bodies(n)%x + kx3(n), self%bodies(n)%vx + kvx3(n), t)
		ky4(n)  = self%dt * p_fderiv(self, self%bodies(n)%y + ky3(n), self%bodies(n)%vy + kvy3(n), t)
		DO j = 1, self%nbodies

		IF (n.ne.j) THEN

		!### construct the k4 coefficients ###
		kvx4(n) = kvx4(n) + self%dt * p_sderiv(self, self%bodies(n)%x + kx3(n), self%bodies(n)%y + ky3(n), self%bodies(j)%mass, self%bodies(j)%x + kx3(j), self%bodies(j)%y + ky3(j), t)
		kvy4(n) = kvy4(n) + self%dt * p_sderiv(self, self%bodies(n)%y + ky3(n), self%bodies(n)%x + kx3(n), self%bodies(j)%mass, self%bodies(j)%x + kx3(j), self%bodies(j)%y + ky3(j), t)

		END IF
		END DO

	END DO

	!### loop k4's through all of the bodies ###
	DO n = 1, self%nbodies

		!### using proper weighting evaluate the integral of the next point ###
		!PRINT '("N = ",I3,"has x = ",ES16.4," and y = ",ES16.4)', n, self%bodies(n)%x, self%bodies(n)%y
		self%bodies(n)%x  = self%bodies(n)%x + sixth * (kx1(n) + 2.d0 * kx2(n) + 2.d0 * kx3(n) + kx4(n))
		self%bodies(n)%y  = self%bodies(n)%y + sixth * (ky1(n) + 2.d0 * ky2(n) + 2.d0 * ky3(n) + ky4(n))
		self%bodies(n)%vx = self%bodies(n)%vx + sixth * (kvx1(n) + 2.d0 * kvx2(n) + 2.d0 * kvx3(n) + kvx4(n))
		self%bodies(n)%vy = self%bodies(n)%vy + sixth * (kvy1(n) + 2.d0 * kvy2(n) + 2.d0 * kvy3(n) + kvy4(n))
	END DO

	!### calculate the center of mass to see if it has moved ###
	xcm     = 0.d0
	ycm     = 0.d0

	!### loop through each of the bodies ###
	DO n = 1, self%nbodies
		!### sum up the center of mass ###
		xcm = xcm + self%bodies(n)%mass * self%bodies(n)%x
		ycm = ycm + self%bodies(n)%mass * self%bodies(n)%y
	END DO

	xcm = xcm / self%totmass
	ycm = ycm / self%totmass

	!### compare the new center of mass to the old and fix the fuckers ###
	xcm = self%xcm - xcm
	ycm = self%ycm - ycm

	!### update bodies position by this amount ###
	DO n = 1, self%nbodies
		self%bodies(n)%x = self%bodies(n)%x + xcm
		self%bodies(n)%y = self%bodies(n)%y + ycm
		!PRINT '("N = ",I3,"now has x = ",ES16.4," and y = ",ES16.4)', n, self%bodies(n)%x, self%bodies(n)%y
	END DO

END SUBROUTINE

END MODULE