module m_force

  use m_parameters

  implicit none

  integer*4 :: n_forced_nodes, n_forced_nodes_total

  ! coordinated of the forced nodes
  integer, allocatable :: ifn(:), jfn(:), kfn(:), k_shell(:)

contains

  subroutine m_force_init

    use m_parameters
    use x_fftw

    implicit none

    integer :: i, j, k, n, n_shell
    real*8 :: dissi_spec(kmax)
    real*8 :: eps_v_temp,eta_temp,uvar_temp,sctmp_temp,x_length_temp,u_prime_temp,tau_eta_temp

    ! if flow is not forced, return
    if (flow_type .ne. 1) return

    if (task.ne.'hydro') return

    select case (force_type)

    case (1:2)
       if (force_type == 1) then
       ! Machiels forcing (see article in PRL #79(18) p.3411)
          write(out,*) "Forcing #1: Machiels forcing - setting up"
       else
       ! Fully Controlled Deterministic Forcing Scheme (FC-DFS)
          write(out,*) "Forcing #2: FC-DFS - setting up"
       end if
       call flush(out)

       ! find out how many nodes are we forcing and book them
       n_forced_nodes = 0
       n_forced_nodes_total = 0
       do k = 1,nz
          do j = 1,ny
             do i = 1,nx
                n_shell = nint(sqrt(real(akx(i)**2 + aky(k)**2 + akz(j)**2, 4)))
                if (n_shell .gt. 0 .and. n_shell .le. kfmax) then
                   n_forced_nodes = n_forced_nodes + 1
                end if
             end do
          end do
       end do

       ! reducing to the master process to find out the total number of forced nodes
!!$       write(out,*) 'before reducing'; call flush(out)
       count = 1
       call MPI_REDUCE(n_forced_nodes, n_forced_nodes_total, count,  &
            MPI_INTEGER4,MPI_SUM,0,MPI_COMM_TASK,mpi_err)

       ! writing out the # of forced nodes
       write(out,'(a40,3x,i10)') 'Number of forced nodes for this process:',n_forced_nodes
       if (myid.eq.0) write(out,'(a40,3x,i10)') ' total number:',n_forced_nodes_total
       call flush(out)

       ! allocating arrays for the coordinates of the forced nodes
       allocate(ifn(n_forced_nodes), jfn(n_forced_nodes), kfn(n_forced_nodes), &
            k_shell(n_forced_nodes))

       ! filling up the arrays
       n = 0
       do k = 1,nz
          do j = 1,ny
             do i = 1,nx
                n_shell = nint(sqrt(real(akx(i)**2 + aky(k)**2 + akz(j)**2, 4)))
                if (n_shell .gt. 0 .and. n_shell .le. kfmax) then
                   n = n + 1
                   ifn(n) = i
                   jfn(n) = j
                   kfn(n) = k
                   k_shell(n) = n_shell
                end if
             end do
          end do
       end do

! Restarting setup
       allocate(model_e_spec(kmax)) ; model_e_spec=zip
       k0=(SQRT(two*PI))/t_length
       model_A=16.0*(SQRT(two/PI))*(t_u_prime**2.0)/k0
       do i=1,kmax
         k=real(i,8)
         model_e_spec(i)=model_A*((k/k0)**4.0)*exp(-2.0*((k/k0)**2.0))
       enddo

!!$       ! writing out the nodes
!!$       do n = 1,n_forced_nodes
!!$          write(out,"(3i4)") ifn(n),jfn(n),kfn(n)
!!$       end do
!!$       call flush(out)

    case default

       write(out,*) 'WRONG FORCE TYPE:',force_type
       write(out,*) 'STOPPING'
       call flush(out)
       stop

    end select


    return
  end subroutine m_force_init

!================================================================================

  subroutine force_velocity

    ! adding forcing to the arrays wrk(:,:,:,1:3) that already contain the RHS for velocities


    use m_openmpi
    use m_io
    use m_parameters
    use m_fields
    use m_work
    use x_fftw
    use m_stats

    implicit none

    integer :: i, j, k, n_shell, n

    real*8 :: fac, fac2
    real*8 :: kk
    real*8 :: mean_eps,tau_f
    real*8 :: kfmax8
    real*8 :: F_k(kmax),ff_k(kmax),eps_f(kmax)
    real*8 :: alpha,sum_F_Em,sum_F_Es
    real*8 :: energy_model

    select case (force_type)


    case (1)
       ! Machiels forcing (see article in PRL #79(18) p.3411)
       ! write(out,*) "Machiels forcing"; call flush(out)

       e_spec = zip
       e_spec1 = zip
       hits = 0
       hits1 = 0

       ! need this normalization factor because the FFT is unnormalized
       fac = one / real(nx*ny*nz_all)**2

       ! assembling the total energy in each shell
       do k = 1,nz
          do j = 1,ny
             do i = 1,nx
                n_shell = nint(sqrt(real(akx(i)**2 + aky(k)**2 + akz(j)**2, 4)))
                if (n_shell .gt. 0 .and. n_shell .le. kfmax) then
                   fac2 = fac * (fields(i,j,k,1)**2 + fields(i,j,k,2)**2 + fields(i,j,k,3)**2)
                   if (akx(i).eq.0.d0) fac2 = fac2 * 0.5d0
                   e_spec1(n_shell) = e_spec1(n_shell) + fac2
                end if
             end do
          end do
       end do
       ! reducing the number of hits and energy to two arrays on master node
       count = kfmax
       call MPI_REDUCE(e_spec1,e_spec,count,MPI_REAL8,MPI_SUM,0,MPI_COMM_TASK,mpi_err)

       ! getting the total energy in the region [0:kfmax] by integrating the spectrum
       if (myid.eq.0) energy = sum(e_spec(1:kfmax))

       ! broadcasting the current energy in the forcing range of wavenumbers
       count = 1
       call MPI_BCAST(energy,count,MPI_REAL8,0,MPI_COMM_TASK,mpi_err)

       ! now applying the forcing to the RHS for velocities (wrk(:,:,:,1:3))

       fac = FAMP / energy

       do n = 1,n_forced_nodes
          n_shell = k_shell(n)
          i = ifn(n)
          j = jfn(n)
          k = kfn(n)

          wrk(i,j,k,1) = wrk(i,j,k,1) + fac * fields(i,j,k,1)
          wrk(i,j,k,2) = wrk(i,j,k,2) + fac * fields(i,j,k,2)
          wrk(i,j,k,3) = wrk(i,j,k,3) + fac * fields(i,j,k,3)
          
       end do
       
!------------------------------------------------------------------------
! case (2)
! New forcing scheme added by ksj and edited by J. Kwon, 2014
! Ref. Flow, Turbulence and Combustion 73: 133-167, 2004
! by L. Guichard, J. Reveillon and R. Hauguel
! ------------------------------------------------------------------------
    case (2)
       ! Fully Controlled Deterministic Forcing Scheme (FC-DFS)

       e_spec = zip
       e_spec1 = zip
       hits = 0
       hits1 = 0

       ! need this normalization factor because the FFT is unnormalized
       fac = one / real(nx*ny*nz_all)**2

       ! assembling the total energy in each shell
       do k = 1,nz
          do j = 1,ny
             do i = 1,nx
                n_shell = nint(sqrt(real(akx(i)**2 + aky(k)**2 + akz(j)**2, 4)))
                if (n_shell .gt. 0 .and. n_shell .le. kmax) then
                   fac2 = fac * (fields(i,j,k,1)**2 + fields(i,j,k,2)**2 + fields(i,j,k,3)**2)
                   if (akx(i).eq.0.d0) fac2 = fac2 * 0.5d0
                   e_spec1(n_shell) = e_spec1(n_shell) + fac2
                end if
             end do
          end do
       end do

       ! reducing the number of hits and energy to two arrays on master node
       count = kmax
       call MPI_REDUCE(e_spec1,e_spec,count,MPI_REAL8,MPI_SUM,0,MPI_COMM_TASK,mpi_err)
       count = kmax
       call MPI_BCAST(e_spec,count,MPI_REAL8,0,MPI_COMM_TASK,mpi_err)

       ! getting the total energy by integrating the spectrum
       energy = sum(e_spec(1:kmax))

       mean_eps = zip
       tau_f = zip
       eps_f = zip
       F_k = zip
       ff_k = zip
       energy_model = sum(model_e_spec(1:kmax))

       do i = 1,kmax
          eps_f(i) = e_spec(i) * real(i**2,8) * two * nu
       end do

       mean_eps = sum(eps_f(1:kmax))
       tau_f=c_f*sqrt(nu/mean_eps)
       kfmax8=real(kfmax,8)

       do i=1,kfmax
         kk=real(i,8)
         F_k(i)=(1.-exp(-((kk-kfmax8)**2.)/(delta_k**2.)))/(1.-exp(-(kfmax8**2.)/(delta_k**2.)))
       enddo

       sum_F_Em=zip
       sum_F_Es=zip

       do i=1,kfmax
         sum_F_Em=sum_F_Em+F_k(i)*model_e_spec(i)
         sum_F_Es=sum_F_Es+F_k(i)*e_spec(i)
       enddo

       alpha=(energy_model-energy+tau_f*mean_eps+sum_F_Es)/sum_F_Em

       do i=1,kfmax
         ff_k(i)=F_k(i)/(two*e_spec(i))*(alpha*model_e_spec(i)-e_spec(i))
       enddo

       do n = 1,n_forced_nodes
          n_shell = k_shell(n)
          i = ifn(n)
          j = jfn(n)
          k = kfn(n)

          wrk(i,j,k,1) = wrk(i,j,k,1) + ff_k(n_shell)/tau_f * fields(i,j,k,1)
          wrk(i,j,k,2) = wrk(i,j,k,2) + ff_k(n_shell)/tau_f * fields(i,j,k,2)
          wrk(i,j,k,3) = wrk(i,j,k,3) + ff_k(n_shell)/tau_f * fields(i,j,k,3)
       end do

    case default
       write(out,*) "WRONG FORCE_TYPE:",force_type
       stop
    end select

    return
  end subroutine force_velocity


end module m_force