MODULE hevi_scheme_mod
  USE prec
  USE domain_mod
  USE field_mod
  USE euler_scheme_mod
  USE abort_mod
  IMPLICIT NONE
  PRIVATE

  REAL(rstd), SAVE :: wj(3), bjl(3,3), cjl(3,3), taujj(3)
  REAL(rstd), PARAMETER, DIMENSION(3) :: zero = (/0.,0.,0./)

  PUBLIC :: set_coefs_ark23, set_coefs_ark33, hevi_scheme

CONTAINS
  SUBROUTINE set_coefs_ark23(dt)
    ! ARK2 scheme by Giraldo, Kelly, Constantinescu 2013
    ! See Weller et al., 2013 - ARK2 scheme Fig. 2
    REAL(rstd) :: dt
    REAL(rstd), PARAMETER :: delta=.5/SQRT(2.), gamma=1.-2.*delta
!    REAL(rstd), PARAMETER :: alpha=(3.+SQRT(8.))/6. ! original value in GKC2013
    REAL(rstd), PARAMETER :: alpha=0.7
    REAL(rstd), PARAMETER, DIMENSION(3) :: wj = (/delta,delta,gamma/)
    CALL set_coefs_hevi(dt, &
         (/ zero, (/2.*gamma,0.,0./), (/1-alpha,alpha,0./), wj /), &
         (/ zero, (/gamma,gamma,0./), wj, wj /) )
  END SUBROUTINE set_coefs_ark23

  SUBROUTINE set_coefs_ark33(dt)
    ! Fully-explicit RK3 scheme disguised as ARK
    REAL(rstd) :: dt
    CALL set_coefs_rk(dt, (/ zero, (/.5,0.,0./), (/-1.,2.,0./), (/1./6.,2./3.,1./6./) /) )
  END SUBROUTINE set_coefs_ark33
    
  SUBROUTINE set_coefs_rk(dt, ajl)
    REAL(rstd) :: dt, ajl(3,4)
    CALL set_coefs_hevi(dt,ajl,ajl)
  END SUBROUTINE set_coefs_rk

  SUBROUTINE set_coefs_hevi(dt, ajl_slow, ajl_fast)
    REAL(rstd) :: dt, ajl_slow(3,4), ajl_fast(3,4) ! fast/slow Butcher tableaus
    INTEGER :: j
    DO j=1,3
       bjl(:,j) = dt*(ajl_slow(:,j+1)-ajl_slow(:,j))
       cjl(:,j) = dt*(ajl_fast(:,j+1)-ajl_fast(:,j))
       taujj(j) = dt*ajl_fast(j,j)
    END DO
    wj=dt*ajl_slow(:,4)
  END SUBROUTINE set_coefs_hevi

  SUBROUTINE HEVI_scheme(it, fluxt_zero)
    USE time_mod
    USE disvert_mod
    USE caldyn_hevi_mod
    USE omp_para
    USE checksum_mod
    LOGICAL :: fluxt_zero(ndomain) ! set to .TRUE. to start accumulating fluxes in time
    INTEGER :: it,j,l, ind
    REAL(rstd),POINTER :: hflux(:,:),wflux(:,:),hfluxt(:,:),wfluxt(:,:)

    CALL legacy_to_DEC(f_ps, f_u)

    DO j=1,nb_stage
       CALL caldyn_hevi((j==1) .AND. (MOD(it,itau_out)==0), taujj(j), &
            f_phis, f_ps,f_mass,f_theta_rhodz,f_u,f_q, &
            f_W, f_geopot, f_hflux, f_wflux, &
            f_dps_slow(:,j), f_dmass_slow(:,j), f_dtheta_rhodz_slow(:,j), &
            f_du_slow(:,j), f_du_fast(:,j), &
            f_dPhi_slow(:,j), f_dPhi_fast(:,j), &
            f_dW_slow(:,j), f_dW_fast(:,j) )
       ! accumulate mass fluxes for transport scheme
       
       DO ind=1,ndomain
          IF (.NOT. assigned_domain(ind)) CYCLE
          CALL swap_dimensions(ind)
          hflux=f_hflux(ind);     hfluxt=f_hfluxt(ind)
          wflux=f_wflux(ind);     wfluxt=f_wfluxt(ind)
          CALL accumulate_fluxes(hflux,wflux, hfluxt,wfluxt, wj(j), fluxt_zero(ind))

       END DO
       ! update model state
       DO l=1,j
          IF(caldyn_eta==eta_mass) THEN
             CALL update_2D(bjl(l,j), f_ps, f_dps_slow(:,l))
          ELSE
             CALL update_3D(bjl(l,j), f_mass, f_dmass_slow(:,l))
          END IF
          CALL update_4D(bjl(l,j), f_theta_rhodz, f_dtheta_rhodz_slow(:,l))
          CALL update_3D(bjl(l,j), f_u, f_du_slow(:,l))
          CALL update_3D(cjl(l,j), f_u, f_du_fast(:,l))
          IF(.NOT. hydrostatic) THEN
             CALL abort_acc("HEVI_scheme/!hydrostatic")
             CALL update_3D(bjl(l,j), f_W, f_dW_slow(:,l))
             CALL update_3D(cjl(l,j), f_W, f_dW_fast(:,l))
             CALL update_3D(bjl(l,j), f_geopot, f_dPhi_slow(:,l))
             CALL update_3D(cjl(l,j), f_geopot, f_dPhi_fast(:,l))
          END IF

!$OMP BARRIER       
       END DO
    END DO
    CALL DEC_to_legacy(f_ps, f_u)
  END SUBROUTINE HEVI_scheme
  
  SUBROUTINE update_4D(w, f_y, f_dy)
    USE dimensions
    USE grid_param, ONLY : nqdyn
    REAL(rstd) :: w
    TYPE(t_field) :: f_y(:), f_dy(:)
    REAL(rstd), POINTER :: y(:,:,:), dy(:,:,:)
    INTEGER :: ind, iq
    IF(w /= 0.) THEN
       DO ind=1,ndomain
          IF (.NOT. assigned_domain(ind)) CYCLE
          CALL swap_dimensions(ind)
          dy=f_dy(ind); y=f_y(ind)
          DO iq=1,nqdyn
             CALL compute_update_3D(w,y(:,:,iq),dy(:,:,iq))
          END DO
       ENDDO
    END IF
  END SUBROUTINE update_4D
    
  SUBROUTINE update_3D(w, f_y, f_dy)
    USE dimensions
    REAL(rstd) :: w
    TYPE(t_field) :: f_y(:), f_dy(:)
    REAL(rstd), POINTER :: y(:,:), dy(:,:)
    INTEGER :: ind
    IF(w /= 0.) THEN
       DO ind=1,ndomain
          IF (.NOT. assigned_domain(ind)) CYCLE
          CALL swap_dimensions(ind)
          dy=f_dy(ind); y=f_y(ind)
          CALL compute_update_3D(w,y,dy)
       ENDDO
    END IF
  END SUBROUTINE update_3D
    
  SUBROUTINE compute_update_3D(w, y, dy)
    USE omp_para
    USE disvert_mod
    REAL(rstd) :: w
    REAL(rstd) :: y(:,:), dy(:,:)
    INTENT(INOUT) :: y
    INTENT(IN) :: dy
    INTEGER :: l
    !$acc kernels async default(present)
    DO l=ll_begin,ll_end
       y(:,l)=y(:,l)+w*dy(:,l)
    ENDDO
    !$acc end kernels
  END SUBROUTINE compute_update_3D
  
  SUBROUTINE update_2D(w, f_y, f_dy)
  USE omp_para
    REAL(rstd) :: w
    TYPE(t_field) :: f_y(:), f_dy(:)
    REAL(rstd), POINTER :: y(:), dy(:)
    INTEGER :: ind
    DO ind=1,ndomain
       IF (.NOT. assigned_domain(ind)) CYCLE
       dy=f_dy(ind); y=f_y(ind)
       IF (is_omp_level_master) CALL compute_update_2D(w,y,dy)
    ENDDO
  END SUBROUTINE update_2D
    
  SUBROUTINE compute_update_2D(w, y, dy)
    REAL(rstd) :: w, y(:), dy(:)
    INTENT(INOUT) :: y
    INTENT(IN) :: dy
    !$acc kernels async default(present)
    y(:)=y(:)+w*dy(:)
    !$acc end kernels
  END SUBROUTINE compute_update_2D
  
END MODULE hevi_scheme_mod