MODULE timestep_unstructured_mod
  USE ISO_C_BINDING
  USE OMP_LIB
#ifdef CPP_USING_XIOS
  USE xios
#endif
  USE caldyn_unstructured_mod
  USE transfer_unstructured_mod
  IMPLICIT NONE
  PRIVATE
  SAVE

#ifdef CPP_USING_XIOS
  TYPE(xios_context) :: ctx_hdl
#endif

CONTAINS

#include "unstructured.h90"

#define HASNAN(field) (ANY(.NOT.ABS(field)<1e20))

   SUBROUTINE scalar_laplacian(s,laps) BINDC(scalar_laplacian)
     FIELD_MASS   :: s,laps ! IN, OUT
     FIELD_U      :: grads  ! BUF
     !$OMP PARALLEL NUM_THREADS(nb_threads)
     CALL update_halo(transfer_primal, s)
     CALL compute_scalar_laplacian(s,grads,laps)
     !$OMP END PARALLEL
   END SUBROUTINE scalar_laplacian

   SUBROUTINE curl_laplacian(u,lapu) BINDC(curl_laplacian)
     FIELD_U   :: u,lapu ! IN, OUT
     FIELD_Z   :: qv  ! BUF
     !$OMP PARALLEL NUM_THREADS(nb_threads)
     CALL update_halo(transfer_edge, u)
     CALL compute_curl_laplacian(u,qv,lapu)
     !$OMP END PARALLEL
   END SUBROUTINE curl_laplacian

  !
  SUBROUTINE caldyn_unstructured(tau, mass_col,rhodz,theta_rhodz,u,geopot,w, & ! IN : flow state
                                 theta,ps,pk,hflux,qv, &    ! OUT : diags (except surface geopot : IN)
                                 dmass_col,drhodz,dtheta_rhodz,du_fast,du_slow, &
                                 dPhi_fast, dPhi_slow, dW_fast, dW_slow) BINDC(caldyn_unstructured) ! OUT : tendencies
    NUM, VALUE :: tau
    FIELD_MASS   :: rhodz, drhodz, pk, berni          ! IN, OUT, DIAG
    FIELD_THETA  :: theta_rhodz, dtheta_rhodz, theta  ! IN, OUT, DIAG
    FIELD_GEOPOT :: wflux, w, geopot, &               ! DIAG, INOUT
         dPhi_fast, dPhi_slow, dW_fast, dW_slow       ! OUT
    FIELD_U      :: u,du_fast,du_slow,hflux,qu        ! INOUT,OUT,OUT,DIAG
    FIELD_Z      :: qv                                ! DIAG
    FIELD_PS     :: ps,dmass_col,mass_col             ! OUT,OUT,IN (if eta_mass) or OUT,UNUSED,UNUSED (if eta_lag)
    ! buffers for fields that need to be shared among OpenMP threads
    ! vertical size is allocated as llm+1 even if only llm is needed
    FIELD_GEOPOT :: bufm1, bufm2, bufm3
    FIELD_UL     :: bufu1, bufu2, bufu3,bufu4

    TIME          :: time1,time2
    INTEGER :: ij
    
    !  CALL CPU_TIME(time1)
    time1=OMP_GET_WTIME()
    
    IF(hydrostatic) THEN
       
       !$OMP PARALLEL NUM_THREADS(nb_threads)
       !$OMP DO SCHEDULE(STATIC)
       DO ij=1,edge_num
          du_fast(:,ij)=0.
          du_slow(:,ij)=0.
       END DO
       !$OMP END DO
       CALL compute_theta(mass_col,rhodz,theta_rhodz, theta)
       CALL compute_geopot(rhodz,theta, ps,pk,geopot)
       
       CALL compute_caldyn_fast(tau, pk,berni,theta,geopot, du_fast,u)
       
       CALL compute_pvort_only(rhodz,u,qv,qu)
       CALL compute_caldyn_slow_hydro(rhodz,theta,u, berni,hflux,du_slow)
       CALL compute_coriolis(hflux,theta,qu, bufu1, drhodz,dtheta_rhodz,du_slow)
       IF(caldyn_eta == eta_mass) THEN
          CALL caldyn_vert(drhodz,rhodz,theta,u, dmass_col,wflux,dtheta_rhodz,du_slow,bufu1)
       END IF
       !$OMP END PARALLEL
       
    ELSE ! NH
       DO ij=1,edge_num
          du_fast(:,ij)=0.
          du_slow(:,ij)=0.
       END DO
       DO ij=1,primal_num
          wflux(1,ij)=0.
          wflux(llm+1,ij)=0.
       END DO
       CALL compute_theta(mass_col,rhodz,theta_rhodz, theta)
       CALL compute_caldyn_solver(tau,rhodz,theta, bufm1,bufm2,bufm3, pk,geopot,W,dPhi_fast,dW_fast,du_fast)
       CALL compute_caldyn_fast(tau, pk,berni,theta,geopot, du_fast,u)
       CALL compute_pvort_only(rhodz,u,qv,qu)
       CALL compute_caldyn_slow_NH(u,rhodz,geopot,W, bufm1,bufm2,bufm3,bufu1,bufu2,bufu3,bufu4, hflux,du_slow,dPhi_slow,dW_slow) 
       CALL compute_coriolis(hflux,theta,qu, bufu1, drhodz,dtheta_rhodz,du_slow)
       IF(caldyn_eta == eta_mass) THEN
          CALL caldyn_vert(drhodz,rhodz,theta,u, dmass_col,wflux,dtheta_rhodz,du_slow,bufu1)
          CALL compute_caldyn_vert_NH(rhodz,geopot,W,wflux, bufm1,bufm2,bufm3, du_slow,dPhi_slow,dW_slow)
       END IF
    END IF
    
    time2=OMP_GET_WTIME()
    !  CALL CPU_TIME(time2)
    IF(time2>time1) elapsed = elapsed + time2-time1

    CALL print_trace()
  END SUBROUTINE caldyn_unstructured
  !
  !----------------------------- Time stepping -------------------------------
  
  !
  SUBROUTINE ARK_step(nstep, mass_col,rhodz,theta_rhodz,u,geopot,w, & ! INOUT : flow state
       theta,ps,pk,hflux,qv, &    ! OUT : diags (except surface geopot : IN)
       dmass_col,drhodz,dtheta_rhodz,du_fast,du_slow, &
       dPhi_fast, dPhi_slow, dW_fast, dW_slow) BINDC(ARK_step) ! OUT : tendencies
    INTEGER, VALUE :: nstep ! advance by nstep time steps
    FIELD_PS     :: mass_col, ps                     ! OUT,IN (if eta_mass) or OUT,UNUSED (if eta_lag)
    FIELD_MASS   :: rhodz, pk, berni                 ! IN, DIAG
    FIELD_THETA  :: theta_rhodz, theta               ! IN, DIAG
    FIELD_U      :: u,hflux,qu                       ! INOUT,DIAG*2
    FIELD_GEOPOT :: geopot, w, wflux                 ! IN, INOUT, DIAG
    FIELD_Z      :: qv                               ! DIAG
    NUM2(       primal_num, max_nb_stage) :: dmass_col    ! OUT
    NUM3(llm,   primal_num, max_nb_stage) :: drhodz   ! OUT
    NUM4(llm,   primal_num, nqdyn, max_nb_stage) :: dtheta_rhodz ! OUT
    NUM3(llm,   edge_num,   max_nb_stage) :: du_fast,du_slow ! OUT
    NUM3(llm+1, primal_num, max_nb_stage) :: &
         dPhi_fast, dPhi_slow, dW_fast, dW_slow      ! OUT
    ! buffers for fields that need to be shared among OpenMP threads
    ! vertical size is allocated as llm+1 even if only llm is needed
    FIELD_GEOPOT :: bufm1, bufm2, bufm3
    FIELD_UL     :: bufu1, bufu2, bufu3,bufu4
    TIME       :: time1,time2
    INTEGER :: step, stage, iq, ij
    
    !CALL CPU_TIME(time1)
    time1=OMP_GET_WTIME()
    
    !$OMP PARALLEL NUM_THREADS(nb_threads)
    
    !$OMP DO SCHEDULE(STATIC)
    DO ij=1,primal_num
       wflux(1,ij)=0.
       wflux(llm+1,ij)=0.
    END DO
    !$OMP END DO

    DO step=1, nstep
       DO stage=1, nb_stage
          
          CALL update_halo(transfer_primal, rhodz)
          DO iq=1, nqdyn
             CALL update_halo(transfer_primal, theta_rhodz(:,:,iq))
          END DO

          IF(hydrostatic) THEN
             
             !$OMP DO SCHEDULE(STATIC)
             DO ij=1,edge_num
                du_fast(:,ij,stage)=0.
                du_slow(:,ij,stage)=0.
             END DO
	     !$OMP END DO
             
             CALL compute_theta(mass_col,rhodz,theta_rhodz, theta)
             CALL compute_geopot(rhodz,theta, ps,pk,geopot)
             
             CALL compute_caldyn_fast(tauj(stage), pk,berni,theta,geopot, du_fast(:,:,stage), u)
             CALL update_halo(transfer_edge, u)
             CALL compute_pvort_only(rhodz,u,qv,qu)
             CALL update_halo(transfer_edge, qu)
	    
             CALL compute_caldyn_slow_hydro(rhodz,theta,u, berni,hflux,du_slow(:,:,stage))
             CALL compute_coriolis(hflux,theta,qu, bufu1, drhodz(:,:,stage), &
                  dtheta_rhodz(:,:,:,stage),du_slow(:,:,stage))
             IF(caldyn_eta == eta_mass) THEN
                CALL caldyn_vert(drhodz(:,:,stage),rhodz,theta,u, &
                     dmass_col(:,stage),wflux,dtheta_rhodz(:,:,:,stage),du_slow(:,:,stage),bufu1)
             END IF
             
          ELSE ! NH
             CALL update_halo(transfer_primal, geopot)
             CALL update_halo(transfer_primal, W)
             
             CALL compute_theta(mass_col,rhodz,theta_rhodz, theta)
             CALL compute_caldyn_solver(tauj(stage),rhodz,theta, bufm1,bufm2,bufm3, pk,geopot,W, &
                  dPhi_fast(:,:,stage), dW_fast(:,:,stage), du_fast(:,:,stage))

             CALL compute_caldyn_fast(tauj(stage), pk,berni,theta,geopot, du_fast(:,:,stage),u)
             CALL update_halo(transfer_edge, u)
             CALL compute_pvort_only(rhodz,u,qv,qu)
             CALL update_halo(transfer_edge, qu)

             CALL compute_caldyn_slow_NH(u,rhodz,geopot,W, bufm1,bufm2,bufm3, &
                  bufu1,bufu2,bufu3,bufu4, hflux, &
                  du_slow(:,:,stage), dPhi_slow(:,:,stage), dW_slow(:,:,stage))
             CALL compute_coriolis(hflux,theta,qu, bufu1, &
                  drhodz(:,:,stage),dtheta_rhodz(:,:,:,stage), du_slow(:,:,stage))
             IF(caldyn_eta == eta_mass) THEN
                CALL caldyn_vert(drhodz(:,:,stage),rhodz,theta,u, &
                     dmass_col(:,stage),wflux,dtheta_rhodz(:,:,:,stage), du_slow(:,:,stage), bufu1)
                CALL compute_caldyn_vert_NH(rhodz,geopot,W,wflux, bufm1,bufm2,bufm3, &
                     du_slow(:,:,stage), dPhi_slow(:,:,stage), dW_slow(:,:,stage) )
             END IF
          END IF ! NH
          
          ! FIXME : 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
          
#define UPDATE(clj,field,dfield) update(stage,SIZE(field),clj,field,dfield)
          
          CALL UPDATE(cslj, rhodz, drhodz)
          CALL UPDATE(cslj, theta_rhodz, dtheta_rhodz)
          CALL UPDATE(cslj, u, du_slow)
          CALL UPDATE(cflj, u, du_fast)
          
          IF(.NOT.hydrostatic) THEN
             CALL UPDATE(cslj, geopot, dPhi_slow)
             CALL UPDATE(cflj, geopot, dPhi_fast)
             CALL UPDATE(cslj, W, dW_slow)
             CALL UPDATE(cflj, W, dW_fast)
          END IF
          
       END DO
    END DO
    !$OMP END PARALLEL
    
    time2=OMP_GET_WTIME()
    !  CALL CPU_TIME(time2)
    IF(time2>time1) elapsed = elapsed + time2-time1

    CALL print_trace()
    
  END SUBROUTINE ARK_step
  
  
  SUBROUTINE update(j,sz,clj,field,dfield)
    INTEGER :: j, sz ! stage in ARK scheme, field size
    NUM2(max_nb_stage,max_nb_stage) :: clj ! modified Butcher tableau
    NUM1(sz) :: field
    NUM2(sz, max_nb_stage) :: dfield
    !
    INTEGER :: l, ij
    CALL enter_trace(id_update, 8*sz*(j+1) )
    !  PRINT *, clj(1:j,j)
    SELECT CASE(j)
    CASE(1)
       !$OMP DO SCHEDULE(static)
       DO ij=1,sz
          field(ij) = field(ij) &
               + clj(1,j)*dfield(ij,1)
       END DO
    CASE(2)
       !$OMP DO SCHEDULE(static)
       DO ij=1,sz
          field(ij) = field(ij) &
               + clj(1,j)*dfield(ij,1) &
               + clj(2,j)*dfield(ij,2)
       END DO
    CASE(3)
       !$OMP DO SCHEDULE(static)
       DO ij=1,sz
          field(ij) = field(ij) &
               + clj(1,j)*dfield(ij,1) &
               + clj(2,j)*dfield(ij,2) &
               + clj(3,j)*dfield(ij,3)
          
       END DO
    CASE(4)
       !$OMP DO SCHEDULE(static)
       DO ij=1,sz
          field(ij) = field(ij) &
               + clj(1,j)*dfield(ij,1) &
               + clj(2,j)*dfield(ij,2) &
               + clj(3,j)*dfield(ij,3) &
               + clj(4,j)*dfield(ij,4)
       END DO
    CASE(5)
       !$OMP DO SCHEDULE(static)
       DO ij=1,sz
          field(ij) = field(ij) &
               + clj(1,j)*dfield(ij,1) &
               + clj(2,j)*dfield(ij,2) &
               + clj(3,j)*dfield(ij,3) &
               + clj(4,j)*dfield(ij,4) &
               + clj(5,j)*dfield(ij,5)
       END DO
    END SELECT
    CALL exit_trace()
  END SUBROUTINE update
  
  !---------------------------------------------- XIOS ----------------------------------------

#ifdef CPP_USING_XIOS
  
  SUBROUTINE setup_xios() BINDC(setup_xios)
    ! MPI_INIT / MPI_finalize are assumed to be called BEFORE/AFTER this routine
    ! This is the case when calling from Python after importing mpi4py
    INTEGER :: ierr, mpi_threading_mode
    PRINT *, 'Initialize XIOS and obtain our communicator'
    CALL xios_initialize("icosagcm",return_comm=comm_icosa)
    CALL xios_context_initialize("icosagcm",comm_icosa)
    CALL xios_get_handle("icosagcm",ctx_hdl)
    CALL xios_set_current_context(ctx_hdl)
  END SUBROUTINE setup_xios

  SUBROUTINE call_xios_set_timestep(dt) BINDC(xios_set_timestep)
    DBL, VALUE :: dt
    TYPE(xios_duration)      :: dtime
    dtime%second=dt
    CALL xios_set_timestep(dtime)
  END SUBROUTINE call_xios_set_timestep

  SUBROUTINE call_xios_update_calendar(step) BINDC(xios_update_calendar)
    INDEX, value :: step
    CALL xios_update_calendar(step)
  END SUBROUTINE call_xios_update_calendar

#endif

END MODULE timestep_unstructured_mod