!================================================================================
!  Subroutine that rescales the current velocities.  The rescales velocity
!  will have the spectrum that is defined in the input file via parameters
!  isp_type (spectrum type) and peak_wavenum (peak wavenumber).
!
!  The program is copies from a part of init_velocity.f90.  Some time in the
!  future we should make it one routine that is called in init_velocity.
!   
!  Time-stamp: <2010-01-25 17:01:55 (chumakov)>
!================================================================================

subroutine velocity_rescale

  use m_openmpi
  use m_parameters
  use m_io
  use m_fields
  use x_fftw

  implicit none

  integer :: i, j, k, n

  real,       allocatable   :: rr(:)
  real*8,     allocatable :: e_spec(:), e_spec1(:)
  integer *8, allocatable :: hits(:), hits1(:)

  integer   :: n_shell
  real*8    :: sc_rad1, sc_rad2

  real*8 :: wmag, wmag2, ratio, fac


  ! if Taylor-Green, return
  if (isp_type.eq.-1) return


!================================================================================
  allocate( e_spec(kmax), e_spec1(kmax), rr(nx+2), hits(kmax), hits1(kmax), stat=ierr)
  if (ierr.ne.0) stop "cannot allocate the init_velocity arrays"

  write(out,*) 'Rescaling the velocities'
  call flush(out)

!-------------------------------------------------------------------------------
!     Making the spectrum to be what is prescribed in the input file <...>.in 
!-------------------------------------------------------------------------------

  ! --- first get the energy spectrum (copied from m_stat.f90)

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

  ! zeroing out the arrays
  e_spec1 = zip
  e_spec = zip
  hits = 0
  hits1 = 0

  ! finding the total energy in each shell and number of hits in each shell
  ! storing them in arrays hits1 and e_spec1
  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
              hits1(n_shell) = hits1(n_shell) + 1
              e_spec1(n_shell) = e_spec1(n_shell) + &
                   fac * (fields(i,j,k,1)**2 + fields(i,j,k,2)**2 + fields(i,j,k,3)**2)
           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(hits1,hits,count,MPI_INTEGER8,MPI_SUM,0,MPI_COMM_TASK,mpi_err)
  count = kmax
  call MPI_REDUCE(e_spec1,e_spec,count,MPI_REAL8,MPI_SUM,0,MPI_COMM_TASK,mpi_err)

  ! now the master node counts the energy density in each shell
  if (myid.eq.0) then
     fac = four/three * PI / two
     do k = 1,kmax
        sc_rad1 = real(k,8) + half
        sc_rad2 = real(k,8) - half
        if (k.eq.1) sc_rad2 = 0.d0
        if (hits(k).gt.0) then
           e_spec(k) = e_spec(k) / hits(k) * fac * (sc_rad1**3 - sc_rad2**3)
        else
           e_spec(k) = zip
        end if
     end do
  end if

  ! broadcasting the spectrum
  count = kmax
  call MPI_BCAST(e_spec,count,MPI_REAL8,0,MPI_COMM_TASK,mpi_err)

!-------------------------------------------------------------------------------
!  Now make the spectrum to be as desired
!-------------------------------------------------------------------------------

  ! first, define the desired spectrum
  do k = 1,kmax

     wmag = real(k, 8)
     ratio = wmag / peak_wavenum

     if (isp_type.eq.0) then
        ! Plain Kolmogorov spectrum
        e_spec1(k) = wmag**(-5.d0/3.d0)

     else if (isp_type.eq.1) then
        ! Exponential spectrum
        e_spec1(k) =  ratio**3 / peak_wavenum * exp(-3.0D0*ratio)

     else if (isp_type.eq.3) then
        ! Von Karman spectrum
        fac = two * PI * ratio
        e_spec1(k) = fac**4 / (one + fac**2)**3

     else
        write(out,*) "ERROR: WRONG INITIAL SPECTRUM TYPE: ",isp_type
        call flush(out)
        stop

     end if
  end do

!  normalize it so it has the unit total energy
  e_spec1 = e_spec1 / sum(e_spec1(1:kmax))

  ! now go over all Fourier shells and multiply the velocities in a shell by
  ! the sqrt of ratio of the resired to the current spectrum
  do k = 1,nz
     do j = 1,ny
        do i = 1,nx+2

           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 .and. e_spec(n_shell) .gt. zip) then
              fields(i,j,k,1:3) = fields(i,j,k,1:3) * sqrt(e_spec1(n_shell)/e_spec(n_shell))
           else
              fields(i,j,k,1:3) = zip
           end if

        end do
     end do
  end do

  write(out,*) "Rescaled the velocities."
  call flush(out)

  ! deallocate work arrays
  deallocate(e_spec, e_spec1, rr, hits, hits1, stat=ierr)

  return
end subroutine velocity_rescale