MODULE euler_scheme_mod
  USE field_mod
  IMPLICIT NONE
  PRIVATE

  TYPE(t_field),POINTER,SAVE,PUBLIC :: f_geopot(:), f_q(:), &
       f_rhodz(:), f_mass(:), f_massm1(:), f_massm2(:), f_dmass(:), & ! current and previous time steps + tendency of mass,
       f_phis(:), f_ps(:),f_psm1(:), f_psm2(:), f_dps(:), &           ! column-integrated mass
       f_u(:),f_um1(:),f_um2(:), f_du(:), f_W(:), &                   ! 'horizontal' and vertical (NH) momentum
       f_theta_rhodz(:),f_theta_rhodzm1(:),f_theta_rhodzm2(:), f_dtheta_rhodz(:), & ! mass-weighted theta/entropy
       f_hflux(:), f_wflux(:), f_hfluxt(:), f_wfluxt(:), &            ! accumulated mass fluxes
       f_du_slow(:,:), f_du_fast(:,:), &                              ! slow/fast trend arrays 
       f_dmass_slow(:,:), f_dps_slow(:,:), f_dtheta_rhodz_slow(:,:), &! for HEVI time scheme 
       f_dPhi_slow(:,:), f_dPhi_fast(:,:), &                          ! geopotential tendencies (NH)
       f_dW_slow(:,:), f_dW_fast(:,:)                                 ! vertical momentum tendencies (NH)

  INTEGER, PARAMETER, PUBLIC :: explicit=1, hevi=2, euler=1, rk4=2, mlf=3, rk25=4, ark23=6, ark33=7

  INTEGER,SAVE, PUBLIC :: nb_stage, matsuno_period, scheme, scheme_family
  !$OMP THREADPRIVATE(nb_stage, matsuno_period, scheme, scheme_family)

  PUBLIC :: euler_scheme, accumulate_fluxes, legacy_to_DEC, DEC_to_legacy
CONTAINS

  SUBROUTINE Euler_scheme(with_dps,fluxt_zero)
    USE icosa
    USE disvert_mod
    USE omp_para
    USE trace
    LOGICAL :: with_dps
    LOGICAL, OPTIONAL :: fluxt_zero(ndomain) ! set to .TRUE. to start accumulating fluxes in time
    REAL(rstd),POINTER :: ps(:), dps(:)
    REAL(rstd),POINTER :: u(:,:) , du(:,:)
    REAL(rstd),POINTER :: mass(:,:), dmass(:,:)
    REAL(rstd),POINTER :: theta_rhodz(:,:), dtheta_rhodz(:,:)
    REAL(rstd),POINTER :: hflux(:,:),wflux(:,:),hfluxt(:,:),wfluxt(:,:)
    INTEGER :: ind
    INTEGER :: i,j,ij,l
    CALL trace_start("Euler_scheme")  

    DO ind=1,ndomain
       IF (.NOT. assigned_domain(ind)) CYCLE
       CALL swap_dimensions(ind)
       CALL swap_geometry(ind)

       IF(with_dps) THEN ! update ps/mass
          IF(caldyn_eta==eta_mass) THEN ! update ps
             ps=f_ps(ind) ; dps=f_dps(ind) ;              
             IF (is_omp_first_level) THEN
                !$SIMD
                DO ij=ij_begin,ij_end
                   ps(ij)=ps(ij)+dt*dps(ij)
                ENDDO
             ENDIF
          ELSE ! update mass
             mass=f_mass(ind) ; dmass=f_dmass(ind) ;              
             DO l=ll_begin,ll_end
                !$SIMD
                DO ij=ij_begin,ij_end
                   mass(ij,l)=mass(ij,l)+dt*dmass(ij,l)
                ENDDO
             END DO
          END IF

          hflux=f_hflux(ind);     hfluxt=f_hfluxt(ind)
          wflux=f_wflux(ind);     wfluxt=f_wfluxt(ind)
          CALL accumulate_fluxes(hflux,wflux,hfluxt,wfluxt,dt,fluxt_zero(ind))
       END IF ! update ps/mass

       u=f_u(ind) ; theta_rhodz=f_theta_rhodz(ind)
       du=f_du(ind) ; dtheta_rhodz=f_dtheta_rhodz(ind)

       DO l=ll_begin,ll_end
          !$SIMD
          DO ij=ij_begin,ij_end
             u(ij+u_right,l)=u(ij+u_right,l)+dt*du(ij+u_right,l)
             u(ij+u_lup,l)=u(ij+u_lup,l)+dt*du(ij+u_lup,l)
             u(ij+u_ldown,l)=u(ij+u_ldown,l)+dt*du(ij+u_ldown,l)
             theta_rhodz(ij,l)=theta_rhodz(ij,l)+dt*dtheta_rhodz(ij,l)
          ENDDO
       ENDDO
    ENDDO

    CALL trace_end("Euler_scheme")  

  END SUBROUTINE Euler_scheme

  SUBROUTINE accumulate_fluxes(hflux,wflux, hfluxt,wfluxt, tau,fluxt_zero)
    USE dimensions
    USE omp_para
    USE disvert_mod
    IMPLICIT NONE
    REAL(rstd), INTENT(IN)    :: hflux(3*iim*jjm,llm), wflux(iim*jjm,llm+1)
    REAL(rstd), INTENT(INOUT) :: hfluxt(3*iim*jjm,llm), wfluxt(iim*jjm,llm+1)
    REAL(rstd), INTENT(IN) :: tau
    LOGICAL, INTENT(INOUT) :: fluxt_zero
    INTEGER :: l,i,j,ij

    IF(fluxt_zero) THEN

       fluxt_zero=.FALSE.

       DO l=ll_begin,ll_end
          !$SIMD
          DO ij=ij_begin_ext,ij_end_ext
             hfluxt(ij+u_right,l) = tau*hflux(ij+u_right,l)
             hfluxt(ij+u_lup,l) = tau*hflux(ij+u_lup,l)
             hfluxt(ij+u_ldown,l) = tau*hflux(ij+u_ldown,l)
          ENDDO
       ENDDO

       IF(caldyn_eta==eta_mass) THEN ! no need for vertical fluxes if eta_lag
          DO l=ll_begin,ll_endp1
             !$SIMD
             DO ij=ij_begin,ij_end
                wfluxt(ij,l) = tau*wflux(ij,l)
             ENDDO
          ENDDO
       END IF

    ELSE

       DO l=ll_begin,ll_end
          !$SIMD
          DO ij=ij_begin_ext,ij_end_ext
             hfluxt(ij+u_right,l) = hfluxt(ij+u_right,l)+tau*hflux(ij+u_right,l)
             hfluxt(ij+u_lup,l) = hfluxt(ij+u_lup,l)+tau*hflux(ij+u_lup,l)
             hfluxt(ij+u_ldown,l) = hfluxt(ij+u_ldown,l)+tau*hflux(ij+u_ldown,l)
          ENDDO
       ENDDO

       IF(caldyn_eta==eta_mass) THEN ! no need for vertical fluxes if eta_lag
          DO l=ll_begin,ll_endp1
             !$SIMD
             DO ij=ij_begin,ij_end
                wfluxt(ij,l) = wfluxt(ij,l)+tau*wflux(ij,l)
             ENDDO
          ENDDO
       END IF

    END IF

  END SUBROUTINE accumulate_fluxes

  SUBROUTINE legacy_to_DEC(f_ps, f_u)
    USE icosa
    USE disvert_mod
    USE omp_para
    USE trace
    TYPE(t_field),POINTER :: f_ps(:), f_u(:)
    REAL(rstd), POINTER :: ps(:), u(:,:)
    INTEGER :: ind,ij,l
    CALL trace_start("legacy_to_DEC")

    DO ind=1,ndomain
       IF (.NOT. assigned_domain(ind)) CYCLE
       CALL swap_dimensions(ind)
       CALL swap_geometry(ind)

       IF(caldyn_eta==eta_mass .AND. is_omp_first_level) THEN ! update ps
          ps=f_ps(ind)
          !$SIMD
          DO ij=ij_begin,ij_end
             ps(ij)=(ps(ij)-ptop)/g ! convert ps to column-integrated mass
          ENDDO
       END IF

       u=f_u(ind)
       DO l=ll_begin,ll_end
          !$SIMD
          DO ij=ij_begin,ij_end
             u(ij+u_right,l)=u(ij+u_right,l)*de(ij+u_right)
             u(ij+u_lup,l)=u(ij+u_lup,l)*de(ij+u_lup)
             u(ij+u_ldown,l)=u(ij+u_ldown,l)*de(ij+u_ldown)
          ENDDO
       ENDDO
    ENDDO

    CALL trace_end("legacy_to_DEC")  
  END SUBROUTINE Legacy_to_DEC

  SUBROUTINE DEC_to_legacy(f_ps, f_u)
    USE icosa
    USE disvert_mod
    USE omp_para
    USE trace
    TYPE(t_field),POINTER :: f_ps(:), f_u(:)
    REAL(rstd), POINTER :: ps(:), u(:,:)
    INTEGER :: ind,ij,l
    CALL trace_start("legacy_to_DEC")

    DO ind=1,ndomain
       IF (.NOT. assigned_domain(ind)) CYCLE
       CALL swap_dimensions(ind)
       CALL swap_geometry(ind)

       IF(caldyn_eta==eta_mass .AND. is_omp_first_level) THEN
          ps=f_ps(ind)
          !$SIMD
          DO ij=ij_begin,ij_end
             ps(ij)=ptop+ps(ij)*g ! convert column-integrated mass to ps
          ENDDO
       ENDIF
       
       u=f_u(ind)
       DO l=ll_begin,ll_end
          !$SIMD
          DO ij=ij_begin,ij_end
             u(ij+u_right,l)=u(ij+u_right,l)/de(ij+u_right)
             u(ij+u_lup,l)=u(ij+u_lup,l)/de(ij+u_lup)
             u(ij+u_ldown,l)=u(ij+u_ldown,l)/de(ij+u_ldown)
          ENDDO
       ENDDO
    ENDDO

    CALL trace_end("DEC_to_legacy")  
  END SUBROUTINE DEC_to_legacy

END MODULE euler_scheme_mod