MODULE exner_mod

CONTAINS

  SUBROUTINE exner(f_ps,f_p,f_pks,f_pk)
  USE icosa
  IMPLICIT NONE
    TYPE(t_field), POINTER :: f_ps(:)
    TYPE(t_field), POINTER :: f_p(:)
    TYPE(t_field), POINTER :: f_pks(:)
    TYPE(t_field), POINTER :: f_pk(:)
  
    REAL(rstd), POINTER :: ps(:)
    REAL(rstd), POINTER :: p(:,:)
    REAL(rstd), POINTER :: pks(:)
    REAL(rstd), POINTER :: pk(:,:)
    INTEGER :: ind

    DO ind=1,ndomain
      CALL swap_dimensions(ind)
      CALL swap_geometry(ind)
      ps=f_ps(ind)
      p=f_p(ind)
      pks=f_pks(ind)
      pk=f_pk(ind)
      CALL compute_exner(ps, p, pks, pk, 0)
    ENDDO
  
  END SUBROUTINE exner
  
  SUBROUTINE compute_exner(ps,p,pks,pk,offset)
  USE icosa
  USE disvert_mod
  USE pression_mod
  IMPLICIT NONE
    REAL(rstd),INTENT(IN) :: ps(iim*jjm)
    REAL(rstd),INTENT(IN) :: p(iim*jjm,llm+1)
    REAL(rstd),INTENT(OUT) :: pks(iim*jjm)
    REAL(rstd),INTENT(OUT) :: pk(iim*jjm,llm)
    INTEGER,INTENT(IN) :: offset
    INTEGER :: i,j,ij,l
    REAL(rstd) :: alpha(iim*jjm,llm),beta(iim*jjm,llm)
    REAL(rstd) :: delta 
    
    IF(.FALSE.) THEN          ! LMD-Z style calculation of Exner pressure
       !! Compute Alpha and Beta
       
       ! for llm layer
       !$OMP DO
       DO j=jj_begin-offset,jj_end+offset
          DO i=ii_begin-offset,ii_end+offset
             ij=(j-1)*iim+i
             alpha(ij,llm) = 0.
             beta (ij,llm) = 1./ (1+ 2*kappa)
          ENDDO
       ENDDO
       
       ! for other layer   
       DO l = llm-1 , 2 , -1
          !$OMP DO
          DO j=jj_begin-offset,jj_end+offset
             DO i=ii_begin-offset,ii_end+offset
                ij=(j-1)*iim+i
                delta = p(ij,l)* (1+2*kappa) + p(ij,l+1)* ( beta(ij,l+1)- (1+2*kappa) )
                alpha(ij,l)  = - p(ij,l+1) / delta * alpha(ij,l+1)
                beta (ij,l)  =   p(ij,l  ) / delta   
             ENDDO
          ENDDO
       ENDDO
       
       !! Compute pk
       
       ! for first layer
       !$OMP DO
       DO j=jj_begin-offset,jj_end+offset
          DO i=ii_begin-offset,ii_end+offset
             ij=(j-1)*iim+i
             pks(ij) = cpp * ( ps(ij)/preff ) ** kappa
             pk(ij,1) = ( p(ij,1)*pks(ij) - 0.5*alpha(ij,2)*p(ij,2) )     /    &
                  (  p(ij,1)* (1.+kappa) + 0.5*( beta(ij,2)-(1.+2*kappa) )* p(ij,2) )
          ENDDO
       ENDDO
       
       ! for other layers
       
       DO l = 2, llm
          !$OMP DO
          DO j=jj_begin-offset,jj_end+offset
             DO i=ii_begin-offset,ii_end+offset
                ij=(j-1)*iim+i
                pk(ij,l) = alpha(ij,l) + beta(ij,l) * pk(ij,l-1)
             ENDDO
          ENDDO
       ENDDO
       
    ELSE ! Simple calculation of Exner pressure based on centered average
       ! surface : pks
       !$OMP DO
       DO j=jj_begin-offset,jj_end+offset
          DO i=ii_begin-offset,ii_end+offset
             ij=(j-1)*iim+i
             pks(ij) = cpp * ( ps(ij)/preff ) ** kappa
          ENDDO
       ENDDO

       ! 3D : pk
       DO l = 1, llm
          !$OMP DO
          DO j=jj_begin-offset,jj_end+offset
             DO i=ii_begin-offset,ii_end+offset
                ij=(j-1)*iim+i
                pk(ij,l) = cpp * ((.5/preff)*(p(ij,l)+p(ij,l+1))) ** kappa
             ENDDO
          ENDDO
       ENDDO
    END IF
    
  END SUBROUTINE compute_exner
  
END MODULE exner_mod