MODULE hevi_scheme_mod
  USE prec
  USE domain_mod
  USE field_mod
  USE euler_scheme_mod
  USE caldyn_kernels_base_mod, ONLY : DEC
  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 :: alpha=(3.+SQRT(8.))/6., delta=.5/SQRT(2.), gamma=1.-2.*delta
    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
    ! To check correctness of caldyn_hevi
    REAL(rstd) :: dt
    REAL(rstd), PARAMETER, DIMENSION(3,4) :: &
         ajl=(/ zero, (/.5,0.,0./), (/-1.,2.,0./), (/1./6.,2./3.,1./6./) /)
    CALL set_coefs_hevi(dt, ajl, ajl)
  END SUBROUTINE set_coefs_ark33
    
  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
    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(:,:)

    IF(DEC) 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_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) )
       ! 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(bjl(l,j), f_mass, f_dmass_slow(:,l))
          END IF
          CALL update(bjl(l,j), f_theta_rhodz, f_dtheta_rhodz_slow(:,l))
          CALL update(bjl(l,j), f_u, f_du_slow(:,l))
          CALL update(cjl(l,j), f_u, f_du_fast(:,l))
       END DO
    END DO
    IF(DEC) CALL DEC_to_legacy(f_ps, f_u)
  END SUBROUTINE HEVI_scheme
  
  SUBROUTINE update(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(w,y,dy)
       ENDDO
    END IF
  END SUBROUTINE update
    
  SUBROUTINE compute_update(w, y, dy)
    USE omp_para
    USE disvert_mod
    REAL(rstd) :: w
    REAL(rstd) :: y(:,:), dy(:,:)
    INTENT(INOUT) :: y
    INTENT(IN) :: dy
    INTEGER :: l
    DO l=ll_begin,ll_end
       y(:,l)=y(:,l)+w*dy(:,l)
    ENDDO
  END SUBROUTINE compute_update
  
  SUBROUTINE update_2D(w, f_y, f_dy)
    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)
       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
    y(:)=y(:)+w*dy(:)
  END SUBROUTINE compute_update_2D
  
END MODULE hevi_scheme_mod