MODULE compute_theta_mod
  USE grid_param
  USE disvert_mod, ONLY : mass_dak, mass_dbk, caldyn_eta, eta_mass, ptop
  IMPLICIT NONE
  PRIVATE

#include "../unstructured/unstructured.h90"

  PUBLIC :: compute_theta_unst, compute_theta_hex, compute_theta_manual

CONTAINS

  ! Python/Fortran differences to be resolved at some point :

  ! the Fortran driver prognoses ps/mass_col or rhodz depending on caldyn_eta
  ! the Python driver prognoses rhodz even if caldyn_eta==eta_mass
  ! so that the DOFs are the same whatever caldyn_eta

  ! in the Fortran driver dak, dbk are 1D and based on pressure
  ! => m = mass_dak(l)+(mass_col(ij)*g+ptop)*mass_dbk(l)
  !    rhodz(CELL) = m/g 
  ! in the Python driver dak, dbk are 2D and based on mass
  ! => rhodz(CELL) = MASS_DAK(CELL) + mass_col(HIDX(CELL))*MASS_DBK(CELL)

#ifdef BEGIN_DYSL

KERNEL(theta)
  IF(caldyn_eta==eta_mass) THEN ! Compute mass
    ! compute mass_col from rhodz
    SEQUENCE_C1
      PROLOGUE(0)
        mass_col(HIDX(CELL))=0.
      END_BLOCK
      BODY('1,llm')
          mass_col(HIDX(CELL)) = mass_col(HIDX(CELL)) + rhodz(CELL)
      END_BLOCK
    END_BLOCK
    FORALL_CELLS_EXT()
      ON_PRIMAL
        rhodz(CELL) = MASS_DAK(CELL) + mass_col(HIDX(CELL))*MASS_DBK(CELL)
      END_BLOCK
    END_BLOCK
  END IF
  DO iq=1,nqdyn
    FORALL_CELLS_EXT()
      ON_PRIMAL
        theta(CELL,iq) = theta_rhodz(CELL,iq)/rhodz(CELL)
      END_BLOCK
    END_BLOCK
  END DO
END_BLOCK

KERNEL(compute_theta)
  IF(caldyn_eta==eta_mass) THEN ! Compute mass
     FORALL_CELLS_EXT()
       ON_PRIMAL
         m = MASS_DAK(CELL)+(mass_col(HIDX(CELL))*g+ptop)*MASS_DBK(CELL)
         rhodz(CELL) = m/g
        END_BLOCK
     END_BLOCK
  END IF  
  DO iq=1,nqdyn
     FORALL_CELLS_EXT()
       ON_PRIMAL
         theta(CELL,iq) = theta_rhodz(CELL,iq)/rhodz(CELL)
       END_BLOCK
     END_BLOCK
  END DO
END_BLOCK

#endif END_DYSL

  SUBROUTINE compute_theta_unst(mass_col,theta_rhodz, rhodz,theta)
    USE ISO_C_BINDING, only : C_DOUBLE, C_FLOAT
    USE data_unstructured_mod, ONLY : enter_trace, exit_trace, id_theta
    FIELD_PS :: mass_col
    FIELD_MASS :: rhodz
    FIELD_THETA :: theta, theta_rhodz
    DECLARE_INDICES
    NUM :: m
    START_TRACE(id_theta, 3,0,0) ! primal, dual, edge  
#define MASS_DAK(l,ij) mass_dak(l)
#define MASS_DBK(l,ij) mass_dbk(l)
#include "../kernels_unst/theta.k90"
#undef MASS_DAK
#undef MASS_DBK
    STOP_TRACE
  END SUBROUTINE compute_theta_unst

  SUBROUTINE compute_theta_hex(mass_col,theta_rhodz, rhodz,theta)
    USE icosa
    USE trace, ONLY : trace_start, trace_end
    USE omp_para, ONLY : ll_begin, ll_end
    REAL(rstd),INTENT(IN)    :: mass_col(iim*jjm)
    REAL(rstd),INTENT(IN)    :: theta_rhodz(iim*jjm,llm,nqdyn)
    REAL(rstd),INTENT(INOUT) :: rhodz(iim*jjm,llm)
    REAL(rstd),INTENT(OUT)   :: theta(iim*jjm,llm,nqdyn)
    INTEGER :: ij,l,iq
    REAL(rstd) :: m
    CALL trace_start("compute_theta")
#define MASS_DAK(ij,l) mass_dak(l)
#define MASS_DBK(ij,l) mass_dbk(l)
#include "../kernels_hex/compute_theta.k90"
#undef MASS_DAK
#undef MASS_DBK
    CALL trace_end("compute_theta")
  END SUBROUTINE compute_theta_hex
  
  SUBROUTINE compute_theta_manual(ps,theta_rhodz, rhodz,theta)
    USE icosa
    USE trace, ONLY : trace_start, trace_end
    USE omp_para, ONLY : ll_begin, ll_end
    REAL(rstd),INTENT(IN)    :: ps(iim*jjm)
    REAL(rstd),INTENT(IN)    :: theta_rhodz(iim*jjm,llm,nqdyn)
    REAL(rstd),INTENT(INOUT) :: rhodz(iim*jjm,llm)
    REAL(rstd),INTENT(OUT)   :: theta(iim*jjm,llm,nqdyn)
    INTEGER :: ij,l,iq
    REAL(rstd) :: m
    CALL trace_start("compute_theta")

    IF(caldyn_eta==eta_mass) THEN ! Compute mass
       DO l = ll_begin,ll_end
          !DIR$ SIMD
          DO ij=ij_begin_ext,ij_end_ext
             m = mass_dak(l)+(ps(ij)*g+ptop)*mass_dbk(l) ! ps is actually Ms
             rhodz(ij,l) = m/g
          END DO
       END DO
    END IF

    DO l = ll_begin,ll_end
       DO iq=1,nqdyn
          !DIR$ SIMD
          DO ij=ij_begin_ext,ij_end_ext
             theta(ij,l,iq) = theta_rhodz(ij,l,iq)/rhodz(ij,l)
          END DO
       END DO
    END DO

    CALL trace_end("compute_theta")
  END SUBROUTINE compute_theta_manual
  
END MODULE compute_theta_mod