MODULE SPECTRAL_FIELDS_OPER_MOD

!   Purpose.
!   --------
!     This module contains the basic operators associated
!     to the SPECTRAL_FIELD type.
!
!   Author.
!   -------
!     Yannick Tremolet
!
!   Modifications.
!   -------------- 
!     Original    22-Jul-2004 Created from SPECTRAL_FIELDS_MOD.
!     C. Fischer  20-May-2005 Updated for Aladin
! ------------------------------------------------------------------

USE PARKIND1, ONLY: JPIM, JPRB
USE SPECTRAL_FIELDS_DATA, ONLY: SPECTRAL_FIELD, &
                              & MYLEVS, MYMS, LELAM 

IMPLICIT NONE
PRIVATE
PUBLIC ASSIGNMENT(=), OPERATOR(.EQV.), OPERATOR(.NEQV.), &
     & SELF_ADD, SELF_SUB, SELF_MUL

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

INTERFACE ASSIGNMENT (=)
MODULE PROCEDURE assign_sp_sp, assign_scalar_sp, assign_ar_sp, assign_sp_ar
END INTERFACE

INTERFACE OPERATOR (.EQV.)
MODULE PROCEDURE equiv_spec
END INTERFACE

INTERFACE OPERATOR (.NEQV.)
MODULE PROCEDURE nequiv_spec
END INTERFACE

INTERFACE SELF_ADD  ! What we really want here is ASSIGNMENT (+=)
MODULE PROCEDURE SELF_ADD_SP
END INTERFACE

INTERFACE SELF_SUB  ! What we really want here is ASSIGNMENT (-=)
MODULE PROCEDURE SELF_SUB_SP
END INTERFACE

INTERFACE SELF_MUL  ! What we really want here is ASSIGNMENT (*=)
MODULE PROCEDURE SELF_MUL_SP, SELF_MUL_SCAL
END INTERFACE

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

SUBROUTINE assign_scalar_sp(ydsp,pval)
TYPE (SPECTRAL_FIELD), INTENT(INOUT) :: ydsp
REAL(KIND=JPRB), INTENT(IN) :: pval

ydsp%sp2d(:,:)  =pval
ydsp%sp3d(:,:,:)=pval
IF (LELAM) ydsp%sp1d(:,:)  =pval

END SUBROUTINE assign_scalar_sp

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

SUBROUTINE assign_sp_sp(ydsp1,ydsp2)
TYPE (SPECTRAL_FIELD), INTENT(INOUT) :: ydsp1
TYPE (SPECTRAL_FIELD), INTENT(IN)    :: ydsp2
INTEGER(KIND=JPIM) :: jf,jm,im,ilen,js1,js2

IF (ydsp1.NEQV.ydsp2) CALL ABOR1('Incompatible spectral fields')

IF ((ydsp1%nsmax<ydsp2%nsmax).AND.(.NOT.LELAM)) THEN      ! Truncate in IFS
  DO jf=1,ydsp1%ns2d
    DO jm=1,ydsp1%NUMP
      im=MYMS(jm)   
      ilen=2*(ydsp1%NSMAX-im+1)-1
      js1=ydsp1%NASM0(im)
      js2=ydsp2%NASM0(im)
      ydsp1%sp2d(js1:js1+ilen,jf) = ydsp2%sp2d(js2:js2+ilen,jf)
    ENDDO
  ENDDO
  DO jf=1,ydsp1%ns3d
    DO jm=1,ydsp1%NUMP
      im=MYMS(jm)   
      ilen=2*(ydsp1%NSMAX-im+1)-1
      js1=ydsp1%NASM0(im)
      js2=ydsp2%NASM0(im)
      ydsp1%sp3d(:,js1:js1+ilen,jf) = ydsp2%sp3d(:,js2:js2+ilen,jf)
    ENDDO
  ENDDO 
ELSEIF ((ydsp1%nsmax>ydsp2%nsmax).AND.(.NOT.LELAM)) THEN  ! Pad with zeros in IFS
  ydsp1%sp2d(:,:) = 0.0_JPRB
  DO jf=1,ydsp2%ns2d
    DO jm=1,ydsp2%NUMP
      im=MYMS(jm)   
      ilen=2*(ydsp2%NSMAX-im+1)-1
      js1=ydsp1%NASM0(im)
      js2=ydsp2%NASM0(im)
      ydsp1%sp2d(js1:js1+ilen,jf) = ydsp2%sp2d(js2:js2+ilen,jf)
    ENDDO
  ENDDO
  ydsp1%sp3d(:,:,:) = 0.0_JPRB
  DO jf=1,ydsp2%ns3d
    DO jm=1,ydsp2%NUMP
      im=MYMS(jm)   
      ilen=2*(ydsp2%NSMAX-im+1)-1
      js1=ydsp1%NASM0(im)
      js2=ydsp2%NASM0(im)
      ydsp1%sp3d(:,js1:js1+ilen,jf) = ydsp2%sp3d(:,js2:js2+ilen,jf)
    ENDDO
  ENDDO 
ELSEIF ((((ydsp1%nsmax<ydsp2%nsmax).AND.(ydsp1%nmsmax<=ydsp2%nmsmax)).OR. &
 &((ydsp1%nsmax<=ydsp2%nsmax).AND.(ydsp1%nmsmax<ydsp2%nmsmax))) &
 &.AND.(LELAM)) THEN      ! Truncate in Aladin
  DO jf=1,ydsp1%ns2d
    DO jm=1,ydsp1%NUMP
      im=MYMS(jm)   
      ilen=ydsp1%ncpl4m(im)-1
      js1=ydsp1%NESM0(im)
      js2=ydsp2%NESM0(im)
      ydsp1%sp2d(js1:js1+ilen,jf) = ydsp2%sp2d(js2:js2+ilen,jf)
    ENDDO
  ENDDO
  DO jf=1,ydsp1%ns3d
    DO jm=1,ydsp1%NUMP
      im=MYMS(jm)   
      ilen=ydsp1%ncpl4m(im)-1
      js1=ydsp1%NESM0(im)
      js2=ydsp2%NESM0(im)
      ydsp1%sp3d(:,js1:js1+ilen,jf) = ydsp2%sp3d(:,js2:js2+ilen,jf)
    ENDDO
  ENDDO 
ELSEIF ((((ydsp1%nsmax>ydsp2%nsmax).AND.(ydsp1%nmsmax>=ydsp2%nmsmax)).OR. &
 &((ydsp1%nsmax>=ydsp2%nsmax).AND.(ydsp1%nmsmax>ydsp2%nmsmax))) &
 &.AND.(LELAM)) THEN  ! Pad with zeros in Aladin
  ydsp1%sp2d(:,:) = 0.0_JPRB
  DO jf=1,ydsp2%ns2d
    DO jm=1,ydsp2%NUMP
      im=MYMS(jm)   
      ilen=ydsp2%ncpl4m(im)-1
      js1=ydsp1%NESM0(im)
      js2=ydsp2%NESM0(im)
      ydsp1%sp2d(js1:js1+ilen,jf) = ydsp2%sp2d(js2:js2+ilen,jf)
    ENDDO
  ENDDO
  ydsp1%sp3d(:,:,:) = 0.0_JPRB
  DO jf=1,ydsp2%ns3d
    DO jm=1,ydsp2%NUMP
      im=MYMS(jm)   
      ilen=ydsp2%ncpl4m(im)-1
      js1=ydsp1%NESM0(im)
      js2=ydsp2%NESM0(im)
      ydsp1%sp3d(:,js1:js1+ilen,jf) = ydsp2%sp3d(:,js2:js2+ilen,jf)
    ENDDO
  ENDDO 
ELSEIF (((ydsp1%nsmax/=ydsp2%nsmax).OR.(ydsp1%nmsmax/=ydsp2%nmsmax)).AND. &
 & (LELAM)) THEN  ! Unfeasible cases left in Aladin
  CALL ABOR1('ASSIGN_SP_SP error: mismatch in Aladin spectral truncations')
ELSE                                   ! Copy
  ydsp1%sp2d(:,:)  =ydsp2%sp2d(:,:)
  ydsp1%sp3d(:,:,:)=ydsp2%sp3d(:,:,:)
ENDIF

! mean wind always copied if present in input
DO jf=1,ydsp2%ns1d
  ydsp1%sp1d(:,jf)=ydsp2%sp1d(:,jf)
ENDDO

END SUBROUTINE assign_sp_sp

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

SUBROUTINE assign_ar_sp(ydsp,pflat)
TYPE (SPECTRAL_FIELD), INTENT(INOUT) :: ydsp
REAL(KIND=JPRB), INTENT(IN) :: pflat(:)
INTEGER(KIND=JPIM) :: i2d,i3d,ishape2d(2),ishape3d(3)
INTEGER(KIND=JPIM) :: ii,im,ioff,il,jsg,jf,jm,js,jn,jir,jl,ilen

IF (SIZE(pflat)==ydsp%nspsizel) THEN             ! local array (IFS and Aladin)

  i2d=ydsp%ns2d*ydsp%nspec2
  i3d=ydsp%nflevl*ydsp%ns3d*ydsp%nspec2
  ishape2d(1)=ydsp%nspec2
  ishape2d(2)=ydsp%ns2d
  ishape3d(1)=ydsp%nflevl
  ishape3d(2)=ydsp%nspec2
  ishape3d(3)=ydsp%ns3d
  ydsp%sp2d(:,:)  =RESHAPE(pflat(    1:    i2d),ishape2d)
  ydsp%sp3d(:,:,:)=RESHAPE(pflat(i2d+1:i2d+i3d),ishape3d)

ELSEIF ((SIZE(pflat)==ydsp%nspsizeg).AND.(.NOT.LELAM)) THEN         ! global array in IFS

  ioff=0
  DO jf=1,ydsp%ns2d
    DO jm=1,ydsp%NUMP
      im=MYMS(jm)
      jsg=ydsp%NASM0G(im)-1
      js =ydsp%NASM0 (im)-1
      ilen=2*(ydsp%nsmax-im+1)
      ydsp%sp2d(js+1:js+ilen,jf)=pflat(ioff+jsg+1:ioff+jsg+ilen)
    ENDDO
    ioff=ioff+ydsp%nspec2g
  ENDDO
  ioff=ydsp%ns2g*ydsp%nspec2g
  DO jf=1,ydsp%ns3d
    DO jm=1,ydsp%NUMP
      im=MYMS(jm)
      DO jn=im,ydsp%nsmax
        DO jir=0,1
          jsg=ydsp%NASM0G(im)+(jn-im)*2+jir
          js =ydsp%NASM0 (im)+(jn-im)*2+jir
          ii=ioff+(jsg-1)*ydsp%nflevg
          DO jl=1,ydsp%nflevl
            il=MYLEVS(jl)
            ydsp%sp3d(jl,js,jf)=pflat(ii+il)
          ENDDO
        ENDDO
      ENDDO
    ENDDO
    ioff=ioff+ydsp%nspec2g*ydsp%nflevg
  ENDDO

ELSEIF ((SIZE(pflat)==ydsp%nspsizeg).AND.(LELAM)) THEN         ! global array in Aladin

  ioff=0
  DO jf=1,ydsp%ns2d
    DO jm=1,ydsp%NUMP
      im=MYMS(jm)
      jsg=ydsp%NESM0G(im)-1
      js =ydsp%NESM0 (im)-1
      ilen=ydsp%ncpl4m(im)
      ydsp%sp2d(js+1:js+ilen,jf)=pflat(ioff+jsg+1:ioff+jsg+ilen)
    ENDDO
    ioff=ioff+ydsp%nspec2g
  ENDDO
  ioff=ydsp%ns2g*ydsp%nspec2g
  DO jf=1,ydsp%ns3d
    DO jm=1,ydsp%NUMP
      im=MYMS(jm)
      DO jn=1,ydsp%ncpl4m(im)
        jsg=ydsp%NESM0G(im)+jn-1
        js =ydsp%NESM0 (im)+jn-1
        ii=ioff+(jsg-1)*ydsp%nflevg
        DO jl=1,ydsp%nflevl
          il=MYLEVS(jl)
          ydsp%sp3d(jl,js,jf)=pflat(ii+il)
        ENDDO
      ENDDO
    ENDDO
    ioff=ioff+ydsp%nspec2g*ydsp%nflevg
  ENDDO

ELSE
  CALL ABOR1('assign_ar_sp: cannot assign array to spectral field')
ENDIF

!cf mean wind must be copied apart, but for conversions from spectral_field
!cf to control_vector and back, 0. seems to be the most appropriate value

IF (LELAM) ydsp%sp1d(:,:)=0.0_JPRB

END SUBROUTINE assign_ar_sp

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

SUBROUTINE assign_sp_ar(pflat,ydsp)
REAL(KIND=JPRB), INTENT(OUT) :: pflat(:)
TYPE (SPECTRAL_FIELD), INTENT(IN) :: ydsp
INTEGER(KIND=JPIM) :: i2d,i3d,ishape2d(1),ishape3d(1)

IF (SIZE(pflat)==ydsp%nspsizel) THEN         ! local array (IFS and Aladin)

  i2d=ydsp%ns2d*ydsp%nspec2
  i3d=ydsp%nflevl*ydsp%ns3d*ydsp%nspec2
  ishape2d(1)=i2d
  ishape3d(1)=i3d
  pflat(    1:    i2d)=RESHAPE(ydsp%sp2d(:,:)  ,ishape2d)
  pflat(i2d+1:i2d+i3d)=RESHAPE(ydsp%sp3d(:,:,:),ishape3d)

ELSEIF (SIZE(pflat)==ydsp%nspsizeg) THEN     ! global array
  CALL ABOR1('Cannot assign spectral field to global array. '// &
           & 'Please use ALLGATHER_SPEC')
ELSE
  CALL ABOR1('Vector wrong length')
ENDIF

!cf information from ydsp%sp1d is not retrieved here - external copy required

END SUBROUTINE assign_sp_ar

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

LOGICAL FUNCTION equiv_spec(ydsp1,ydsp2)
TYPE(SPECTRAL_FIELD), INTENT(IN) :: ydsp1
TYPE(SPECTRAL_FIELD), INTENT(IN) :: ydsp2
LOGICAL :: ll
INTEGER(KIND=JPIM) :: jf

ll = .true.
ll = ll .AND. (ydsp1%ns2g ==ydsp2%ns2g)
ll = ll .AND. (ydsp1%ns3d ==ydsp2%ns3d)
ll = ll .AND. (ydsp1%ns1d ==ydsp2%ns1d)
IF (ll) THEN
  DO jf=1,ydsp1%ns2g+ydsp1%ns3d
    ll = ll .AND. (ydsp1%ngrib(jf)==ydsp2%ngrib(jf))
  ENDDO
ENDIF

equiv_spec=ll

END FUNCTION equiv_spec

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

LOGICAL FUNCTION nequiv_spec(ydsp1,ydsp2)
TYPE(SPECTRAL_FIELD), INTENT(IN) :: ydsp1
TYPE(SPECTRAL_FIELD), INTENT(IN) :: ydsp2

nequiv_spec = .NOT. equiv_spec(ydsp1,ydsp2)

END FUNCTION nequiv_spec

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

SUBROUTINE SELF_ADD_SP(YDA,YDB)
TYPE(SPECTRAL_FIELD), INTENT(INOUT) :: YDA
TYPE(SPECTRAL_FIELD), INTENT(IN)    :: YDB
INTEGER(KIND=JPIM) :: js,jl,jf

IF (YDA.NEQV.YDB) CALL ABOR1('Incompatible spectral fields')
IF (YDA%NSMAX/=YDB%NSMAX) CALL ABOR1('Incompatible spectral resolutions')
IF ((YDA%NMSMAX/=YDB%NMSMAX).AND.LELAM) CALL ABOR1('Incompatible spectral resolutions LAM')

!yt!$OMP PARALLEL DO SCHEDULE(STATIC) PRIVATE(jf,js)
DO jf=1,YDA%NS2D
  DO js=1,YDA%NSPEC2
    YDA%SP2D(js,jf) = YDA%SP2D(js,jf) + YDB%SP2D(js,jf)
  ENDDO
ENDDO
!yt!$OMP END PARALLEL DO

!yt!$OMP PARALLEL DO SCHEDULE(STATIC) PRIVATE(jf,js,jl)
DO jf=1,YDA%NS3D
  DO js=1,YDA%NSPEC2
    DO jl=1,YDA%NFLEVL
      YDA%SP3D(jl,js,jf) = YDA%SP3D(jl,js,jf) + YDB%SP3D(jl,js,jf)
    ENDDO
  ENDDO
ENDDO
!yt!$OMP END PARALLEL DO

IF (LELAM) THEN
  DO jf=1,YDA%NS1D
    DO jl=1,YDA%NFLEVL
      YDA%SP1D(jl,jf) = YDA%SP1D(jl,jf) + YDB%SP1D(jl,jf)
    ENDDO
  ENDDO
ENDIF

END SUBROUTINE SELF_ADD_SP
! ------------------------------------------------------------------

SUBROUTINE SELF_SUB_SP(YDA,YDB)
TYPE(SPECTRAL_FIELD), INTENT(INOUT) :: YDA
TYPE(SPECTRAL_FIELD), INTENT(IN)    :: YDB
INTEGER(KIND=JPIM) :: js,jl,jf

IF (YDA.NEQV.YDB) CALL ABOR1('Incompatible spectral fields')
IF (YDA%NSMAX/=YDB%NSMAX) CALL ABOR1('Incompatible spectral resolutions')
IF ((YDA%NMSMAX/=YDB%NMSMAX).AND.LELAM) CALL ABOR1('Incompatible spectral resolutions LAM')

!yt!$OMP PARALLEL DO SCHEDULE(STATIC) PRIVATE(jf,js)
DO jf=1,YDA%NS2D
  DO js=1,YDA%NSPEC2
    YDA%SP2D(js,jf) = YDA%SP2D(js,jf) - YDB%SP2D(js,jf)
  ENDDO
ENDDO
!yt!$OMP END PARALLEL DO

!yt!$OMP PARALLEL DO SCHEDULE(STATIC) PRIVATE(jf,js,jl)
DO jf=1,YDA%NS3D
  DO js=1,YDA%NSPEC2
    DO jl=1,YDA%NFLEVL
      YDA%SP3D(jl,js,jf) = YDA%SP3D(jl,js,jf) - YDB%SP3D(jl,js,jf)
    ENDDO
  ENDDO
ENDDO
!yt!$OMP END PARALLEL DO

IF (LELAM) THEN
  DO jf=1,YDA%NS1D
    DO jl=1,YDA%NFLEVL
      YDA%SP1D(jl,jf) = YDA%SP1D(jl,jf) - YDB%SP1D(jl,jf)
    ENDDO
  ENDDO
ENDIF

END SUBROUTINE SELF_SUB_SP
! ------------------------------------------------------------------

SUBROUTINE SELF_MUL_SP(YDA,YDB)
TYPE(SPECTRAL_FIELD), INTENT(INOUT) :: YDA
TYPE(SPECTRAL_FIELD), INTENT(IN)    :: YDB
INTEGER(KIND=JPIM) :: js,jl,jf

IF (YDA.NEQV.YDB) CALL ABOR1('Incompatible spectral fields')
IF (YDA%NSMAX/=YDB%NSMAX) CALL ABOR1('Incompatible spectral resolutions')
IF ((YDA%NMSMAX/=YDB%NMSMAX).AND.LELAM) CALL ABOR1('Incompatible spectral resolutions LAM')

!yt!$OMP PARALLEL DO SCHEDULE(STATIC) PRIVATE(jf,js)
DO jf=1,YDA%NS2D
  DO js=1,YDA%NSPEC2
    YDA%SP2D(js,jf) = YDA%SP2D(js,jf) * YDB%SP2D(js,jf)
  ENDDO
ENDDO
!yt!$OMP END PARALLEL DO

!yt!$OMP PARALLEL DO SCHEDULE(STATIC) PRIVATE(jf,js,jl)
DO jf=1,YDA%NS3D
  DO js=1,YDA%NSPEC2
    DO jl=1,YDA%NFLEVL
      YDA%SP3D(jl,js,jf) = YDA%SP3D(jl,js,jf) * YDB%SP3D(jl,js,jf)
    ENDDO
  ENDDO
ENDDO
!yt!$OMP END PARALLEL DO

IF (LELAM) THEN
  DO jf=1,YDA%NS1D
    DO jl=1,YDA%NFLEVL
      YDA%SP1D(jl,jf) = YDA%SP1D(jl,jf) * YDB%SP1D(jl,jf)
    ENDDO
  ENDDO
ENDIF

END SUBROUTINE SELF_MUL_SP
! ------------------------------------------------------------------

SUBROUTINE SELF_MUL_SCAL(YDA,PB)
TYPE(SPECTRAL_FIELD), INTENT(INOUT) :: YDA
REAL(KIND=JPRB),      INTENT(IN)    :: PB
INTEGER(KIND=JPIM) :: js,jl,jf

!yt!$OMP PARALLEL DO SCHEDULE(STATIC) PRIVATE(jf,js)
DO jf=1,YDA%NS2D
  DO js=1,YDA%NSPEC2
    YDA%SP2D(js,jf) = PB * YDA%SP2D(js,jf)
  ENDDO
ENDDO
!yt!$OMP END PARALLEL DO

!yt!$OMP PARALLEL DO SCHEDULE(STATIC) PRIVATE(jf,js,jl)
DO jf=1,YDA%NS3D
  DO js=1,YDA%NSPEC2
    DO jl=1,YDA%NFLEVL
      YDA%SP3D(jl,js,jf) = PB * YDA%SP3D(jl,js,jf)
    ENDDO
  ENDDO
ENDDO
!yt!$OMP END PARALLEL DO

IF (LELAM) THEN
  DO jf=1,YDA%NS1D
    DO jl=1,YDA%NFLEVL
      YDA%SP1D(jl,jf) = PB * YDA%SP1D(jl,jf)
    ENDDO
  ENDDO
ENDIF

END SUBROUTINE SELF_MUL_SCAL
! ------------------------------------------------------------------

END MODULE SPECTRAL_FIELDS_OPER_MOD
