MODULE compute_theta_mod
  USE grid_param, ONLY : llm
  IMPLICIT NONE
  PRIVATE

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

  PUBLIC :: compute_theta

CONTAINS

#ifdef BEGIN_DYSL

KERNEL(theta)
  IF(caldyn_eta==eta_mass) THEN ! Compute mass
    ! FIXME : here mass_col is computed from rhodz
    ! so that the DOFs are the same whatever caldyn_eta
    ! in DYNAMICO mass_col is prognosed rather than 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
        ! FIXME : formula below (used in DYNAMICO) is for dak, dbk based on
        ! pressure rather than mass
        !        m = mass_dak(CELL)+(mass_col(HIDX(CELL))*g+ptop)*mass_dbk(CELL)
        !        rhodz(CELL) = m/g 
        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

#endif END_DYSL

SUBROUTINE compute_theta_unst(mass_col,rhodz,theta_rhodz, theta)
  USE ISO_C_BINDING, only : C_DOUBLE, C_FLOAT
  USE data_unstructured_mod, ONLY : id_theta,primal_num,dual_num,edge_num, &
    enter_trace, exit_trace
  USE grid_param, ONLY : nqdyn
  USE disvert_mod, ONLY : mass_dak, mass_dbk, caldyn_eta, eta_mass, ptop
  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

  SUBROUTINE compute_theta(ps,theta_rhodz, rhodz,theta)
    USE icosa
    USE trace, ONLY : trace_start, trace_end
    USE omp_para, ONLY : ll_begin, ll_end
    USE disvert_mod, ONLY : mass_dak, mass_dbk, caldyn_eta, eta_mass, ptop
    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
  
END MODULE compute_theta_mod