module time_integration
        use params
        use forces
        use energy
        use vitalrates

        implicit none
      contains
              subroutine time_average_array(tavrgvel,n,vel,pn)
                      ! Compute the average instantenous velocity of the particles.
                      integer :: n,pn
                      real(kind=rkd) :: tavrgvel(n),vel(pn,2)

                      tavrgvel(1:n-1)=tavrgvel(2:n)
                      tavrgvel(n)=sum(sqrt(vel(:,1)**2 +vel(:,2)**2))/real(n,kind=rkd)
              end subroutine time_average_array

              function time_average(vel,n)
                      ! Compute the time average of the velocity.
                      integer :: n
                      real(kind=rkd) :: vel(n)
                      real(kind=rkd) :: time_average

                      time_average=sum(vel)/real(n,kind=rkd)
              end function time_average

              subroutine comp_f(pop,pn,res,rn,alc,an)
                      integer :: pn,rn,an
                      real(kind=rkd),intent(inout) :: pop(pn,21),res(rn,11),alc(an,11)

                      call reset_forces(pop,pn,res,rn,alc,an)
                      call comp_fgrav(pop,pn,res,rn,alc,an)
                      call comp_flj(pop,pn,res,rn,alc,an)
                      call comp_fhl(pop,pn,res,rn,alc,an)
              end subroutine comp_f

              subroutine comp_langevin(pop,pn,res,rn,alc,an,dt)
                      ! Implement the Langevin dynamics.
                      integer :: pn,rn,an
                      real(kind=rkd),intent(inout) :: pop(pn,21),res(rn,11),alc(an,11)
                      real(kind=rkd),intent(in) :: dt
                      real(kind=rkd) :: m,vel(pn,2)
                      real(kind=rkd) :: lgam,lran,tmp1,tmp2,tmp3,tmp4
                      integer :: i,n

                      ! Generate the bias corrected gaussian random set.
                      vel=corrected_gaussian(pop,pn,res,rn,alc,an,dt)

                      do i=1,pn

                      n=int(pop(i,1))

                      !
                      ! If we have a blank space, then skip the computations.
                      !
                      if (n<1) CYCLE

                      ! This concerns only the individual person particles.
                      !if (n < 3) then
                      !  cycle
                      !end if

                      ! Fetch the variable values.
                      lgam=lpars(n,1)
                      lran=lpars(n,2)
                      m=gpars(n)

                      ! We have a friction force against the velocity
                      ! and a random kick on the xy -plane.

                      ! Compute the Langevin force.
                      pop(i,6)=pop(i,6) - lgam*pop(i,4) + lran*vel(i,1)

                      ! Compute the Langevin force.
                      pop(i,7)=pop(i,7) - lgam*pop(i,5) + lran*vel(i,2)
                      end do

              end subroutine comp_langevin

              function corrected_gaussian(pop,pn,res,rn,alc,an,dt)
                      ! Velocity correction to prevent the movement of the
                      ! center of the mass.
                      integer :: pn,rn,an
                      real(kind=rkd),intent(inout) :: pop(pn,21),res(rn,11),alc(an,11)
                      real(kind=rkd) :: corrected_gaussian(pn,2)
                      real(kind=rkd),intent(in) :: dt
                      real(kind=rkd) :: m
                      real(kind=rkd) :: lgam,lran,tmp1,tmp2,tmp3,tmp4
                      integer :: i,n
                      real(kind=rkd) :: meanv(3) ! Velocity correction

                      ! Initalize the counter.
                      meanv(1:3)=0

                      ! Compute the sum of the velocities.
                      do i=1,pn

                      n=int(pop(i,1))

                      !
                      ! If we have a blank space, then skip the computations.
                      !
                      if (n<1) CYCLE

                      ! Increase particle n for the velocity correction.
                      meanv(3)=meanv(3)+1.0D+00

                      ! We sample the random number from the normal
                      ! distribution using the method from Stickler B,
                      ! Schachinger E: Basic Concepts in Computational
                      ! Physics, 2ed, 2016, Springer, p.198.

                      ! Add the velocity to the sum
                      ! for the velocity correction.
                      ! Do the process on x-direction.
                      call random_number(tmp1)
                      call random_number(tmp2)
                      tmp3=cos(2*pi*tmp2)*sqrt(-2*log(tmp1))
                      call random_number(tmp1)
                      call random_number(tmp2)
                      tmp4=sin(2*pi*tmp2)*sqrt(-2*log(tmp1))

                      corrected_gaussian(i,1)=tmp3*tmp4
                      meanv(1)=meanv(1)+tmp3*tmp4

                      ! Repeat the process on y-dirextion.
                      call random_number(tmp1)
                      call random_number(tmp2)
                      tmp3=cos(2*pi*tmp2)*sqrt(-2*log(tmp1))
                      call random_number(tmp1)
                      call random_number(tmp2)
                      tmp4=sin(2*pi*tmp2)*sqrt(-2*log(tmp1))

                      corrected_gaussian(i,2)=tmp3*tmp4
                      meanv(2)=meanv(2)+tmp3*tmp4

                      end do

                      ! Calculate the mean bias (x,y).
                      meanv(1)=meanv(1)/meanv(3)
                      meanv(2)=meanv(2)/meanv(3)

                      ! Substract the mean bias.
                      do i=1,pn

                      n=int(pop(i,1))

                      !
                      ! If we have a blank space, then skip the computations.
                      !
                      if (n<1) CYCLE

                      ! Substract the mean bias from the velocities.
                      corrected_gaussian(i,1)=corrected_gaussian(i,1)-meanv(1)
                      corrected_gaussian(i,2)=corrected_gaussian(i,2)-meanv(2)
                      end do
              end function corrected_gaussian

              subroutine comp_x(pop,pn,res,rn,alc,an,dt)
                      integer :: pn,rn,an
                      real(kind=rkd),intent(inout) :: pop(pn,21),res(rn,11),alc(an,11)
                      real(kind=rkd) :: m
                      real(kind=rkd),intent(in) :: dt
                      integer :: i,n

                      do i=1,pn
                      n=int(pop(i,1))

                      !
                      ! If we have a blank space, then skip the computations.
                      !
                      if (n<1) CYCLE

                      m=gpars(n)
                      pop(i,2)=pop(i,2) + dt * (pop(i,4) + 0.5D+00 / m * pop(i,6) * dt)
                      pop(i,3)=pop(i,3) + dt * (pop(i,5) + 0.5D+00 / m * pop(i,7) * dt)
                      if (abs(pop(i,2)) > sbox(1)/2.0D+00 .or. abs(pop(i,3)) > sbox(2)/2.0D+00) then
                        pop(i,1)=0.0D+00
                      end if
                      end do
              end subroutine comp_x

              subroutine comp_x_quenching(pop,pn,res,rn,alc,an,dt)
                      integer :: pn,rn,an
                      real(kind=rkd),intent(inout) :: pop(pn,21),res(rn,11),alc(an,11)
                      real(kind=rkd),intent(in) :: dt
                      real(kind=rkd) :: f
                      integer :: i,n

                      do i=1,pn
                      n=int(pop(i,1))

                      !
                      ! If we have a blank space, then skip the computations.
                      !
                      if (n<1) CYCLE

                      ! Compute the total force.
                      f = sqrt(pop(i,6)**2 + pop(i,7)**2)
                      ! Compute the unit component (sin or cos) of the total force vector.
                      ! The velocity is set to the component fraction of the time step.
                      pop(i,2) = pop(i,2) + pop(i,6)/f * dt * quench_s
                      pop(i,3) = pop(i,3) + pop(i,7)/f * dt * quench_s
                      end do
              end subroutine comp_x_quenching

              subroutine comp_v(pop,pn,res,rn,alc,an,dt)
                      integer :: pn,rn,an
                      real(kind=rkd),intent(inout) :: pop(pn,21),res(rn,11),alc(an,11)
                      real(kind=rkd),intent(in) :: dt
                      real(kind=rkd) :: m
                      integer :: i,n

                      do i=1,pn
                      n=int(pop(i,1))

                      !
                      ! If we have a blank space, then skip the computations.
                      !
                      if (n<1) CYCLE

                      m=gpars(n)
                      pop(i,4)=pop(i,4) + dt * 0.5D+00 / m * (pop(i,6)+pop(i,8))
                      pop(i,5)=pop(i,5) + dt * 0.5D+00 / m * (pop(i,7)+pop(i,9))
                      end do
              end subroutine comp_v

              function center_of_mass(pop,pn,res,rn,alc,an,dt)
                      integer :: pn,rn,an
                      real(kind=rkd),intent(inout) :: pop(pn,21),res(rn,11),alc(an,11)
                      real(kind=rkd) :: m,mtot,w(2)
                      real(kind=rkd) :: center_of_mass(2)
                      real(kind=rkd),intent(in) :: dt
                      integer :: i,n

                      w(1:2)=0
                      mtot=0

                      do i=1,pn
                      n=pop(i,1)
                      if (n<1) CYCLE

                      m=gpars(n)
                      w(1)=w(1)+m*pop(i,2)
                      w(2)=w(2)+m*pop(i,3)
                      mtot=mtot+m
                      end do

                      center_of_mass(1)=w(1)/mtot
                      center_of_mass(2)=w(2)/mtot
              end function center_of_mass

              subroutine comp_pressure(pop,pn,res,rn,alc,an,dt)
                ! Compute the pressure tensor PV, and calculate pressure P.
                ! The simulation is 2D, so the pressure is actually
                ! the surface pressure P=P/V=P/A.
                ! The method is from Hoover Wm. G: Molecular Dynamics (Springer 1986)
                ! and Thmopson, Plimpton, Mattson (2019).
                ! The order of the components is xx,yy,xy. Because yx is not
                ! independent, we save time and space by not outputting it.
                      integer :: pn,rn,an
                      real(kind=rkd),intent(inout) :: pop(pn,21),res(rn,11),alc(an,11)
                      real(kind=rkd) :: m,v(2)
                      real(kind=rkd),intent(in) :: dt
                      integer :: i,n

                      do i=1,pn
                      n=pop(i,1)
                      !
                      ! Skip, if the array cell is empty.
                      !
                      if (n<1) CYCLE
                      !
                      ! Compute the tensors kinetic parts: xx, yy, xy.
                      !
                      m=gpars(n)
                      v=pop(i,4:5)
                      pop(i,11) = abs(pop(i,11)) + abs(m*v(1)*v(1))
                      pop(i,12) = abs(pop(i,12)) + abs(m*v(2)*v(2))
                      pop(i,13) = abs(pop(i,13)) + abs(m*v(1)*v(2))
                      !
                      ! Compute the kinetic energy.
                      !
                      pop(i,14) = ekinetic(m,v) 
                      end do
              end subroutine comp_pressure

              subroutine output_p_test(pop,pn,res,rn,alc,an,t)
                      integer :: pn,rn,an,i,j
                      real(kind=rkd),intent(inout) :: pop(pn,21),res(rn,11),alc(an,11),t
                      real(kind=rkd) :: ekin,egr,elj,ehl

                      !
                      ! Compute the kinetic and potential energy
                      ! to verify the sanity of the system.
                      !
                      ekin = comp_ekin(pop,pn,res,rn,alc,an)
                      egr = comp_pgrav(pop,pn,res,rn,alc,an)
                      elj = comp_plj(pop,pn,res,rn,alc,an)
                      ehl = comp_phl(pop,pn,res,rn,alc,an)

                      if (t < 0.0D+00) then
                              print*,"t",char(9),"ekin",char(9),"egrav",char(9),&
                                      "elj",char(9),"ehl",char(9),&
                                      "x1",char(9),"y1",char(9),"vx1",char(9),"vy1",char(9),&
                                      "fx1",char(9),"fy1",char(9),"x2",char(9),"y2",char(9),"vx2",char(9),&
                                      "vy2",char(9),"fx2",char(9),"fy2"!,char(9),"x3",char(9),"x3",char(9),&
                                      !"vx3",char(9),"vy3",char(9),"fx3",char(9),"fy3",char(9),"x4",char(9),&
                                      !"x4",char(9),"vx4",char(9),"vy4",char(9),"fx4",char(9),"fy4"
                      else
                              print*,t,ekin,egr,elj,ehl,(pop(i,2:7),i=1,pn)
                      end if
              end subroutine output_p_test

              subroutine output_p(pop,pn,res,rn,alc,an,t)
                      integer :: pn,rn,an,i,j,n
                      real(kind=rkd),intent(inout) :: pop(pn,21),res(rn,11),alc(an,11),t

                      !
                      ! Count all the particles.
                      !
                      j=1
                      n=0
                      do i=1,pn
                      if (int(pop(i,1))>0) then
                              n=n+1
                      end if
                      end do

                      !
                      ! Print the xyz header: the number of particles and the comment line.
                      !
                      if (n>99999) then
                        print '(I6)',n
                      else if (n>9999) then
                        print '(I5)',n
                      else if (n>999) then
                        print '(I4)',n
                      else if (n>99) then
                        print '(I3)',n
                      elseif (n>9) then
                        print '(I2)',n
                      else
                        print '(I1)',n
                      end if
                      print '(A5,E16.8)',"Time=",t
                      
                      !
                      ! Print the individual particles.
                      !
                      do i=1,pn
                      if (int(pop(i,1)) < 1) CYCLE
                      !
                      ! Augment the array by adding 3rd dimension, though it is not used.
                      !
                      ! The format:
                      ! 1: Particle type [1]
                      ! 2,3,4: position (x,y,z=0) [2,3]
                      ! 5,6,7: (x,y,z=0) [4,5]
                      ! 8,9,10: (x,y,z=0) [6,7]
                      ! 11,12: [8,9]
                      ! 13: ID [10]
                      ! 14:
                      ! 15:
                      ! 16:
                      ! 17:
                      ! 18: Age [15]
                      ! 19:
                      ! 20:
                      ! 21:
                      ! 22:
                      ! 23: Age at giving birth
                      ! 24: Bond target (negative: target is the mother) [21]


                      print '(A4,X,F16.4,X,F16.4,X,F16.4,X,E16.4,X,E16.4,X,E16.4,X,E16.4,X,E16.4,X,&
                        E16.4,E16.4,E16.4,I8,E16.4,E16.4,E16.4,E16.4,I8,E16.4,E16.4,E16.4,E16.4,E16.4,E16.4)', &
                        form_ident(pop(i,:)),pop(i,2),pop(i,3),0.0D+00,pop(i,4),pop(i,5),0.0D+00, &
                        pop(i,6),pop(i,7),0.0D+00,pop(i,8:9),int(pop(i,10)),pop(i,11:14),int(pop(i,15)),pop(i,16:21)
                      end do
              end subroutine output_p

              subroutine output_popg(pop,pn,res,rn,alc,an,t)
                      integer :: pn,rn,an,i,j,n
                      real(kind=rkd),intent(inout) :: pop(pn,21),res(rn,11),alc(an,11),t

                      !
                      ! Open the state freeze file.
                      !
                      open(21,file='popg_out.txt',action='WRITE')
                      !
                      ! Count all the particles.
                      !
                      j=1
                      n=0
                      do i=1,pn
                      if (int(pop(i,1))>0) then
                              n=n+1
                      end if
                      end do
                      !
                      ! Print the popg.txt header: the number of particles.
                      !
                      write (21,*) n
                      !
                      ! Print the individual particles.
                      !
                      do i=1,pn
                      if (int(pop(i,1)) < 1) CYCLE
                      ! Print the position (x,y,z=0) and the velocities (vx,vy,vz=0).
                      !write '(A4,X,F16.4,X,F16.4,X,F16.4,F20.4,X,F20.4,X,F20.4)', &
                      write (21,'(I1,X,E16.8,X,E16.8,X,E16.8,X,E16.8,X,E16.8,X,E16.8,X,E16.8,X,E16.8,X,&
                        I8,E16.8,E16.8,E16.8,E16.8,E16.8,E16.8,E16.8,E16.8,E16.8,E16.8,I8)')  &
                        int(pop(i,1)), pop(i,2:9), int(pop(i,10)), pop(i,11:20), int(pop(i,21)) 
                      end do
                      close(21)
              end subroutine output_popg

              subroutine time_integrate_vsv(pop,pn,res,rn,alc,an,t,dt,tmax)
                      !
                      ! Velocity-Störmer-Verlet Method
                      !

                      integer :: pn,rn,an,outc=0
                      real(kind=rkd),intent(inout) :: pop(pn,21),res(rn,11),alc(an,11)
                      real(kind=rkd),intent(inout) :: t
                      real(kind=rkd),intent(in) :: dt,tmax
                      real(kind=rkd) :: e,exy(2),p,vitc=0.0D+00

                      !
                      ! Compute the initial forces.
                      !
                      if (pn>1) then
                              call comp_f(pop,pn,res,rn,alc,an)
                      end if
                      !
                      ! We use an indefinite DO loop that can be more machine
                      ! efficient compared to WHILE loop (Barlow & Barnett, Metcalf & Reid).
                      ! 
                      do
                      outc = outc + 1
                      vitc = vitc + 1.0D+00 ! Vitc is a counter, but real to save in casting.
                      t=t+dt
                      if (t > tmax) EXIT
                      !
                      ! If we start a new simulation, we need to remove
                      ! excess repulsions that are the results of a
                      ! random population. This prevents the catastrophic
                      ! initial explosion.
                      !
                      if (t < quench_t) then
                        call comp_x_quenching(pop,pn,res,rn,alc,an,dt)
                      else
                        call distsum(pop,pn)
                        if (vitc > vitcc) then
                                call popdeath(pop,pn,res,rn,alc,an,vitcc,dt)
                                call popbirth(pop,pn,res,rn,alc,an,vitcc,dt)
                                call distzero(pop,pn)
                                vitc=0.0D+00
                        end if
                        call comp_langevin(pop,pn,res,rn,alc,an,dt)
                        call comp_x(pop,pn,res,rn,alc,an,dt)
                      end if
                      !
                      ! The forces can be calculated only if there are
                      ! more than 1 particles. 
                      !
                      if (pn>1) then
                              call comp_f(pop,pn,res,rn,alc,an)
                      end if
                      !
                      ! The initial quenching, see above.
                      !
                      if (t >= quench_t) then
                        call comp_v(pop,pn,res,rn,alc,an,dt)
                      end if
                      !
                      ! Print a line of data for every outcc:th iteration.
                      !
                      if (tavrgcc > 0) then
                              call time_average_array(tavrg,tavrgcc,pop(1:1,4:5),pn)
                      end if
                      !
                      ! Output during the simulationa
                      !
                      if (outc > outcc) then
                        !
                        ! If tavrgcc is set, then print the time average,
                        ! else print the simulation snapshot.
                        !
                        if(tavrgcc > 0) then
                          print*,time_average(tavrg,tavrgcc)
                        else
                          call comp_pressure(pop,pn,res,rn,alc,an,t)
                          call output_p(pop,pn,res,rn,alc,an,t)
                        end if
                        outc = 0
                      end if
                      end do
                      !
                      ! Print the snapshot of the final step.
                      call output_popg(pop,pn,res,rn,alc,an,t)
              end subroutine time_integrate_vsv

end module time_integration
