MODULE caldyn_gcm_mod
  USE icosa
  USE transfert_mod
  PRIVATE

  INTEGER, PARAMETER :: energy=1, enstrophy=2
  TYPE(t_field),POINTER :: f_out_u(:), f_qu(:), f_qv(:)
  REAL(rstd),SAVE,POINTER :: out_u(:,:), p(:,:), qu(:,:)
  !$OMP THREADPRIVATE(out_u, p, qu)

  TYPE(t_field),POINTER :: f_buf_i(:), f_buf_ulon(:), f_buf_ulat(:), f_buf_u3d(:)
  TYPE(t_field),POINTER :: f_buf_v(:), f_buf_s(:), f_buf_p(:)

! temporary shared variable for caldyn
  TYPE(t_field),POINTER :: f_theta(:)
  TYPE(t_field),POINTER :: f_pk(:)
  TYPE(t_field),POINTER :: f_geopot(:)
  TYPE(t_field),POINTER :: f_wwuu(:)
  TYPE(t_field),POINTER :: f_planetvel(:)
   
  INTEGER :: caldyn_conserv
 !$OMP THREADPRIVATE(caldyn_conserv) 
 
  TYPE(t_message) :: req_ps, req_mass, req_theta_rhodz, req_u, req_qu

  PUBLIC init_caldyn, caldyn_BC, caldyn, write_output_fields, &
       req_ps, req_mass

CONTAINS
  
  SUBROUTINE init_caldyn
    USE icosa
    USE exner_mod
    USE mpipara
    IMPLICIT NONE
    CHARACTER(len=255) :: def
    INTEGER            :: ind
    REAL(rstd),POINTER :: planetvel(:)
  
    def='energy'
    CALL getin('caldyn_conserv',def)
    SELECT CASE(TRIM(def))
    CASE('energy')
       caldyn_conserv=energy
    CASE('enstrophy')
       caldyn_conserv=enstrophy
    CASE DEFAULT
       IF (is_mpi_root) PRINT *,'Bad selector for variable caldyn_conserv : <', &
            TRIM(def),'> options are <energy>, <enstrophy>'
       STOP
    END SELECT
    IF (is_mpi_root) PRINT *, 'caldyn_conserv=',def

    CALL allocate_caldyn

    DO ind=1,ndomain
       IF (.NOT. assigned_domain(ind)) CYCLE
       CALL swap_dimensions(ind)
       CALL swap_geometry(ind)
       planetvel=f_planetvel(ind)
       CALL compute_planetvel(planetvel)
    END DO

  END SUBROUTINE init_caldyn

  SUBROUTINE allocate_caldyn
  USE icosa
  IMPLICIT NONE

    CALL allocate_field(f_out_u,field_u,type_real,llm) 
    CALL allocate_field(f_qu,field_u,type_real,llm) 
    CALL allocate_field(f_qv,field_z,type_real,llm) 
  
    CALL allocate_field(f_buf_i,   field_t,type_real,llm)
    CALL allocate_field(f_buf_p,   field_t,type_real,llm+1) 
    CALL allocate_field(f_buf_u3d, field_t,type_real,3,llm)  ! 3D vel at cell centers
    CALL allocate_field(f_buf_ulon,field_t,type_real,llm)
    CALL allocate_field(f_buf_ulat,field_t,type_real,llm)
    CALL allocate_field(f_buf_v,   field_z,type_real,llm)
    CALL allocate_field(f_buf_s,   field_t,type_real)

    CALL allocate_field(f_theta, field_t,type_real,llm,  name='theta')   ! potential temperature
    CALL allocate_field(f_pk,    field_t,type_real,llm,  name='pk')
    CALL allocate_field(f_geopot,field_t,type_real,llm+1,name='geopot')  ! geopotential
    CALL allocate_field(f_wwuu,  field_u,type_real,llm+1,name='wwuu')
    CALL allocate_field(f_planetvel,  field_u,type_real, name='planetvel') ! planetary velocity at r=a

  END SUBROUTINE allocate_caldyn

  SUBROUTINE caldyn_BC(f_phis, f_wflux)
    USE icosa
    USE mpipara
    USE omp_para
    IMPLICIT NONE
    TYPE(t_field),POINTER :: f_phis(:)
    TYPE(t_field),POINTER :: f_wflux(:)
    REAL(rstd),POINTER  :: phis(:)
    REAL(rstd),POINTER  :: wflux(:,:)
    REAL(rstd),POINTER  :: geopot(:,:)
    REAL(rstd),POINTER  :: wwuu(:,:)

    INTEGER :: ind,i,j,ij,l

    IF (omp_first) THEN
       DO ind=1,ndomain
          IF (.NOT. assigned_domain(ind)) CYCLE
          CALL swap_dimensions(ind)
          CALL swap_geometry(ind)
          geopot=f_geopot(ind)
          phis=f_phis(ind)
          wflux=f_wflux(ind)
          wwuu=f_wwuu(ind)
          
          DO ij=ij_begin_ext,ij_end_ext
              ! lower BCs : geopot=phis, wflux=0, wwuu=0
              geopot(ij,1) = phis(ij)
              wflux(ij,1) = 0.
              wwuu(ij+u_right,1)=0   
              wwuu(ij+u_lup,1)=0   
              wwuu(ij+u_ldown,1)=0
              ! top BCs : wflux=0, wwuu=0
              wflux(ij,llm+1)  = 0.
              wwuu(ij+u_right,llm+1)=0   
              wwuu(ij+u_lup,llm+1)=0   
              wwuu(ij+u_ldown,llm+1)=0
          ENDDO
       END DO
    ENDIF

!    !$OMP BARRIER
  END SUBROUTINE caldyn_BC
   
  SUBROUTINE caldyn(write_out,f_phis, f_ps, f_mass, f_theta_rhodz, f_u, f_q, &
       f_hflux, f_wflux, f_dps, f_dmass, f_dtheta_rhodz, f_du)
    USE icosa
    USE disvert_mod, ONLY : caldyn_eta, eta_mass
    USE vorticity_mod
    USE kinetic_mod
    USE theta2theta_rhodz_mod
    USE wind_mod
    USE mpipara
    USE trace
    USE omp_para
    USE output_field_mod
    IMPLICIT NONE
    LOGICAL,INTENT(IN)    :: write_out
    TYPE(t_field),POINTER :: f_phis(:)
    TYPE(t_field),POINTER :: f_ps(:)
    TYPE(t_field),POINTER :: f_mass(:)
    TYPE(t_field),POINTER :: f_theta_rhodz(:)
    TYPE(t_field),POINTER :: f_u(:)
    TYPE(t_field),POINTER :: f_q(:)
    TYPE(t_field),POINTER :: f_hflux(:), f_wflux(:)
    TYPE(t_field),POINTER :: f_dps(:)
    TYPE(t_field),POINTER :: f_dmass(:)
    TYPE(t_field),POINTER :: f_dtheta_rhodz(:)
    TYPE(t_field),POINTER :: f_du(:)
    
    REAL(rstd),POINTER :: ps(:), dps(:)
    REAL(rstd),POINTER :: mass(:,:), theta_rhodz(:,:), dtheta_rhodz(:,:)
    REAL(rstd),POINTER :: u(:,:), du(:,:), hflux(:,:), wflux(:,:)
    REAL(rstd),POINTER :: qu(:,:)
    REAL(rstd),POINTER :: qv(:,:)

! temporary shared variable
    REAL(rstd),POINTER  :: theta(:,:)  
    REAL(rstd),POINTER  :: pk(:,:)
    REAL(rstd),POINTER  :: geopot(:,:)
    REAL(rstd),POINTER  :: convm(:,:) 
    REAL(rstd),POINTER  :: wwuu(:,:)
        
    INTEGER :: ind
    LOGICAL,SAVE :: first=.TRUE.
!$OMP THREADPRIVATE(first)
    
    ! MPI messages need to be sent at first call to caldyn
    ! This is needed only once : the next ones will be sent by timeloop
    IF (first) THEN 
      first=.FALSE.
      IF(caldyn_eta==eta_mass) THEN
         CALL init_message(f_ps,req_i1,req_ps)
      ELSE
         CALL init_message(f_mass,req_i1,req_mass)
      END IF
      CALL init_message(f_theta_rhodz,req_i1,req_theta_rhodz)
      CALL init_message(f_u,req_e1_vect,req_u)
      CALL init_message(f_qu,req_e1_scal,req_qu)
!      IF(caldyn_eta==eta_mass) THEN
!         CALL send_message(f_ps,req_ps) 
!         CALL wait_message(req_ps)  
!      ELSE
!         CALL send_message(f_mass,req_mass) 
!         CALL wait_message(req_mass)  
!      END IF
    ENDIF
    
    CALL trace_start("caldyn")

      IF(caldyn_eta==eta_mass) THEN
         CALL send_message(f_ps,req_ps) 
         CALL wait_message(req_ps)  
      ELSE
         CALL send_message(f_mass,req_mass) 
         CALL wait_message(req_mass)  
      END IF

    CALL send_message(f_u,req_u)
    CALL wait_message(req_u)
    CALL send_message(f_theta_rhodz,req_theta_rhodz) 
    CALL wait_message(req_theta_rhodz) 
    
!    CALL wait_message(req_u)
!    CALL wait_message(req_theta_rhodz) 

    SELECT CASE(caldyn_conserv)
    CASE(energy) ! energy-conserving
       DO ind=1,ndomain
          IF (.NOT. assigned_domain(ind)) CYCLE
          CALL swap_dimensions(ind)
          CALL swap_geometry(ind)
          ps=f_ps(ind)
          u=f_u(ind)
          theta_rhodz = f_theta_rhodz(ind)
          mass=f_mass(ind)
          theta = f_theta(ind)
          qu=f_qu(ind)
          qv=f_qv(ind)
          CALL compute_pvort(ps,u,theta_rhodz, mass,theta,qu,qv)
       ENDDO

       CALL send_message(f_qu,req_qu)
       CALL wait_message(req_qu)

       DO ind=1,ndomain
          IF (.NOT. assigned_domain(ind)) CYCLE
          CALL swap_dimensions(ind)
          CALL swap_geometry(ind)
          ps=f_ps(ind)
          u=f_u(ind)
          theta_rhodz=f_theta_rhodz(ind)
          mass=f_mass(ind)
          theta = f_theta(ind)
          qu=f_qu(ind)
          pk = f_pk(ind)
          geopot = f_geopot(ind)  
          CALL compute_geopot(ps,mass,theta, pk,geopot)
          hflux=f_hflux(ind)
          convm = f_dmass(ind)
          dtheta_rhodz=f_dtheta_rhodz(ind)
          du=f_du(ind)
          CALL compute_caldyn_horiz(u,mass,qu,theta,pk,geopot, hflux,convm,dtheta_rhodz,du)
          IF(caldyn_eta==eta_mass) THEN
             wflux=f_wflux(ind)
             wwuu=f_wwuu(ind)
             dps=f_dps(ind)
             CALL compute_caldyn_vert(u,theta,mass,convm, wflux,wwuu, dps, dtheta_rhodz, du)
          END IF
       ENDDO       
       
    CASE(enstrophy) ! enstrophy-conserving
       DO ind=1,ndomain
          IF (.NOT. assigned_domain(ind)) CYCLE
          CALL swap_dimensions(ind)
          CALL swap_geometry(ind)
          ps=f_ps(ind)
          u=f_u(ind)
          theta_rhodz=f_theta_rhodz(ind)
          mass=f_mass(ind)
          theta = f_theta(ind)
          qu=f_qu(ind)
          qv=f_qv(ind)
          CALL compute_pvort(ps,u,theta_rhodz, mass,theta,qu,qv)
          pk = f_pk(ind)
          geopot = f_geopot(ind)  
          CALL compute_geopot(ps,mass,theta, pk,geopot)
          hflux=f_hflux(ind)
          convm = f_dmass(ind)
          dtheta_rhodz=f_dtheta_rhodz(ind)
          du=f_du(ind)
          CALL compute_caldyn_horiz(u,mass,qu,theta,pk,geopot, hflux,convm,dtheta_rhodz,du)
          IF(caldyn_eta==eta_mass) THEN
             wflux=f_wflux(ind)
             wwuu=f_wwuu(ind)
             dps=f_dps(ind)
             CALL compute_caldyn_vert(u,theta,mass,convm, wflux,wwuu, dps, dtheta_rhodz, du)
          END IF
       ENDDO
       
    CASE DEFAULT
       STOP
    END SELECT

!!$OMP BARRIER
    IF (write_out) THEN

       IF (is_mpi_root) PRINT *,'CALL write_output_fields'

! ---> for openMP test to fix later
!       CALL write_output_fields(f_ps, f_phis, f_dps, f_u, f_theta_rhodz, f_q, &
!            f_buf_i, f_buf_v, f_buf_u3d, f_buf_ulon, f_buf_ulat, f_buf_s, f_buf_p)
       CALL un2ulonlat(f_u, f_buf_ulon, f_buf_ulat)
       CALL output_field("ulon",f_buf_ulon)
       CALL output_field("ulat",f_buf_ulat)

       CALL output_field("ps",f_ps)
       CALL output_field("dps",f_dps)
       CALL output_field("mass",f_mass)
       CALL output_field("dmass",f_dmass)
       CALL output_field("vort",f_qv)
       CALL output_field("theta",f_theta)
       CALL output_field("exner",f_pk)
       CALL output_field("pv",f_qv)
 
    END IF
    
    !    CALL check_mass_conservation(f_ps,f_dps)
    CALL trace_end("caldyn")
!!$OMP BARRIER
    
END SUBROUTINE caldyn

SUBROUTINE compute_planetvel(planetvel)
  USE wind_mod
  REAL(rstd),INTENT(OUT)  :: planetvel(iim*3*jjm)
  REAL(rstd) :: ulon(iim*3*jjm)
  REAL(rstd) :: ulat(iim*3*jjm)
  REAL(rstd) :: lon,lat
  INTEGER :: ij
  DO ij=ij_begin_ext,ij_end_ext
     CALL xyz2lonlat(xyz_e(ij+u_right,:),lon,lat)
     ulon(ij+u_right)=a*omega*cos(lat)
     ulat(ij+u_right)=0

     CALL xyz2lonlat(xyz_e(ij+u_lup,:),lon,lat)
     ulon(ij+u_lup)=a*omega*cos(lat)
     ulat(ij+u_lup)=0

     CALL xyz2lonlat(xyz_e(ij+u_ldown,:),lon,lat)
     ulon(ij+u_ldown)=a*omega*cos(lat)
     ulat(ij+u_ldown)=0
  END DO
  CALL compute_wind2D_perp_from_lonlat_compound(ulon, ulat, planetvel)
END SUBROUTINE compute_planetvel
    
SUBROUTINE compute_pvort(ps,u,theta_rhodz, rhodz,theta,qu,qv)
  USE icosa
  USE disvert_mod, ONLY : mass_dak, mass_dbk, caldyn_eta, eta_mass
  USE exner_mod
  USE trace
  USE omp_para
  IMPLICIT NONE
  REAL(rstd),INTENT(IN)  :: u(iim*3*jjm,llm)
  REAL(rstd),INTENT(IN)  :: ps(iim*jjm)
  REAL(rstd),INTENT(IN)  :: theta_rhodz(iim*jjm,llm)
  REAL(rstd),INTENT(INOUT) :: rhodz(iim*jjm,llm)
  REAL(rstd),INTENT(OUT) :: theta(iim*jjm,llm)
  REAL(rstd),INTENT(OUT) :: qu(iim*3*jjm,llm)
  REAL(rstd),INTENT(OUT) :: qv(iim*2*jjm,llm)
  
  INTEGER :: i,j,ij,l
  REAL(rstd) :: etav,hv, m
!  REAL(rstd) :: qv(2*iim*jjm,llm)     ! potential velocity  
  
  CALL trace_start("compute_pvort")  

  IF(caldyn_eta==eta_mass) THEN
!     CALL wait_message(req_ps)  
  ELSE
!     CALL wait_message(req_mass)
  END IF
!  CALL wait_message(req_theta_rhodz) 

  IF(caldyn_eta==eta_mass) THEN ! Compute mass & theta
     DO l = ll_begin,ll_end
!        CALL test_message(req_u) 
!DIR$ SIMD
        DO ij=ij_begin_ext,ij_end_ext
           m = ( mass_dak(l)+ps(ij)*mass_dbk(l) )/g
           rhodz(ij,l) = m
           theta(ij,l) = theta_rhodz(ij,l)/rhodz(ij,l)
        ENDDO
     ENDDO
  ELSE ! Compute only theta
     DO l = ll_begin,ll_end
!        CALL test_message(req_u) 
!DIR$ SIMD
       DO ij=ij_begin_ext,ij_end_ext
         theta(ij,l) = theta_rhodz(ij,l)/rhodz(ij,l)
       ENDDO
     ENDDO
  END IF

!  CALL wait_message(req_u)   
  
!!! Compute shallow-water potential vorticity
  DO l = ll_begin,ll_end
!DIR$ SIMD
     DO ij=ij_begin_ext,ij_end_ext
          etav= 1./Av(ij+z_up)*(  ne_rup        * u(ij+u_rup,l)        * de(ij+u_rup)         &
                                + ne_left * u(ij+t_rup+u_left,l) * de(ij+t_rup+u_left)  &
                                - ne_lup        * u(ij+u_lup,l)        * de(ij+u_lup) )                               

          hv =  Riv2(ij,vup)          * rhodz(ij,l)            &
              + Riv2(ij+t_rup,vldown) * rhodz(ij+t_rup,l)     &
              + Riv2(ij+t_lup,vrdown) * rhodz(ij+t_lup,l)
     
          qv(ij+z_up,l) = ( etav+fv(ij+z_up) )/hv
          
          etav = 1./Av(ij+z_down)*(  ne_ldown         * u(ij+u_ldown,l)          * de(ij+u_ldown)          &
                                   + ne_right * u(ij+t_ldown+u_right,l)  * de(ij+t_ldown+u_right)  &
                                   - ne_rdown         * u(ij+u_rdown,l)          * de(ij+u_rdown) )
       
          hv =  Riv2(ij,vdown)        * rhodz(ij,l)          &
              + Riv2(ij+t_ldown,vrup) * rhodz(ij+t_ldown,l)  &
              + Riv2(ij+t_rdown,vlup) * rhodz(ij+t_rdown,l)
                       
          qv(ij+z_down,l) =( etav+fv(ij+z_down) )/hv

      ENDDO

!DIR$ SIMD
      DO ij=ij_begin,ij_end
         qu(ij+u_right,l) = 0.5*(qv(ij+z_rdown,l)+qv(ij+z_rup,l))  
         qu(ij+u_lup,l) = 0.5*(qv(ij+z_up,l)+qv(ij+z_lup,l))  
         qu(ij+u_ldown,l) = 0.5*(qv(ij+z_ldown,l)+qv(ij+z_down,l))  
      END DO
      
   ENDDO

   CALL trace_end("compute_pvort")
                                                       
  END SUBROUTINE compute_pvort
  
  SUBROUTINE compute_geopot(ps,rhodz,theta, pk,geopot) 
  USE icosa
  USE disvert_mod
  USE exner_mod
  USE trace
  USE omp_para
  IMPLICIT NONE
    REAL(rstd),INTENT(INOUT) :: ps(iim*jjm)
    REAL(rstd),INTENT(IN)    :: rhodz(iim*jjm,llm)
    REAL(rstd),INTENT(IN)    :: theta(iim*jjm,llm)    ! potential temperature
    REAL(rstd),INTENT(INOUT) :: pk(iim*jjm,llm)       ! Exner function
    REAL(rstd),INTENT(INOUT) :: geopot(iim*jjm,llm+1) ! geopotential
    
    INTEGER :: i,j,ij,l
    REAL(rstd) :: p_ik, exner_ik

    CALL trace_start("compute_geopot")

    IF(caldyn_eta==eta_mass) THEN

!!! Compute exner function and geopotential
       DO l = 1,llm
!         !$OMP DO SCHEDULE(STATIC) 
          !DIR$ SIMD
          DO ij=ij_begin_ext,ij_end_ext         
             p_ik = ptop + mass_ak(l) + mass_bk(l)*ps(ij) ! FIXME : leave ps for the moment ; change ps to Ms later
             !         p_ik = ptop + g*(mass_ak(l)+ mass_bk(l)*ps(i,j))
             exner_ik = cpp * (p_ik/preff) ** kappa
             pk(ij,l) = exner_ik
             ! specific volume v = kappa*theta*pi/p = dphi/g/rhodz
             geopot(ij,l+1) = geopot(ij,l) + (g*kappa)*rhodz(ij,l)*theta(ij,l)*exner_ik/p_ik
          ENDDO
       ENDDO

    ELSE 
       ! We are using a Lagrangian vertical coordinate
       ! Pressure must be computed first top-down (temporarily stored in pk)
       ! Then Exner pressure and geopotential are computed bottom-up
       ! Notice that the computation below should work also when caldyn_eta=eta_mass          

       IF(boussinesq) THEN ! compute only geopotential : pressure pk will be computed in compute_caldyn_horiz
          ! specific volume 1 = dphi/g/rhodz
          DO l = 1,llm
!             !$OMP DO SCHEDULE(STATIC) 
             !DIR$ SIMD
             DO ij=ij_begin_ext,ij_end_ext         
                geopot(ij,l+1) = geopot(ij,l) + g*rhodz(ij,l)
             ENDDO
          ENDDO
       ELSE ! non-Boussinesq, compute geopotential and Exner pressure
          ! uppermost layer
          !DIR$ SIMD
          DO ij=ij_begin_ext,ij_end_ext         
             pk(ij,llm) = ptop + (.5*g)*rhodz(ij,llm)
          END DO
          ! other layers
          DO l = llm-1, 1, -1

!          !$OMP DO SCHEDULE(STATIC) 
             !DIR$ SIMD
             DO ij=ij_begin_ext,ij_end_ext         
                pk(ij,l) = pk(ij,l+1) + (.5*g)*(rhodz(ij,l)+rhodz(ij,l+1))
             END DO
          END DO
          ! surface pressure (for diagnostics)
          DO ij=ij_begin_ext,ij_end_ext         
             ps(ij) = pk(ij,1) + (.5*g)*rhodz(ij,1)
          END DO

          ! specific volume v = kappa*theta*pi/p = dphi/g/rhodz
          DO l = 1,llm

!             !$OMP DO SCHEDULE(STATIC) 
             !DIR$ SIMD
             DO ij=ij_begin_ext,ij_end_ext         
                p_ik = pk(ij,l)
                exner_ik = cpp * (p_ik/preff) ** kappa
                geopot(ij,l+1) = geopot(ij,l) + (g*kappa)*rhodz(ij,l)*theta(ij,l)*exner_ik/p_ik 
                pk(ij,l) = exner_ik
             ENDDO
          ENDDO
       END IF

    END IF

  CALL trace_end("compute_geopot")

  END SUBROUTINE compute_geopot

  SUBROUTINE compute_caldyn_horiz(u,rhodz,qu,theta,pk,geopot, hflux,convm, dtheta_rhodz, du)
  USE icosa
  USE disvert_mod
  USE exner_mod
  USE trace
  USE omp_para
  IMPLICIT NONE
    REAL(rstd),INTENT(IN)  :: u(iim*3*jjm,llm)    ! prognostic "velocity"
    REAL(rstd),INTENT(IN)  :: rhodz(iim*jjm,llm)
    REAL(rstd),INTENT(IN)  :: qu(iim*3*jjm,llm)
    REAL(rstd),INTENT(IN)  :: theta(iim*jjm,llm)  ! potential temperature
    REAL(rstd),INTENT(INOUT) :: pk(iim*jjm,llm) ! Exner function
    REAL(rstd),INTENT(IN)  :: geopot(iim*jjm,llm+1)    ! geopotential

    REAL(rstd),INTENT(OUT) :: hflux(iim*3*jjm,llm) ! hflux in kg/s
    REAL(rstd),INTENT(OUT) :: convm(iim*jjm,llm)  ! mass flux convergence
    REAL(rstd),INTENT(OUT) :: dtheta_rhodz(iim*jjm,llm)
    REAL(rstd),INTENT(OUT) :: du(iim*3*jjm,llm)

    REAL(rstd) :: cor_NT(iim*jjm,llm)  ! NT coriolis force u.(du/dPhi)
    REAL(rstd) :: urel(3*iim*jjm,llm)  ! relative velocity
    REAL(rstd) :: Ftheta(3*iim*jjm,llm) ! theta flux
    REAL(rstd) :: berni(iim*jjm,llm)  ! Bernoulli function
    
    INTEGER :: i,j,ij,l
    REAL(rstd) :: ww,uu

    CALL trace_start("compute_caldyn_horiz")

!    CALL wait_message(req_theta_rhodz) 

  DO l = ll_begin, ll_end
!!!  Compute mass and theta fluxes
!    IF (caldyn_conserv==energy) CALL test_message(req_qu) 
!DIR$ SIMD
    DO ij=ij_begin_ext,ij_end_ext         
        hflux(ij+u_right,l)=0.5*(rhodz(ij,l)+rhodz(ij+t_right,l))*u(ij+u_right,l)*le(ij+u_right)
        hflux(ij+u_lup,l)=0.5*(rhodz(ij,l)+rhodz(ij+t_lup,l))*u(ij+u_lup,l)*le(ij+u_lup)
        hflux(ij+u_ldown,l)=0.5*(rhodz(ij,l)+rhodz(ij+t_ldown,l))*u(ij+u_ldown,l)*le(ij+u_ldown)

        Ftheta(ij+u_right,l)=0.5*(theta(ij,l)+theta(ij+t_right,l))*hflux(ij+u_right,l)
        Ftheta(ij+u_lup,l)=0.5*(theta(ij,l)+theta(ij+t_lup,l))*hflux(ij+u_lup,l)
        Ftheta(ij+u_ldown,l)=0.5*(theta(ij,l)+theta(ij+t_ldown,l))*hflux(ij+u_ldown,l)
    ENDDO

!!! compute horizontal divergence of fluxes
!DIR$ SIMD
    DO ij=ij_begin,ij_end         
        ! convm = -div(mass flux), sign convention as in Ringler et al. 2012, eq. 21
        convm(ij,l)= -1./Ai(ij)*(ne_right*hflux(ij+u_right,l)   +  &
                                ne_rup*hflux(ij+u_rup,l)       +  &  
                                ne_lup*hflux(ij+u_lup,l)       +  &  
                                ne_left*hflux(ij+u_left,l)     +  &  
                                ne_ldown*hflux(ij+u_ldown,l)   +  &  
                                ne_rdown*hflux(ij+u_rdown,l))    

        ! signe ? attention d (rho theta dz)
        ! dtheta_rhodz =  -div(flux.theta)
        dtheta_rhodz(ij,l)=-1./Ai(ij)*(ne_right*Ftheta(ij+u_right,l)    +  &
                                 ne_rup*Ftheta(ij+u_rup,l)        +  &  
                                 ne_lup*Ftheta(ij+u_lup,l)        +  &  
                                 ne_left*Ftheta(ij+u_left,l)      +  &  
                                 ne_ldown*Ftheta(ij+u_ldown,l)    +  &  
                                 ne_rdown*Ftheta(ij+u_rdown,l))
    ENDDO

 END DO

!!! Compute potential vorticity (Coriolis) contribution to du
 
    SELECT CASE(caldyn_conserv)
    CASE(energy) ! energy-conserving TRiSK

!       CALL wait_message(req_qu)

        DO l=ll_begin,ll_end
!DIR$ SIMD
          DO ij=ij_begin,ij_end         
    
                uu = wee(ij+u_right,1,1)*hflux(ij+u_rup,l)*(qu(ij+u_right,l)+qu(ij+u_rup,l))+                             &
                     wee(ij+u_right,2,1)*hflux(ij+u_lup,l)*(qu(ij+u_right,l)+qu(ij+u_lup,l))+                             &
                     wee(ij+u_right,3,1)*hflux(ij+u_left,l)*(qu(ij+u_right,l)+qu(ij+u_left,l))+                           &
                     wee(ij+u_right,4,1)*hflux(ij+u_ldown,l)*(qu(ij+u_right,l)+qu(ij+u_ldown,l))+                         &
                     wee(ij+u_right,5,1)*hflux(ij+u_rdown,l)*(qu(ij+u_right,l)+qu(ij+u_rdown,l))+                         & 
                     wee(ij+u_right,1,2)*hflux(ij+t_right+u_ldown,l)*(qu(ij+u_right,l)+qu(ij+t_right+u_ldown,l))+         &
                     wee(ij+u_right,2,2)*hflux(ij+t_right+u_rdown,l)*(qu(ij+u_right,l)+qu(ij+t_right+u_rdown,l))+         &
                     wee(ij+u_right,3,2)*hflux(ij+t_right+u_right,l)*(qu(ij+u_right,l)+qu(ij+t_right+u_right,l))+         &
                     wee(ij+u_right,4,2)*hflux(ij+t_right+u_rup,l)*(qu(ij+u_right,l)+qu(ij+t_right+u_rup,l))+             &
                     wee(ij+u_right,5,2)*hflux(ij+t_right+u_lup,l)*(qu(ij+u_right,l)+qu(ij+t_right+u_lup,l))
                du(ij+u_right,l) = .5*uu/de(ij+u_right)
                
                uu = wee(ij+u_lup,1,1)*hflux(ij+u_left,l)*(qu(ij+u_lup,l)+qu(ij+u_left,l)) +                        &
                     wee(ij+u_lup,2,1)*hflux(ij+u_ldown,l)*(qu(ij+u_lup,l)+qu(ij+u_ldown,l)) +                      &
                     wee(ij+u_lup,3,1)*hflux(ij+u_rdown,l)*(qu(ij+u_lup,l)+qu(ij+u_rdown,l)) +                      &
                     wee(ij+u_lup,4,1)*hflux(ij+u_right,l)*(qu(ij+u_lup,l)+qu(ij+u_right,l)) +                      &
                     wee(ij+u_lup,5,1)*hflux(ij+u_rup,l)*(qu(ij+u_lup,l)+qu(ij+u_rup,l)) +                          & 
                     wee(ij+u_lup,1,2)*hflux(ij+t_lup+u_right,l)*(qu(ij+u_lup,l)+qu(ij+t_lup+u_right,l)) +          &
                     wee(ij+u_lup,2,2)*hflux(ij+t_lup+u_rup,l)*(qu(ij+u_lup,l)+qu(ij+t_lup+u_rup,l)) +              &
                     wee(ij+u_lup,3,2)*hflux(ij+t_lup+u_lup,l)*(qu(ij+u_lup,l)+qu(ij+t_lup+u_lup,l)) +              &
                     wee(ij+u_lup,4,2)*hflux(ij+t_lup+u_left,l)*(qu(ij+u_lup,l)+qu(ij+t_lup+u_left,l)) +            &
                     wee(ij+u_lup,5,2)*hflux(ij+t_lup+u_ldown,l)*(qu(ij+u_lup,l)+qu(ij+t_lup+u_ldown,l))
                du(ij+u_lup,l) = .5*uu/de(ij+u_lup)

                
                uu = wee(ij+u_ldown,1,1)*hflux(ij+u_rdown,l)*(qu(ij+u_ldown,l)+qu(ij+u_rdown,l)) +                      &
                     wee(ij+u_ldown,2,1)*hflux(ij+u_right,l)*(qu(ij+u_ldown,l)+qu(ij+u_right,l)) +                      &
                     wee(ij+u_ldown,3,1)*hflux(ij+u_rup,l)*(qu(ij+u_ldown,l)+qu(ij+u_rup,l)) +                          &
                     wee(ij+u_ldown,4,1)*hflux(ij+u_lup,l)*(qu(ij+u_ldown,l)+qu(ij+u_lup,l)) +                          &
                     wee(ij+u_ldown,5,1)*hflux(ij+u_left,l)*(qu(ij+u_ldown,l)+qu(ij+u_left,l)) +                        & 
                     wee(ij+u_ldown,1,2)*hflux(ij+t_ldown+u_lup,l)*(qu(ij+u_ldown,l)+qu(ij+t_ldown+u_lup,l)) +          &
                     wee(ij+u_ldown,2,2)*hflux(ij+t_ldown+u_left,l)*(qu(ij+u_ldown,l)+qu(ij+t_ldown+u_left,l)) +        &
                     wee(ij+u_ldown,3,2)*hflux(ij+t_ldown+u_ldown,l)*(qu(ij+u_ldown,l)+qu(ij+t_ldown+u_ldown,l)) +      &
                     wee(ij+u_ldown,4,2)*hflux(ij+t_ldown+u_rdown,l)*(qu(ij+u_ldown,l)+qu(ij+t_ldown+u_rdown,l)) +      &
                     wee(ij+u_ldown,5,2)*hflux(ij+t_ldown+u_right,l)*(qu(ij+u_ldown,l)+qu(ij+t_ldown+u_right,l))
                du(ij+u_ldown,l) = .5*uu/de(ij+u_ldown)
                
          ENDDO
       ENDDO
       
    CASE(enstrophy) ! enstrophy-conserving TRiSK
  
        DO l=ll_begin,ll_end
!DIR$ SIMD
          DO ij=ij_begin,ij_end         

                uu = wee(ij+u_right,1,1)*hflux(ij+u_rup,l)+                           &
                     wee(ij+u_right,2,1)*hflux(ij+u_lup,l)+                           &
                     wee(ij+u_right,3,1)*hflux(ij+u_left,l)+                          &
                     wee(ij+u_right,4,1)*hflux(ij+u_ldown,l)+                         &
                     wee(ij+u_right,5,1)*hflux(ij+u_rdown,l)+                         & 
                     wee(ij+u_right,1,2)*hflux(ij+t_right+u_ldown,l)+                 &
                     wee(ij+u_right,2,2)*hflux(ij+t_right+u_rdown,l)+                 &
                     wee(ij+u_right,3,2)*hflux(ij+t_right+u_right,l)+                 &
                     wee(ij+u_right,4,2)*hflux(ij+t_right+u_rup,l)+                   &
                     wee(ij+u_right,5,2)*hflux(ij+t_right+u_lup,l)
                du(ij+u_right,l) = qu(ij+u_right,l)*uu/de(ij+u_right)
                
                
                uu = wee(ij+u_lup,1,1)*hflux(ij+u_left,l)+                        &
                     wee(ij+u_lup,2,1)*hflux(ij+u_ldown,l)+                       &
                     wee(ij+u_lup,3,1)*hflux(ij+u_rdown,l)+                       &
                     wee(ij+u_lup,4,1)*hflux(ij+u_right,l)+                       &
                     wee(ij+u_lup,5,1)*hflux(ij+u_rup,l)+                         & 
                     wee(ij+u_lup,1,2)*hflux(ij+t_lup+u_right,l)+                 &
                     wee(ij+u_lup,2,2)*hflux(ij+t_lup+u_rup,l)+                   &
                     wee(ij+u_lup,3,2)*hflux(ij+t_lup+u_lup,l)+                   &
                     wee(ij+u_lup,4,2)*hflux(ij+t_lup+u_left,l)+                  &
                     wee(ij+u_lup,5,2)*hflux(ij+t_lup+u_ldown,l)
                du(ij+u_lup,l) = qu(ij+u_lup,l)*uu/de(ij+u_lup)

                uu = wee(ij+u_ldown,1,1)*hflux(ij+u_rdown,l)+                      &
                     wee(ij+u_ldown,2,1)*hflux(ij+u_right,l)+                      &
                     wee(ij+u_ldown,3,1)*hflux(ij+u_rup,l)+                        &
                     wee(ij+u_ldown,4,1)*hflux(ij+u_lup,l)+                        &
                     wee(ij+u_ldown,5,1)*hflux(ij+u_left,l)+                       & 
                     wee(ij+u_ldown,1,2)*hflux(ij+t_ldown+u_lup,l)+                &
                     wee(ij+u_ldown,2,2)*hflux(ij+t_ldown+u_left,l)+               &
                     wee(ij+u_ldown,3,2)*hflux(ij+t_ldown+u_ldown,l)+              &
                     wee(ij+u_ldown,4,2)*hflux(ij+t_ldown+u_rdown,l)+              &
                     wee(ij+u_ldown,5,2)*hflux(ij+t_ldown+u_right,l)
                du(ij+u_ldown,l) = qu(ij+u_ldown,l)*uu/de(ij+u_ldown)
     
          ENDDO
       ENDDO
       
    CASE DEFAULT
       STOP
    END SELECT
  
!!! Compute bernouilli term = Kinetic Energy + geopotential
    IF(boussinesq) THEN
       ! first use hydrostatic balance with theta*rhodz to find pk (Lagrange multiplier=pressure) 
       ! uppermost layer
       !DIR$ SIMD
       DO ij=ij_begin_ext,ij_end_ext         
          pk(ij,llm) = ptop + (.5*g)*theta(ij,llm)*rhodz(ij,llm)
       END DO
       ! other layers
       DO l = llm-1, 1, -1
!          !$OMP DO SCHEDULE(STATIC) 
          !DIR$ SIMD
          DO ij=ij_begin_ext,ij_end_ext         
             pk(ij,l) = pk(ij,l+1) + (.5*g)*(theta(ij,l)*rhodz(ij,l)+theta(ij,l+1)*rhodz(ij,l+1))
          END DO
       END DO
       ! surface pressure (for diagnostics) FIXME
       ! DO ij=ij_begin_ext,ij_end_ext         
       !   ps(ij) = pk(ij,1) + (.5*g)*theta(ij,1)*rhodz(ij,1)
       ! END DO
       ! now pk contains the Lagrange multiplier (pressure)

       DO l=ll_begin,ll_end
!DIR$ SIMD
          DO ij=ij_begin,ij_end         
                
                berni(ij,l) = pk(ij,l) + &
                       1/(4*Ai(ij))*(le(ij+u_right)*de(ij+u_right)*u(ij+u_right,l)**2 +    &
                                     le(ij+u_rup)*de(ij+u_rup)*u(ij+u_rup,l)**2 +          &
                                     le(ij+u_lup)*de(ij+u_lup)*u(ij+u_lup,l)**2 +          &
                                     le(ij+u_left)*de(ij+u_left)*u(ij+u_left,l)**2 +       &
                                     le(ij+u_ldown)*de(ij+u_ldown)*u(ij+u_ldown,l)**2 +    &
                                     le(ij+u_rdown)*de(ij+u_rdown)*u(ij+u_rdown,l)**2 )  
                ! from now on pk contains the vertically-averaged geopotential
                pk(ij,l) = .5*(geopot(ij,l)+geopot(ij,l+1))
          ENDDO
       ENDDO

    ELSE ! compressible

       DO l=ll_begin,ll_end
!DIR$ SIMD
            DO ij=ij_begin,ij_end         
                
                berni(ij,l) = .5*(geopot(ij,l)+geopot(ij,l+1)) &
                     + 1/(4*Ai(ij))*(le(ij+u_right)*de(ij+u_right)*u(ij+u_right,l)**2 +    &
                                     le(ij+u_rup)*de(ij+u_rup)*u(ij+u_rup,l)**2 +          &
                                     le(ij+u_lup)*de(ij+u_lup)*u(ij+u_lup,l)**2 +          &
                                     le(ij+u_left)*de(ij+u_left)*u(ij+u_left,l)**2 +       &
                                     le(ij+u_ldown)*de(ij+u_ldown)*u(ij+u_ldown,l)**2 +    &
                                     le(ij+u_rdown)*de(ij+u_rdown)*u(ij+u_rdown,l)**2 )  
          ENDDO
       ENDDO

    END IF ! Boussinesq/compressible

!!! Add gradients of Bernoulli and Exner functions to du
    DO l=ll_begin,ll_end
!DIR$ SIMD
       DO ij=ij_begin,ij_end         
        
             du(ij+u_right,l) = du(ij+u_right,l) + 1/de(ij+u_right) * (             &
                               0.5*(theta(ij,l)+theta(ij+t_right,l))                &
                                  *( ne_right*pk(ij,l)+ne_left*pk(ij+t_right,l))    &
                                   + ne_right*berni(ij,l)+ne_left*berni(ij+t_right,l) )
        
        
             du(ij+u_lup,l) = du(ij+u_lup,l) +  1/de(ij+u_lup) * (                  &
                               0.5*(theta(ij,l)+theta(ij+t_lup,l))                  &
                                  *( ne_lup*pk(ij,l)+ne_rdown*pk(ij+t_lup,l))       &
                                   + ne_lup*berni(ij,l)+ne_rdown*berni(ij+t_lup,l) )
                
             du(ij+u_ldown,l) = du(ij+u_ldown,l) + 1/de(ij+u_ldown) * (             &
                               0.5*(theta(ij,l)+theta(ij+t_ldown,l))                &
                                  *( ne_ldown*pk(ij,l)+ne_rup*pk(ij+t_ldown,l))     &
                                   + ne_ldown*berni(ij,l)+ne_rup*berni(ij+t_ldown,l) )

       ENDDO
    ENDDO
 
    CALL trace_end("compute_caldyn_horiz")

END SUBROUTINE compute_caldyn_horiz

SUBROUTINE compute_caldyn_vert(u,theta,rhodz,convm, wflux,wwuu, dps,dtheta_rhodz,du)
  USE icosa
  USE disvert_mod
  USE exner_mod
  USE trace
  USE omp_para
  IMPLICIT NONE
    REAL(rstd),INTENT(IN)  :: u(iim*3*jjm,llm)
    REAL(rstd),INTENT(IN)  :: theta(iim*jjm,llm)
    REAL(rstd),INTENT(IN)  :: rhodz(iim*jjm,llm)
    REAL(rstd),INTENT(INOUT)  :: convm(iim*jjm,llm)  ! mass flux convergence

    REAL(rstd),INTENT(INOUT) :: wflux(iim*jjm,llm+1) ! vertical mass flux (kg/m2/s)
    REAL(rstd),INTENT(INOUT) :: wwuu(iim*3*jjm,llm+1)
    REAL(rstd),INTENT(OUT) :: du(iim*3*jjm,llm)
    REAL(rstd),INTENT(OUT) :: dtheta_rhodz(iim*jjm,llm)
    REAL(rstd),INTENT(OUT) :: dps(iim*jjm)

! temporary variable    
    INTEGER :: i,j,ij,l
    REAL(rstd) :: p_ik, exner_ik

!    REAL(rstd) :: wwuu(iim*3*jjm,llm+1) ! tmp var, don't know why but gain 30% on the whole code in opemp
                                        ! need to be understood

!    wwuu=wwuu_out
  CALL trace_start("compute_caldyn_vert")

!!$OMP BARRIER   
!!! cumulate mass flux convergence from top to bottom
  DO  l = llm-1, 1, -1
!    IF (caldyn_conserv==energy) CALL test_message(req_qu) 

!!$OMP DO SCHEDULE(STATIC) 
!DIR$ SIMD
    DO ij=ij_begin,ij_end         
        convm(ij,l) = convm(ij,l) + convm(ij,l+1)
    ENDDO
  ENDDO

! IMPLICIT FLUSH on convm
!!!!!!!!!!!!!!!!!!!!!!!!!

! compute dps
  IF (omp_first) THEN
!DIR$ SIMD
     DO ij=ij_begin,ij_end         
        ! dps/dt = -int(div flux)dz
        dps(ij) = convm(ij,1) * g 
    ENDDO
  ENDIF
  
!!! Compute vertical mass flux (l=1,llm+1 done by caldyn_BC)
  DO l=ll_beginp1,ll_end
!    IF (caldyn_conserv==energy) CALL test_message(req_qu) 
!DIR$ SIMD
      DO ij=ij_begin,ij_end         
        ! w = int(z,ztop,div(flux)dz) + B(eta)dps/dt
        ! => w>0 for upward transport
        wflux( ij, l ) = bp(l) * convm( ij, 1 ) - convm( ij, l ) 
    ENDDO
  ENDDO

  DO l=ll_begin,ll_endm1
!DIR$ SIMD
     DO ij=ij_begin,ij_end         
        dtheta_rhodz(ij, l ) = dtheta_rhodz(ij, l )  -   0.5 * (  wflux(ij,l+1) * (theta(ij,l) + theta(ij,l+1)))
    ENDDO
  ENDDO
  
  DO l=ll_beginp1,ll_end
!DIR$ SIMD
      DO ij=ij_begin,ij_end         
        dtheta_rhodz(ij, l ) = dtheta_rhodz(ij, l )  +   0.5 * ( wflux(ij,l  ) * (theta(ij,l-1) + theta(ij,l) ) )
    ENDDO
  ENDDO
  
! Compute vertical transport
  DO l=ll_beginp1,ll_end
!DIR$ SIMD
      DO ij=ij_begin,ij_end         
        wwuu(ij+u_right,l) = 0.5*( wflux(ij,l) + wflux(ij+t_right,l)) * (u(ij+u_right,l) - u(ij+u_right,l-1))
        wwuu(ij+u_lup,l) = 0.5* ( wflux(ij,l) + wflux(ij+t_lup,l)) * (u(ij+u_lup,l) - u(ij+u_lup,l-1))
        wwuu(ij+u_ldown,l) = 0.5*( wflux(ij,l) + wflux(ij+t_ldown,l)) * (u(ij+u_ldown,l) - u(ij+u_ldown,l-1))
     ENDDO
   ENDDO

 !--> flush wwuu  
 ! !$OMP BARRIER

! Add vertical transport to du
  DO l=ll_begin,ll_end
!DIR$ SIMD
     DO ij=ij_begin,ij_end         
        du(ij+u_right, l )   = du(ij+u_right,l)  - (wwuu(ij+u_right,l+1)+ wwuu(ij+u_right,l)) / (rhodz(ij,l)+rhodz(ij+t_right,l))
        du(ij+u_lup, l )     = du(ij+u_lup,l)    - (wwuu(ij+u_lup,l+1)  + wwuu(ij+u_lup,l))   / (rhodz(ij,l)+rhodz(ij+t_lup,l))
        du(ij+u_ldown, l )   = du(ij+u_ldown,l)  - (wwuu(ij+u_ldown,l+1)+ wwuu(ij+u_ldown,l)) / (rhodz(ij,l)+rhodz(ij+t_ldown,l))
    ENDDO
  ENDDO      

!  DO l=ll_beginp1,ll_end
!!DIR$ SIMD
!      DO ij=ij_begin,ij_end         
!        wwuu_out(ij+u_right,l) = wwuu(ij+u_right,l)
!        wwuu_out(ij+u_lup,l)   = wwuu(ij+u_lup,l)
!        wwuu_out(ij+u_ldown,l) = wwuu(ij+u_ldown,l)
!     ENDDO
!   ENDDO
  
  CALL trace_end("compute_caldyn_vert")

  END SUBROUTINE compute_caldyn_vert

!-------------------------------- Diagnostics ----------------------------

  SUBROUTINE check_mass_conservation(f_ps,f_dps)
  USE icosa
  USE mpipara
  IMPLICIT NONE
    TYPE(t_field),POINTER :: f_ps(:)
    TYPE(t_field),POINTER :: f_dps(:)
    REAL(rstd),POINTER :: ps(:)
    REAL(rstd),POINTER :: dps(:)
    REAL(rstd) :: mass_tot,dmass_tot
    INTEGER :: ind,i,j,ij
    
    mass_tot=0
    dmass_tot=0
    
    CALL transfert_request(f_dps,req_i1)
    CALL transfert_request(f_ps,req_i1)

    DO ind=1,ndomain
      CALL swap_dimensions(ind)
      CALL swap_geometry(ind)

      ps=f_ps(ind)
      dps=f_dps(ind)

      DO j=jj_begin,jj_end
        DO i=ii_begin,ii_end
          ij=(j-1)*iim+i
          IF (domain(ind)%own(i,j)) THEN
            mass_tot=mass_tot+ps(ij)*Ai(ij)/g
            dmass_tot=dmass_tot+dps(ij)*Ai(ij)/g
          ENDIF
        ENDDO
      ENDDO
   
    ENDDO
    IF (is_mpi_root) PRINT*, "mass_tot ", mass_tot,"      dmass_tot ",dmass_tot        

  END SUBROUTINE check_mass_conservation  
  
  SUBROUTINE write_output_fields(f_ps, f_phis, f_dps, f_u, f_theta_rhodz, f_q, &
       f_buf_i, f_buf_v, f_buf_i3, f_buf1_i, f_buf2_i, f_buf_s, f_buf_p)
    USE icosa
    USE vorticity_mod
    USE theta2theta_rhodz_mod
    USE pression_mod
    USE omega_mod
    USE write_field
    USE vertical_interp_mod
    USE wind_mod
    TYPE(t_field),POINTER :: f_ps(:), f_phis(:), f_u(:), f_theta_rhodz(:), f_q(:), f_dps(:), &
         f_buf_i(:), f_buf_v(:), f_buf_i3(:), f_buf1_i(:), f_buf2_i(:), f_buf_s(:), f_buf_p(:)
    
    REAL(rstd) :: out_pression_lev
    CHARACTER(LEN=255) :: str_pression
    CHARACTER(LEN=255) :: physics_type
    
    out_pression_level=0
    CALL getin("out_pression_level",out_pression_level) 
    WRITE(str_pression,*) INT(out_pression_level/100)
    str_pression=ADJUSTL(str_pression)
    
    CALL writefield("ps",f_ps)
    CALL writefield("dps",f_dps)
    CALL writefield("phis",f_phis)
    CALL vorticity(f_u,f_buf_v)
    CALL writefield("vort",f_buf_v)

    CALL w_omega(f_ps, f_u, f_buf_i)
    CALL writefield('omega', f_buf_i)
    IF (out_pression_level<=preff .AND. out_pression_level > 0) THEN
      CALL vertical_interp(f_ps,f_buf_i,f_buf_s,out_pression_level)
      CALL writefield("omega"//TRIM(str_pression),f_buf_s)
    ENDIF
    
    ! Temperature
!    CALL theta_rhodz2temperature(f_ps,f_theta_rhodz,f_buf_i) ; ! FIXME
     
    CALL getin('physics',physics_type)
    IF (TRIM(physics_type)=='dcmip') THEN
      CALL Tv2T(f_buf_i,f_q,f_buf1_i) 
      CALL writefield("T",f_buf1_i)
      IF (out_pression_level<=preff .AND. out_pression_level > 0) THEN
        CALL vertical_interp(f_ps,f_buf1_i,f_buf_s,out_pression_level)
        CALL writefield("T"//TRIM(str_pression),f_buf_s)
      ENDIF
    ELSE
      CALL writefield("T",f_buf_i)
      IF (out_pression_level<=preff .AND. out_pression_level > 0) THEN
        CALL vertical_interp(f_ps,f_buf_i,f_buf_s,out_pression_level)
        CALL writefield("T"//TRIM(str_pression),f_buf_s)
      ENDIF
    ENDIF
   
    ! velocity components
    CALL un2ulonlat(f_u, f_buf1_i, f_buf2_i)
    CALL writefield("ulon",f_buf1_i)
    CALL writefield("ulat",f_buf2_i)

    IF (out_pression_level<=preff .AND. out_pression_level > 0) THEN
      CALL vertical_interp(f_ps,f_buf1_i,f_buf_s,out_pression_level)
      CALL writefield("ulon"//TRIM(str_pression),f_buf_s)
      CALL vertical_interp(f_ps,f_buf2_i,f_buf_s,out_pression_level)
      CALL writefield("ulat"//TRIM(str_pression),f_buf_s)
    ENDIF
    
    ! geopotential ! FIXME
    CALL thetarhodz2geopot(f_ps,f_phis,f_theta_rhodz, f_buf_s,f_buf_p,f_buf1_i,f_buf2_i,f_buf_i)
    CALL writefield("p",f_buf_p)
    CALL writefield("phi",f_geopot)   ! geopotential
    CALL writefield("theta",f_buf1_i) ! potential temperature
    CALL writefield("pk",f_buf2_i)    ! Exner pressure
  
  END SUBROUTINE write_output_fields
  
  SUBROUTINE thetarhodz2geopot(f_ps,f_phis,f_theta_rhodz, f_pks,f_p,f_theta,f_pk,f_phi) 
    USE field_mod
    USE pression_mod
    USE exner_mod
    USE geopotential_mod
    USE theta2theta_rhodz_mod
    TYPE(t_field), POINTER :: f_ps(:), f_phis(:), f_theta_rhodz(:), &  ! IN
         f_pks(:), f_p(:), f_theta(:), f_pk(:), f_phi(:)               ! OUT
    REAL(rstd),POINTER :: pk(:,:), p(:,:), theta(:,:), theta_rhodz(:,:), &
         phi(:,:), phis(:), ps(:), pks(:)
    INTEGER :: ind

    DO ind=1,ndomain
       IF (.NOT. assigned_domain(ind)) CYCLE
       CALL swap_dimensions(ind)
       CALL swap_geometry(ind)
       ps = f_ps(ind)
       p  = f_p(ind)
       CALL compute_pression(ps,p,0)
       pk = f_pk(ind)
       pks = f_pks(ind)
       CALL compute_exner(ps,p,pks,pk,0)
       theta_rhodz = f_theta_rhodz(ind)
       theta = f_theta(ind)
       CALL compute_theta_rhodz2theta(ps, theta_rhodz,theta,0)
       phis = f_phis(ind)
       phi = f_phi(ind)
       CALL compute_geopotential(phis,pks,pk,theta,phi,0)
    END DO

  END SUBROUTINE thetarhodz2geopot
  
  SUBROUTINE Tv2T(f_Tv, f_q, f_T)
  USE icosa
  IMPLICIT NONE
    TYPE(t_field), POINTER :: f_TV(:)
    TYPE(t_field), POINTER :: f_q(:)
    TYPE(t_field), POINTER :: f_T(:)
    
    REAL(rstd),POINTER :: Tv(:,:), q(:,:,:), T(:,:)
    INTEGER :: ind
    
    DO ind=1,ndomain
       IF (.NOT. assigned_domain(ind)) CYCLE
       CALL swap_dimensions(ind)
       CALL swap_geometry(ind)
       Tv=f_Tv(ind)
       q=f_q(ind)
       T=f_T(ind)
       T=Tv/(1+0.608*q(:,:,1))
    END DO
    
  END SUBROUTINE Tv2T
  
END MODULE caldyn_gcm_mod