MODULE SPECTRAL_ARP_MOD

!   Author.
!   -------
!     Martin Leutbecher
!
!   Modifications.
!   -------------- 
!     Original      02-Jan-2009
!     M.Leutbecher  09-Jan-2009 parallelisation of rs2d/rs3d by zonal wavenumber
! ------------------------------------------------------------------

USE PARKIND1, ONLY: JPIM, JPRB
USE SPECTRAL_FIELDS_MOD
USE SPECTRAL_FIELDS_DATA,ONLY: MYMS, LELAM, MYLEVS
USE RANDOM_NUMBERS_MIX

IMPLICIT NONE
PRIVATE
PUBLIC SPECTRAL_ARP, ALLOCATE_ARP, DEALLOCATE_ARP, SET_ARP, SET_ARP2D, EVOLVE_ARP, SUM_ARPS, SET_SEED_ARP

TYPE SPECTRAL_ARP
  !
  !   spectral field sf containing the current state of the auto-regressive process
  !
  TYPE(SPECTRAL_FIELD) :: SF
  !
  !   spectral field sf containing the current state of sum of the 
  !   auto-regressive processes and switch if this should be activated
  !
  LOGICAL              :: LSUM
  TYPE(SPECTRAL_FIELD) :: SFSUM
  !
  !   master seed that generates all other seeds
  !
  INTEGER(KIND=JPIM)   :: ISEED0
  !
  !   initial master seed that generates all other seeds (set by allocate_arp only)
  !
  INTEGER(KIND=JPIM)   :: ISEED0T0
  !
  !   seeds (global variable) for local randomNumberStreams
  !     (depending on [level] spherical harmonics index and field
  !
  INTEGER(KIND=JPIM),       DIMENSION(:,:),   POINTER :: iseed2d
  INTEGER(KIND=JPIM),       DIMENSION(:,:,:), POINTER :: iseed3d
  !
  !   (local) randomNumberStreams for the noise used in the AR(1) processes
  !
  TYPE(randomNumberStream), DIMENSION(:,:),   POINTER :: rs2d
  TYPE(randomNumberStream), DIMENSION(:,:,:), POINTER :: rs3d
  !
  !   standard deviation as function of [level,] total wavenumber and field
  !
  REAL(KIND=JPRB),          DIMENSION(:,:),   POINTER :: sdev2d
  REAL(KIND=JPRB),          DIMENSION(:,:,:), POINTER :: sdev3d
  !   
  !   switch for clipping and
  !   amplitude for clipping the noise (aclip) and for clipping the AR1 (aclipar) as function of [level,] total wavenumber and field
  !
  LOGICAL :: LCLIP
  REAL(KIND=JPRB),          DIMENSION(:,:),   POINTER :: aclip2d
  REAL(KIND=JPRB),          DIMENSION(:,:,:), POINTER :: aclip3d
  REAL(KIND=JPRB),          DIMENSION(:,:),   POINTER :: aclipar2d
  REAL(KIND=JPRB),          DIMENSION(:,:,:), POINTER :: aclipar3d
  !
  !   correlation between successive steps as function of [level,] total wavenumber and field
  !
  REAL(KIND=JPRB),          DIMENSION(:,:),   POINTER :: phi2d
  REAL(KIND=JPRB),          DIMENSION(:,:,:), POINTER :: phi3d
END TYPE SPECTRAL_ARP


! ------------------------------------------------------------------
CONTAINS
! ------------------------------------------------------------------

SUBROUTINE ALLOCATE_ARP(ydarp,ksmax,kmsmax,ks3d,ks2d,ks1d,kgrib,kseed,ldclip,ldsum)

TYPE (SPECTRAL_ARP), INTENT(OUT) :: ydarp
INTEGER(KIND=JPIM), INTENT(IN) :: ksmax, kmsmax, ks3d, ks2d, ks1d, kgrib(ks3d+ks2d), kseed
LOGICAL, OPTIONAL :: ldclip, ldsum

INTEGER(KIND=JPIM) :: inflevg,inflevl,ispec2,ispec2g
INTEGER(KIND=JPIM) :: js, jfld, jl, il
INTEGER(KIND=JPIM) :: jm, im, ijsg, ijs, jnx, ilenx
INTEGER(KIND=JPIM) :: jjj
LOGICAL            :: llclip, llsum
!
!    initial checks
!
IF (LELAM) THEN
  CALL ABOR1('LAM version of allocate_arp not implemented yet.')
ENDIF
IF (PRESENT(LDCLIP)) THEN
  LLCLIP=LDCLIP
ELSE
  LLCLIP=.FALSE.
ENDIF
YDARP%LCLIP=LLCLIP

IF (PRESENT(LDSUM)) THEN
  LLSUM=LDSUM
ELSE
  LLSUM=.FALSE.
ENDIF
YDARP%LSUM=LLSUM

YDARP%ISEED0=kseed
YDARP%ISEED0T0=kseed ! first seed, not changed by set_seed_arp
!
!   allocate the spectral field component and
!        initialize AR(1) with 0.
!
CALL ALLOCATE_SPEC(ydarp%sf,ksmax,kmsmax,ks3d,ks2d,ks1d,kgrib)
ydarp%sf = 0.0_JPRB
IF (LLSUM) THEN
  CALL ALLOCATE_SPEC(ydarp%sfsum,ksmax,kmsmax,min(1,ks3d),min(1,ks2d),min(1,ks1d),kgrib)
  ydarp%sfsum = 0.0_JPRB
ENDIF
!
!   some shortcuts
!
inflevg=ydarp%sf%nflevg
inflevl=ydarp%sf%nflevl
ispec2 =ydarp%sf%nspec2
ispec2g=ydarp%sf%nspec2g
!
!   allocate the seeds (version with one random number generator for each m-wavenumber)
!
ALLOCATE(ydarp%ISEED2D(          0:ksmax, ks2d))
ALLOCATE(ydarp%ISEED3D( inflevg, 0:ksmax, ks3d))
!   allocate the random streams
!     (local to each task)
!
if (ks2d>0) ALLOCATE(YDARP%RS2D(         ydarp%sf%nump, KS2D) )
if (ks3d>0) ALLOCATE(YDARP%RS3D(INFLEVL, ydarp%sf%nump, KS3D) )
!
!   initialize random streams with master seed
!
call set_seed_arp( ydarp, kseed)
!
!   allocate arrays for standard deviation and correlation time scale
!
IF (ks2d>0) THEN
  ALLOCATE( ydarp%sdev2d( 0:ksmax, ks2d ))
  ALLOCATE( ydarp%phi2d(  0:ksmax, ks2d ))
  IF (llclip) THEN
    ALLOCATE(ydarp%aclip2d(0:ksmax, ks2d ))
    ALLOCATE(ydarp%aclipar2d(0:ksmax, ks2d ))
  ENDIF
ENDIF
IF (ks3d>0) THEN
  ALLOCATE( ydarp%sdev3d(inflevg,0:ksmax, ks3d ))
  ALLOCATE( ydarp%phi3d( inflevg,0:ksmax, ks3d ))
  IF (llclip) THEN 
    ALLOCATE( ydarp%aclip3d(inflevg,0:ksmax, ks3d ))
    ALLOCATE( ydarp%aclipar3d(inflevg,0:ksmax, ks3d ))
  ENDIF
ENDIF

END SUBROUTINE allocate_arp
!
!
!
SUBROUTINE SET_SEED_ARP(ydarp,kseed)

TYPE (SPECTRAL_ARP), INTENT(OUT) :: ydarp
INTEGER(KIND=JPIM), INTENT(IN) ::  kseed

INTEGER(KIND=JPIM) :: inflevg,inflevl,ispec2,ispec2g
INTEGER(KIND=JPIM) :: js, jfld, jl, il, ISMAX,is2d,is3d
INTEGER(KIND=JPIM) :: jm, im, ijsg, ijs, jnx, ilenx
INTEGER(KIND=JPIM) :: jjj
LOGICAL            :: llclip

TYPE(randomNumberStream) :: yl_seeds
REAL(KIND=JPRB), DIMENSION(:), ALLOCATABLE :: zseed
!
!    initial checks
!
IF (LELAM) THEN
  CALL ABOR1('LAM version of allocate_arp not implemented yet.')
ENDIF
YDARP%ISEED0=kseed
!
!   some shortcuts
!
inflevg=ydarp%sf%nflevg
inflevl=ydarp%sf%nflevl
ispec2 =ydarp%sf%nspec2
ispec2g=ydarp%sf%nspec2g
ismax  =ydarp%sf%nsmax
is2d   =ydarp%sf%ns2d
is3d   =ydarp%sf%ns3d
!
!   generate seeds (global)
!
CALL INITIALIZE_RANDOM_NUMBERS( KSEED, YL_SEEDS )
if (is2d>0) then
  ALLOCATE(ZSEED( (ismax+1)*IS2D ))
  CALL UNIFORM_DISTRIBUTION (ZSEED, YL_SEEDS)
  YDARP%ISEED2D= FLOOR(RESHAPE(ZSEED, (/(ismax+1), IS2D /) ) * (HUGE(1_JPIM)-1))
  DEALLOCATE(ZSEED)
endif
if (is3d>0) then
  ALLOCATE(ZSEED( INFLEVG*(ismax+1)*IS3D ))
  CALL UNIFORM_DISTRIBUTION (ZSEED, YL_SEEDS)
  YDARP%ISEED3D= FLOOR(RESHAPE(ZSEED, (/INFLEVG, ismax+1, IS3D /) ) * (HUGE(1_JPIM)-1))
  DEALLOCATE(ZSEED)
endif
!
!   (re-)initialise the random streams for 2D fields
!
IF (is2d>0) THEN
  DO JFLD=1, IS2D
    DO jm=1,ydarp%sf%nump
      im=myms(jm)
      CALL INITIALIZE_RANDOM_NUMBERS( ydarp%iseed2d(im,jfld), ydarp%rs2d(jm, jfld))
    ENDDO
  ENDDO
ENDIF
!
!   initialise the random streams for 3D fields
!
IF (is3d>0) THEN
  DO jfld=1, is3d
    DO jm=1,ydarp%sf%nump
      im=myms(jm)
      DO jl=1,ydarp%sf%nflevl
        il=mylevs(jl)
        CALL INITIALIZE_RANDOM_NUMBERS( ydarp%iseed3d(il,im,jfld), ydarp%rs3d(jl,jm, jfld))
      ENDDO
    ENDDO
  ENDDO
ENDIF

END SUBROUTINE SET_SEED_ARP
  
!
!
!
SUBROUTINE deallocate_arp(ydarp)
! ----------------------------------------------------------------------
! deallocate components of ydarp
! ----------------------------------------------------------------------
TYPE (SPECTRAL_ARP), INTENT(INOUT) :: ydarp

CALL deallocate_spec( ydarp%sf )
IF (ydarp%lsum) THEN
  CALL deallocate_spec( ydarp%sfsum )
ENDIF

IF (ASSOCIATED( YDARP%ISEED2D )) DEALLOCATE(YDARP%ISEED2D)
IF (ASSOCIATED( YDARP%ISEED3D )) DEALLOCATE(YDARP%ISEED3D)

IF (ASSOCIATED( YDARP%RS2D    )) DEALLOCATE(YDARP%RS2D)
IF (ASSOCIATED( YDARP%RS3D    )) DEALLOCATE(YDARP%RS3D)

IF (ASSOCIATED( YDARP%SDEV2D  )) DEALLOCATE(YDARP%SDEV2D)
IF (ASSOCIATED( YDARP%SDEV3D  )) DEALLOCATE(YDARP%SDEV3D)

IF (ASSOCIATED( YDARP%ACLIP2D  )) DEALLOCATE(YDARP%ACLIP2D)
IF (ASSOCIATED( YDARP%ACLIP3D  )) DEALLOCATE(YDARP%ACLIP3D)

IF (ASSOCIATED( YDARP%ACLIPAR2D  )) DEALLOCATE(YDARP%ACLIPAR2D)
IF (ASSOCIATED( YDARP%ACLIPAR3D  )) DEALLOCATE(YDARP%ACLIPAR3D)

IF (ASSOCIATED( YDARP%PHI2D   )) DEALLOCATE(YDARP%PHI2D)
IF (ASSOCIATED( YDARP%PHI3D   )) DEALLOCATE(YDARP%PHI3D)

END SUBROUTINE DEALLOCATE_ARP
!
!
!
SUBROUTINE set_arp2d(ydarp,sdev2d,phi2d,aclip2d)
! ----------------------------------------------------------------------
! set standard deviation and correlation 
! ----------------------------------------------------------------------

TYPE (SPECTRAL_ARP), INTENT(INOUT) :: ydarp
REAL(KIND=JPRB), DIMENSION(:,:),   INTENT(IN) :: sdev2d, phi2d
REAL(KIND=JPRB), DIMENSION(:,:),   OPTIONAL, INTENT(IN) :: aclip2d

INTEGER(KIND=JPIM) :: jn , js
REAL(KIND=JPRB)    :: zphi_fac0, zphi

ydarp%sdev2d=sdev2d
ydarp%phi2d =phi2d
if (present(aclip2d)) then
  ydarp%aclip2d =aclip2d
else
  if (ydarp%lclip) then
    ydarp%aclip2d=10._JPRB * ydarp%sdev2d
  endif
endif


if (ydarp%lclip) then
  do js=lbound(ydarp%sdev2d,dim=2), ubound(ydarp%sdev2d,dim=2)
    do jn=lbound(ydarp%sdev2d,dim=1), ubound(ydarp%sdev2d,dim=1)
      zphi=ydarp%phi2d( jn, js)
      zphi_fac0 = 1._JPRB / SQRT( 1._JPRB - zphi*zphi)
      ydarp%aclipar2d(jn,js) = zphi_fac0 * ydarp%aclip2d(jn,js)
    enddo
  enddo
endif

END SUBROUTINE set_arp2d
!
!
!
SUBROUTINE set_arp(ydarp,sdev2d,sdev3d,phi2d,phi3d,aclip2d,aclip3d)
! ----------------------------------------------------------------------
! set standard deviation and correlation 
! ----------------------------------------------------------------------

TYPE (SPECTRAL_ARP), INTENT(INOUT) :: ydarp
REAL(KIND=JPRB), DIMENSION(:,:),   INTENT(IN) :: sdev2d, phi2d
REAL(KIND=JPRB), DIMENSION(:,:,:), INTENT(IN) :: sdev3d, phi3d
REAL(KIND=JPRB), DIMENSION(:,:),    OPTIONAL,INTENT(IN) :: aclip2d
REAL(KIND=JPRB), DIMENSION(:,:,:),  OPTIONAL,INTENT(IN) :: aclip3d

INTEGER(KIND=JPIM) :: jn , js, jlev
REAL(KIND=JPRB)    :: zphi_fac0, zphi

ydarp%sdev2d=sdev2d; ydarp%sdev3d=sdev3d
ydarp%phi2d =phi2d;  ydarp%phi3d =phi3d
if (present(aclip2d)) then
  ydarp%aclip2d=aclip2d
else
  if (ydarp%lclip) then
    ydarp%aclip2d=10._JPRB * ydarp%sdev2d
  endif
endif

if (ydarp%lclip) then
  do js=lbound(ydarp%sdev2d,dim=2), ubound(ydarp%sdev2d,dim=2)
    do jn=lbound(ydarp%sdev2d,dim=1), ubound(ydarp%sdev2d,dim=1)
      zphi=ydarp%phi2d( jn, js)
      zphi_fac0 = 1._JPRB / SQRT( 1._JPRB - zphi*zphi)
      ydarp%aclipar2d(jn,js) = zphi_fac0 * ydarp%aclip2d(jn,js)
    enddo
  enddo
endif

if (present(aclip3d)) then
  ydarp%aclip3d=aclip3d
else
  if (ydarp%lclip) then
    ydarp%aclip3d=10._JPRB * ydarp%sdev3d
  endif
endif

if (ydarp%lclip) then
  do js=lbound(ydarp%sdev3d,dim=3), ubound(ydarp%sdev3d,dim=3)
    do jn=lbound(ydarp%sdev3d,dim=2), ubound(ydarp%sdev3d,dim=2)
      do jlev=lbound(ydarp%sdev3d,dim=1), ubound(ydarp%sdev3d,dim=1)
        zphi=ydarp%phi3d(jlev, jn, js)
        zphi_fac0 = 1._JPRB / SQRT( 1._JPRB - zphi*zphi)
        ydarp%aclipar3d(jlev,jn,js) = zphi_fac0 * ydarp%aclip3d(jlev,jn,js)
      enddo
    enddo
  enddo
endif

END SUBROUTINE set_arp
!
!
!
SUBROUTINE evolve_arp(ydarp,karp,ldinit)
! ----------------------------------------------------------------------
! evolve the autoregressive process by karp steps
! ----------------------------------------------------------------------
TYPE (SPECTRAL_ARP),          INTENT(INOUT) :: ydarp
INTEGER(KIND=JPIM), OPTIONAL, INTENT(IN)    :: karp
LOGICAL, OPTIONAL,            INTENT(IN)    :: ldinit
!
!   local variables
!
LOGICAL :: llinit
INTEGER(KIND=JPIM) :: isteps,jstep
INTEGER(KIND=JPIM) :: jmloc, im,jir,jn, inm, ilen1, ilen2, idxnar
INTEGER(KIND=JPIM) :: jlev,jfld,ilev
REAL(KIND=JPRB)    :: znoise, zphi, zphi_fac0, zlimit
REAL(KIND=JPRB), DIMENSION(:,:),   ALLOCATABLE :: znar2d
REAL(KIND=JPRB), DIMENSION(:,:,:), ALLOCATABLE :: znar3d
!
!
!
IF (PRESENT(ldinit)) THEN
  llinit=ldinit
ELSE
  llinit=.FALSE.
ENDIF
IF (PRESENT(karp)) THEN
  isteps=karp
ELSE
  isteps=1
ENDIF

ALLOCATE(znar2d(                 ydarp%sf%ns2d,2*(1+ydarp%sf%nsmax)))
ALLOCATE(znar3d(ydarp%sf%nflevl, ydarp%sf%ns3d,2*(1+ydarp%sf%nsmax)))

DO jstep=1, isteps
  !
  !   generate new noise and add to spectral field
  !
  DO jmloc=1,ydarp%sf%nump
    im=myms(jmloc)
    !
    !   generate all random numbers for M=IM and N=M, M+1, ..., NSMAX
    !    real and imaginary part (factor min(1,im))
    !
    ilen1=ydarp%sf%nsmax-im+1 ! number of N wavenumbers for M=im
    ilen2=(1+MIN(1,im))*ilen1 ! number of real and imaginary coefficients for M=im
    DO jfld=1,ydarp%sf%ns2d
      CALL GAUSSIAN_DISTRIBUTION(  znar2d(jfld,     1:ilen2), ydarp%rs2d(jmloc,     jfld) )
    ENDDO
    DO jfld=1,ydarp%sf%ns3d
      DO jlev=1,ydarp%sf%nflevl
        CALL GAUSSIAN_DISTRIBUTION(znar3d(jlev,jfld,1:ilen2), ydarp%rs3d(jlev,jmloc,jfld) )
      ENDDO
    ENDDO
    DO jir=0,MIN(1,im)
      DO jn=im,ydarp%sf%nsmax
        inm=ydarp%sf%nasm0(im)+(jn-im)*2+jir
        idxnar=1+(jn-im)*(1+MIN(1,im))+jir
        !
        !   2D fields
        !
        DO jfld=1,ydarp%sf%ns2d

          znoise = ydarp%sdev2d( jn, jfld ) * znar2d(jfld, idxnar)
          !
          !   clipping of forcing noise
          !
          IF (ydarp%lclip) THEN
            zlimit=ydarp%aclip2d( jn, jfld )
            znoise=MAX(-zlimit,MIN(zlimit, znoise))
          ENDIF
          zphi   = ydarp%phi2d(  jn, jfld )
          IF (llinit) THEN
            zphi_fac0 = 1._JPRB / SQRT( 1._JPRB - zphi*zphi) ! stdev ratio between AR(1) and
            !    forcing noise
            znoise = znoise * zphi_fac0 
          ENDIF
          ydarp%sf%sp2d( inm, jfld ) = zphi * ydarp%sf%sp2d( inm, jfld ) + znoise
          !
          !   clipping of AR(1)
          !
          IF (ydarp%lclip) THEN
            zlimit=ydarp%aclipar2d( jn, jfld )
            ydarp%sf%sp2d( inm, jfld ) = MAX(-zlimit, MIN(zlimit,ydarp%sf%sp2d( inm, jfld )))
          ENDIF
        ENDDO
        !
        !   3D fields
        !
        DO jfld=1,ydarp%sf%ns3d
          DO jlev=1,ydarp%sf%nflevl
            ilev=mylevs(jlev)
            znoise = ydarp%sdev3d( ilev, jn, jfld ) * znar3d(jlev,jfld,idxnar)
            !
            !   clipping of forcing noise
            !
            IF (ydarp%lclip) THEN
              zlimit=ydarp%aclip3d( ilev, jn, jfld )
              znoise=MAX(-zlimit,MIN(zlimit, znoise))
            ENDIF
            zphi   = ydarp%phi3d(  ilev, jn, jfld )
            IF (llinit) THEN
              zphi_fac0 = 1._JPRB / SQRT( 1._JPRB - zphi*zphi) ! stdev ratio between AR(1) and
                                                               !    forcing noise
              znoise = znoise * zphi_fac0 
            ENDIF
            ydarp%sf%sp3d( jlev, inm, jfld ) = zphi * ydarp%sf%sp3d( jlev, inm, jfld ) + znoise
            !
            !   clipping of AR(1)
            !
            IF (ydarp%lclip) THEN
              zlimit=ydarp%aclipar3d( ilev, jn, jfld )
              ydarp%sf%sp3d( jlev, inm, jfld ) = MAX(-zlimit, MIN(zlimit,ydarp%sf%sp3d( jlev, inm, jfld )))
            ENDIF
          ENDDO ! JLEV
        ENDDO ! JFLD
      ENDDO ! JN
    ENDDO ! JIR
  ENDDO ! JMLOC
ENDDO ! JSTEP

DEALLOCATE(ZNAR2D)
DEALLOCATE(ZNAR3D)

end SUBROUTINE evolve_arp
!
!
!
SUBROUTINE sum_arps(ydarp)
! ----------------------------------------------------------------------
! sum autoregressive processes
! ----------------------------------------------------------------------
TYPE (SPECTRAL_ARP),          INTENT(INOUT) :: ydarp
!
!   local variables
!
INTEGER(KIND=JPIM) :: jmloc, im,jir,jn, inm
INTEGER(KIND=JPIM) :: js,jl,jf,ispec2,iflevl
IF (.NOT.ydarp%lsum) THEN
  CALL abor1('ydarp%lsum has to be true in order to use sum_arps.')
ENDIF

ispec2=ydarp%sf%nspec2
iflevl=ydarp%sf%NFLEVL

if (ydarp%sf%ns2d > 0) then
  ydarp%sfsum%sp2d(1:ispec2,1)=ydarp%sf%sp2d(1:ispec2,1)
endif

!yt!$OMP PARALLEL DO SCHEDULE(STATIC) PRIVATE(jf,js)
DO jf=2,ydarp%sf%NS2D
  DO js=1,ispec2
    ydarp%sfsum%sp2d(js,1) = ydarp%sfsum%sp2d(js,1) + ydarp%sf%sp2d(js,jf)
  ENDDO
ENDDO
!yt!$OMP END PARALLEL DO

if (ydarp%sf%ns3d > 0) then
  ydarp%sfsum%sp3d(1:iflevl,1:ispec2,1)=ydarp%sf%sp3d(1:iflevl,1:ispec2,1)
endif
!yt!$OMP PARALLEL DO SCHEDULE(STATIC) PRIVATE(jf,js,jl)
DO jf=2,ydarp%sf%ns3d
  DO js=1,ispec2
    DO jl=1,iflevl
      ydarp%sfsum%sp3d(jl,js,1) = ydarp%sfsum%sp3d(jl,js,1) + ydarp%sf%sp3d(jl,js,jf)
    ENDDO
  ENDDO
ENDDO
!yt!$OMP END PARALLEL DO

end SUBROUTINE sum_arps

! ------------------------------------------------------------------

END MODULE SPECTRAL_ARP_MOD
