!!******************************************************************************
!!
!!  This file is part of the AMUN source code, a program to perform
!!  Newtonian or relativistic magnetohydrodynamical simulations on uniform or
!!  adaptive mesh.
!!
!!  Copyright (C) 2008-2024 Grzegorz Kowal <grzegorz@amuncode.org>
!!
!!  This program is free software: you can redistribute it and/or modify
!!  it under the terms of the GNU General Public License as published by
!!  the Free Software Foundation, either version 3 of the License, or
!!  (at your option) any later version.
!!
!!  This program is distributed in the hope that it will be useful,
!!  but WITHOUT ANY WARRANTY; without even the implied warranty of
!!  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
!!  GNU General Public License for more details.
!!
!!  You should have received a copy of the GNU General Public License
!!  along with this program.  If not, see <http://www.gnu.org/licenses/>.
!!
!!******************************************************************************
!!
!! module: RANDOM
!!
!!  This module provides subroutines to random number generators.
!!
!!  References:
!!
!!    [1] David Blackman and Sebastiano Vigna,
!!        "Scrambled linear pseudorandom number generators", 2019
!!        http://vigna.di.unimi.it/ftp/papers/ScrambledLinear.pdf
!!    [2] http://prng.di.unimi.it/
!!
!!******************************************************************************
!
module random

  implicit none

! random generator type
!
  character(len = 32), save :: gentype = "same"

! variables to store seeds for random generator
!
  integer                                     , save :: kp = 0, np = 0
  integer                                     , save :: nseeds, lseed
  integer(kind=8)                             , save :: state
  integer(kind=8), dimension(:,:), allocatable, save :: seeds

! integer to real conversion parameters
!
  real(kind=8), parameter :: i64tor8 = 2.0d+00**(-53)

  private

  public :: initialize_random, finalize_random
  public :: nseeds, reset_seeds, set_seeds, get_seeds, gentype
  public :: randuni, randsym, randnorz

  contains
!
!===============================================================================
!
! subroutine INITIALIZE_RANDOM:
! ----------------------------
!
!   Subroutine initializes random number generator.
!
!   Arguments:
!
!     gtype    - the generator type ('same' or 'random');
!     nthreads - the number of threads;
!     nthread  - the thread number;
!     nproc    - the process number;
!     status   - the call status;
!
!===============================================================================
!
  subroutine initialize_random(gtype, nthreads, nthread, nproc, status)

    implicit none

    character(len=*), intent(in)  :: gtype
    integer         , intent(in)  :: nthreads, nthread, nproc
    integer         , intent(out) :: status

    integer :: i

    integer(kind=8), dimension(4) :: s

!-------------------------------------------------------------------------------
!
    status = 0

! set local copies of the thread and process numbers,
! determine the number of seeds
!
    kp     = nthread
    np     = nproc
    nseeds = nthreads
    lseed  = nseeds - 1

! allocate seeds
!
    allocate(seeds(4, 0:lseed), stat = status)

    if (status == 0) then

! prepare seeds depending on the type of generator
!
      write(gentype,'(a)') trim(gtype)
      select case(gentype)
      case('random')
        state = 1234567890
        s(1)  = splitmix64()
        s(2)  = splitmix64()
        s(3)  = splitmix64()
        s(4)  = splitmix64()
        do i = 1, nseeds * np
          call jump(s(:))
        end do
        seeds(:,0) = s(:)
        do i = 1, lseed
          call jump(s(:))
          seeds(:,i) = s(:)
        end do
      case default
        state      = 1234567890
        seeds(1,:) = splitmix64()
        seeds(2,:) = splitmix64()
        seeds(3,:) = splitmix64()
        seeds(4,:) = splitmix64()
      end select

    end if

!-------------------------------------------------------------------------------
!
  end subroutine initialize_random
!
!===============================================================================
!
! subroutine FINALIZE_RANDOM:
! --------------------------
!
!   Subroutine releases memory allocated by random number generator variables.
!
!   Arguments:
!
!     status - the call status;
!
!===============================================================================
!
  subroutine finalize_random(status)

    implicit none

    integer, intent(out) :: status

!-------------------------------------------------------------------------------
!
    status = 0

    if (allocated(seeds)) deallocate(seeds, stat=status)

!-------------------------------------------------------------------------------
!
  end subroutine finalize_random
!
!===============================================================================
!
! subroutine RESET_SEEDS:
! ----------------------
!
!   Subroutine resets the seeds, according to the generator type.
!
!   Arguments:
!
!     rngtype - the generator type ('same' or 'random');
!
!===============================================================================
!
  subroutine reset_seeds(rngtype)

    implicit none

    character(len=*), intent(in) :: rngtype

    integer :: i

    integer(kind=8), dimension(4) :: s

!-------------------------------------------------------------------------------
!
    select case(trim(rngtype))
    case('random')
      state = 1234567890
      s(1)  = splitmix64()
      s(2)  = splitmix64()
      s(3)  = splitmix64()
      s(4)  = splitmix64()
      do i = 1, nseeds * np
        call jump(s(:))
      end do
      seeds(:,0) = s(:)
      do i = 1, lseed
        call jump(s(:))
        seeds(:,i) = s(:)
      end do
      gentype = 'random'
    case default
      state      = 1234567890
      seeds(1,:) = splitmix64()
      seeds(2,:) = splitmix64()
      seeds(3,:) = splitmix64()
      seeds(4,:) = splitmix64()
      gentype = 'same'
    end select

!-------------------------------------------------------------------------------
!
  end subroutine reset_seeds
!
!===============================================================================
!
! subroutine SET_SEEDS:
! --------------------
!
!   Subroutine sets the seeds.
!
!   Arguments:
!
!     n   - the number of seeds;
!     sds - the new seeds;
!     gen - the flag indication if the seeds have to be regenerated;
!
!===============================================================================
!
  subroutine set_seeds(n, sds, gen)

    implicit none

    integer                        , intent(in) :: n
    integer(kind=8), dimension(4,n), intent(in) :: sds
    logical                        , intent(in) :: gen

    integer :: i

    integer(kind=8), dimension(4) :: s

!-------------------------------------------------------------------------------
!
! set the seeds only if the input array and seeds have the same sizes
!
    if (gen .or. n /= nseeds) then
      select case(gentype)
      case('random')
        state = 1234567890
        s(1)  = splitmix64()
        s(2)  = splitmix64()
        s(3)  = splitmix64()
        s(4)  = splitmix64()
        do i = 1, nseeds * np
          call jump(s(:))
        end do
        seeds(:,0) = s(:)
        do i = 1, lseed
          call jump(s(:))
          seeds(:,i) = s(:)
        end do
      case default
        state      = 1234567890
        seeds(1,:) = splitmix64()
        seeds(2,:) = splitmix64()
        seeds(3,:) = splitmix64()
        seeds(4,:) = splitmix64()
      end select
    else
      seeds(:,:) = sds(:,:)
    end if

!-------------------------------------------------------------------------------
!
  end subroutine set_seeds
!
!===============================================================================
!
! subroutine GET_SEEDS:
! --------------------
!
!   Subroutine returns the seeds.
!
!   Arguments:
!
!     sds - the return seeds array;
!
!===============================================================================
!
  subroutine get_seeds(sds)

    implicit none

    integer(kind=8), dimension(4,nseeds), intent(out) :: sds

!-------------------------------------------------------------------------------
!
    if (size(sds,1) == 4 .and. size(sds,2) == nseeds) then
      sds(:,:) = seeds(:,:)
    end if

!-------------------------------------------------------------------------------
!
  end subroutine get_seeds
!
!===============================================================================
!
! function SPLITMIX64:
! -------------------
!
!   Pseudorandom Number Generator for floating point numbers based on
!   SplitMix64 algorithm [1], used for initialization of xoshiro256p()
!   generator.
!
!   References:
!
!     [1] David Blackman and Sebastiano Vigna,
!         "Scrambled linear pseudorandom number generators", 2019
!         http://vigna.di.unimi.it/ftp/papers/ScrambledLinear.pdf
!
!===============================================================================
!
  integer(kind=8) function splitmix64()

    implicit none

    integer(kind=8), parameter :: a = -7046029254386353131_8
    integer(kind=8), parameter :: b = -4658895280553007687_8
    integer(kind=8), parameter :: c = -7723592293110705685_8

!-------------------------------------------------------------------------------
!
    state      = state + a
    splitmix64 = state
    splitmix64 = ieor(splitmix64, ishft(splitmix64, -30)) * b
    splitmix64 = ieor(splitmix64, ishft(splitmix64, -27)) * c
    splitmix64 = ieor(splitmix64, ishft(splitmix64, -31))

    return

!-------------------------------------------------------------------------------
!
  end function splitmix64
!
!===============================================================================
!
! function XOSHIRO256P:
! --------------------
!
!   Pseudorandom Number Generator for floating point numbers based on
!   xoshiro256+ algorithm [1]. It return a 256-bit random integer.
!
!   References:
!
!     [1] David Blackman and Sebastiano Vigna,
!         "Scrambled linear pseudorandom number generators", 2019
!         http://vigna.di.unimi.it/ftp/papers/ScrambledLinear.pdf
!
!===============================================================================
!
  integer(kind=8) function xoshiro256p()

    implicit none

    integer(kind=8) :: temp

!-------------------------------------------------------------------------------
!
    xoshiro256p = ishft(seeds(1,kp) + seeds(4,kp), -11)
    temp        = ishft(seeds(2,kp), 17)
    seeds(3,kp) = ieor(seeds(3,kp), seeds(1,kp))
    seeds(4,kp) = ieor(seeds(4,kp), seeds(2,kp))
    seeds(2,kp) = ieor(seeds(2,kp), seeds(3,kp))
    seeds(1,kp) = ieor(seeds(1,kp), seeds(4,kp))
    seeds(3,kp) = ieor(seeds(3,kp), temp)
    seeds(4,kp) = ior(ishft(seeds(4,kp), 45), ishft(seeds(4,kp), -19))

    return

!-------------------------------------------------------------------------------
!
  end function xoshiro256p
!
!===============================================================================
!
! subroutine JUMP:
! ---------------
!
!   The jump function for the xoshiro256p() generator used to generate
!   non-overlapping seed subsequences for parallel computations.
!
!   Arguments:
!
!     sds - temporary seeds array;
!
!   References:
!
!     [1] David Blackman and Sebastiano Vigna,
!         "Scrambled linear pseudorandom number generators", 2019
!         http://vigna.di.unimi.it/ftp/papers/ScrambledLinear.pdf
!
!===============================================================================
!
  subroutine jump(sds)

    implicit none

    integer(kind=8), dimension(4), intent(inout) :: sds

    integer(kind=8), dimension(4) :: s
    integer(kind=8)               :: temp
    integer                       :: i, b

    integer(kind=8), dimension(4), parameter :: jmp = [ 1733541517147835066_8, &
                                                       -3051731464161248980_8, &
                                                       -6244198995065845334_8, &
                                                        4155657270789760540_8 ]

!-------------------------------------------------------------------------------
!
    s(:) = 0_8
    do i = 1, 4
      do b = 1, 64
        if (iand(jmp(i), ishft(1_8, b)) /= 0_8) then
          s(1) = ieor(s(1), sds(1))
          s(2) = ieor(s(2), sds(2))
          s(3) = ieor(s(3), sds(3))
          s(4) = ieor(s(4), sds(4))
        end if
        temp   = ishft(sds(2), 17)
        sds(3) = ieor(sds(3), sds(1))
        sds(4) = ieor(sds(4), sds(2))
        sds(2) = ieor(sds(2), sds(3))
        sds(1) = ieor(sds(1), sds(4))
        sds(3) = ieor(sds(3), temp)
        sds(4) = ior(ishft(sds(4), 45), ishft(sds(4), -19))
      end do
    end do
    sds(:) = s(:)

    return

!-------------------------------------------------------------------------------
!
  end subroutine jump
!
!===============================================================================
!
! function RANDUNI:
! ----------------
!
!   Function generates uniformly distributed floating point real random number
!   in range 0.0..1.0 using xoshiro256p() generator.
!
!===============================================================================
!
  real(kind=8) function randuni()

    implicit none

!-------------------------------------------------------------------------------
!
    randuni = i64tor8 * xoshiro256p()

    return

!-------------------------------------------------------------------------------
!
  end function randuni
!
!===============================================================================
!
! function RANDSYM:
! ----------------
!
!   Function generates uniformly distributed floating point random number
!   in range -1.0..1.0 using xoshiro256p() generator.
!
!===============================================================================
!
  real(kind=8) function randsym()

    implicit none

!-------------------------------------------------------------------------------
!
    randsym = 2.0d+00 * (i64tor8 * xoshiro256p() - 5.0d-01)

    return

!-------------------------------------------------------------------------------
!
  end function randsym
!
!===============================================================================
!
! function RANDNORZ:
! -----------------
!
!   Function generates complex floating point number with a normal
!   distribution (μ = 0, σ = 1) using Marsaglia polar method [1].
!
!   References:
!
!     [1]  Marsaglia, G. & Bray, T. A.,
!         "A Convenient Method for Generating Normal Variables",
!         SIAM Review, vol. 6, no. 3, pp. 260-264,
!         https://doi.org/10.1137/1006063
!
!===============================================================================
!
  complex(kind=8) function randnorz()

    implicit none

    logical      :: c
    real(kind=8) :: u, v, s

!-------------------------------------------------------------------------------
!
    c = .true.
    do while(c)
      u = randsym()
      v = randsym()
      s = u**2 + v**2
      c = s <= 0.0d+00 .or. s >= 1.0d+00
    end do

    randnorz = sqrt(- 2.0d+00 * log(s) / s) * cmplx(u, v, kind=8)

    return

!-------------------------------------------------------------------------------
!
  end function randnorz

!===============================================================================
!
end module random