MODULE timeloop_gcm_mod
  USE icosa
  USE disvert_mod
  USE trace
  USE omp_para
  USE euler_scheme_mod
  USE explicit_scheme_mod
  USE hevi_scheme_mod
  IMPLICIT NONE
  PRIVATE
  
  INTEGER, PARAMETER :: itau_sync=10
  TYPE(t_message),SAVE :: req_ps0, req_mass0, req_theta_rhodz0, req_u0, req_q0, req_W0, req_geopot0
  LOGICAL, SAVE :: positive_theta

  PUBLIC :: init_timeloop, timeloop

CONTAINS
  
  SUBROUTINE init_timeloop
    USE dissip_gcm_mod
    USE observable_mod
    USE caldyn_mod
    USE etat0_mod
    USE guided_mod
    USE advect_tracer_mod
    USE check_conserve_mod
    USE output_field_mod
    USE theta2theta_rhodz_mod
    USE sponge_mod

    CHARACTER(len=255) :: def

    CALL init_caldyn
    
!    IF (xios_output) itau_out=1
    IF (.NOT. enable_io) itau_out=HUGE(itau_out)

    positive_theta=.FALSE.
    CALL getin('positive_theta',positive_theta)
    IF(positive_theta .AND. nqtot<1) THEN
       PRINT *, 'nqtot must be >0 if positive_theta is .TRUE.'
       STOP
    END IF

    def='ARK2.3'
    CALL getin('time_scheme',def)
    SELECT CASE (TRIM(def))
    CASE('euler')
       scheme_family=explicit
       scheme=euler
       nb_stage=1
    CASE ('runge_kutta')
       scheme_family=explicit
       scheme=rk4
       nb_stage=4
    CASE ('RK2.5')
       scheme_family=explicit
       scheme=rk25
       nb_stage=5
    CASE ('leapfrog_matsuno')
       scheme_family=explicit
       scheme=mlf
       matsuno_period=5
       CALL getin('matsuno_period',matsuno_period)
       nb_stage=matsuno_period+1
    CASE('ARK2.3')
       scheme_family=hevi
       scheme=ark23
       nb_stage=3
       CALL set_coefs_ark23(dt)
    CASE('ARK3.3')
       scheme_family=hevi
       scheme=ark33
       nb_stage=3
       CALL set_coefs_ark33(dt)
    CASE ('none')
       nb_stage=0
    CASE default
       PRINT*,'Bad selector for variable scheme : <', TRIM(def),             &
            ' > options are <euler>, <runge_kutta>, <leapfrog_matsuno>,<RK2.5>,<ARK2.3>'
       STOP
    END SELECT
    
    ! Time-independant orography
    CALL allocate_field(f_phis,field_t,type_real,name='phis')
    ! Model state at current time step
    CALL allocate_field(f_ps,field_t,type_real, name='ps')
    CALL allocate_field(f_mass,field_t,type_real,llm,name='mass')
    CALL allocate_field(f_rhodz,field_t,type_real,llm,name='rhodz')
    CALL allocate_field(f_theta_rhodz,field_t,type_real,llm,nqdyn,name='theta_rhodz')
    CALL allocate_field(f_u,field_u,type_real,llm,name='u')
    CALL allocate_field(f_geopot,field_t,type_real,llm+1,name='geopot')
    CALL allocate_field(f_W,field_t,type_real,llm+1,name='W')
    CALL allocate_field(f_q,field_t,type_real,llm,nqtot,'q')
    ! Mass fluxes
    CALL allocate_field(f_hflux,field_u,type_real,llm)    ! instantaneous mass fluxes
    CALL allocate_field(f_hfluxt,field_u,type_real,llm)   ! mass "fluxes" accumulated in time
    CALL allocate_field(f_wflux,field_t,type_real,llm+1)  ! vertical mass fluxes
    CALL allocate_field(f_wfluxt,field_t,type_real,llm+1,name='wfluxt')
    
    SELECT CASE(scheme_family)
    CASE(explicit)
       ! Trends
       CALL allocate_field(f_dps,field_t,type_real,name='dps')
       CALL allocate_field(f_dmass,field_t,type_real,llm, name='dmass')
       CALL allocate_field(f_dtheta_rhodz,field_t,type_real,llm,nqdyn,name='dtheta_rhodz')
       CALL allocate_field(f_du,field_u,type_real,llm,name='du')
       ! Model state at previous time step (RK/MLF)
       CALL allocate_field(f_psm1,field_t,type_real,name='psm1')
       CALL allocate_field(f_massm1,field_t,type_real,llm, name='massm1')
       CALL allocate_field(f_theta_rhodzm1,field_t,type_real,llm,nqdyn,name='theta_rhodzm1')
       CALL allocate_field(f_um1,field_u,type_real,llm,name='um1')
    CASE(hevi)
       ! Trends
       CALL allocate_fields(nb_stage,f_dps_slow, field_t,type_real,name='dps_slow')
       CALL allocate_fields(nb_stage,f_dmass_slow, field_t,type_real,llm, name='dmass_slow')
       CALL allocate_fields(nb_stage,f_dtheta_rhodz_slow, field_t,type_real,llm,nqdyn,name='dtheta_rhodz_fast')
       CALL allocate_fields(nb_stage,f_du_slow, field_u,type_real,llm,name='du_slow')
       CALL allocate_fields(nb_stage,f_du_fast, field_u,type_real,llm,name='du_fast')
       CALL allocate_fields(nb_stage,f_dW_slow, field_t,type_real,llm+1,name='dW_slow')
       CALL allocate_fields(nb_stage,f_dW_fast, field_t,type_real,llm+1,name='dW_fast')
       CALL allocate_fields(nb_stage,f_dPhi_slow, field_t,type_real,llm+1,name='dPhi_slow')
       CALL allocate_fields(nb_stage,f_dPhi_fast, field_t,type_real,llm+1,name='dPhi_fast')
       f_dps => f_dps_slow(:,1)
       f_du => f_du_slow(:,1)
       f_dtheta_rhodz => f_dtheta_rhodz_slow(:,1)
    END SELECT

    SELECT CASE(scheme)
    CASE(mlf)
       ! Model state 2 time steps ago (MLF)
       CALL allocate_field(f_psm2,field_t,type_real)
       CALL allocate_field(f_massm2,field_t,type_real,llm)
       CALL allocate_field(f_theta_rhodzm2,field_t,type_real,llm,nqdyn)
       CALL allocate_field(f_um2,field_u,type_real,llm)
    END SELECT

    CALL init_theta2theta_rhodz
    CALL init_dissip
    CALL init_sponge
    CALL init_observable
    CALL init_guided
    CALL init_advect_tracer
    CALL init_check_conserve

    CALL etat0(f_ps,f_mass,f_phis,f_theta_rhodz,f_u, f_geopot,f_W, f_q)

    CALL transfert_request(f_phis,req_i0) 
    CALL transfert_request(f_phis,req_i1) 

    CALL init_message(f_ps,req_i0,req_ps0)
    CALL init_message(f_mass,req_i0,req_mass0)
    CALL init_message(f_theta_rhodz,req_i0,req_theta_rhodz0)
    CALL init_message(f_u,req_e0_vect,req_u0)
    CALL init_message(f_q,req_i0,req_q0)
    CALL init_message(f_geopot,req_i0,req_geopot0)
    CALL init_message(f_W,req_i0,req_W0)

  END SUBROUTINE init_timeloop
  
  SUBROUTINE timeloop
    USE dissip_gcm_mod
    USE sponge_mod
    USE observable_mod
    USE etat0_mod
    USE guided_mod
    USE caldyn_mod
    USE advect_tracer_mod
    USE diagflux_mod
    USE physics_mod
    USE mpipara
    USE transfert_mod
    USE check_conserve_mod
    USE xios_mod
    USE output_field_mod
    USE write_etat0_mod
    USE restart_mod
    USE checksum_mod
    USE explicit_scheme_mod
    USE hevi_scheme_mod
    REAL(rstd),POINTER :: rhodz(:,:), mass(:,:), ps(:)

    REAL(rstd) :: adv_over_out ! ratio itau_adv/itau_out
    INTEGER :: ind, it,i,j,l,n,  stage
    LOGICAL :: fluxt_zero(ndomain) ! set to .TRUE. to start accumulating mass fluxes in time
    LOGICAL, PARAMETER :: check_rhodz=.FALSE.
    INTEGER :: start_clock, stop_clock, rate_clock
    LOGICAL,SAVE :: first_physic=.TRUE.
    !$OMP THREADPRIVATE(first_physic)    

    CALL switch_omp_distrib_level
    CALL caldyn_BC(f_phis, f_geopot, f_wflux) ! set constant values in first/last interfaces

    !$OMP BARRIER
    DO ind=1,ndomain
       IF (.NOT. assigned_domain(ind)) CYCLE
       CALL swap_dimensions(ind)
       CALL swap_geometry(ind)
       rhodz=f_rhodz(ind); mass=f_mass(ind); ps=f_ps(ind)
       IF(caldyn_eta==eta_mass) THEN
          CALL compute_rhodz(.TRUE., ps, rhodz) ! save rhodz for transport scheme before dynamics update ps
       ELSE
          DO l=ll_begin,ll_end
             rhodz(:,l)=mass(:,l)
          ENDDO
       END IF
    END DO
    !$OMP BARRIER
    fluxt_zero=.TRUE.

    IF(positive_theta) CALL copy_theta_to_q(f_theta_rhodz,f_rhodz,f_q)
    IF(diagflux_on) THEN
       adv_over_out = itau_adv*(1./itau_out)
    ELSE
       adv_over_out = 0.
    END IF

    CALL check_conserve(f_ps,f_dps,f_u,f_theta_rhodz,f_phis,itau0)  

    Call trace_on

    IF (xios_output) THEN ! we must call update_calendar before any XIOS output
      IF (is_omp_master) CALL xios_update_calendar(1)
    END IF
    
!    IF (write_start) CALL write_etat0(itau0,f_ps, f_phis,f_theta_rhodz,f_u,f_q)
    IF (write_start) CALL write_etat0(itau0,f_ps, f_phis,f_theta_rhodz,f_u,f_q, f_geopot, f_W)
   
    CALL write_output_fields_basic(.TRUE., f_phis, f_ps, f_mass, f_geopot, f_theta_rhodz, f_u, f_W, f_q)

    !$OMP MASTER
    CALL SYSTEM_CLOCK(start_clock)
    CALL SYSTEM_CLOCK(count_rate=rate_clock)
    !$OMP END MASTER   

    DO it=itau0+1,itau0+itaumax

       IF (is_master) CALL print_iteration(it, itau0, itaumax, start_clock, rate_clock)

       IF (xios_output) THEN

          IF(it>itau0+1 .AND. is_omp_master) CALL xios_update_calendar(it-itau0)
       ELSE
          CALL update_time_counter(dt*it)
       END IF

       IF (it==itau0+1 .OR. MOD(it,itau_sync)==0) THEN
          CALL send_message(f_ps,req_ps0)
          CALL wait_message(req_ps0)
          CALL send_message(f_mass,req_mass0)
          CALL wait_message(req_mass0)
          CALL send_message(f_theta_rhodz,req_theta_rhodz0) 
          CALL wait_message(req_theta_rhodz0) 
          CALL send_message(f_u,req_u0)
          CALL wait_message(req_u0)
          CALL send_message(f_q,req_q0) 
          CALL wait_message(req_q0)
          IF(.NOT. hydrostatic) THEN
             CALL send_message(f_geopot,req_geopot0)
             CALL wait_message(req_geopot0)
             CALL send_message(f_W,req_W0)
             CALL wait_message(req_W0)
          END IF
       ENDIF

       CALL guided(it*dt,f_ps,f_theta_rhodz,f_u,f_q)

       SELECT CASE(scheme_family)
       CASE(explicit)
          CALL explicit_scheme(it, fluxt_zero)
       CASE(hevi)
          CALL HEVI_scheme(it, fluxt_zero)
       END SELECT

       IF (MOD(it,itau_dissip)==0) THEN
          
          IF(caldyn_eta==eta_mass) THEN
             !ym flush ps
             !$OMP BARRIER
             DO ind=1,ndomain
                IF (.NOT. assigned_domain(ind)) CYCLE
                CALL swap_dimensions(ind)
                CALL swap_geometry(ind)
                mass=f_mass(ind); ps=f_ps(ind);
                CALL compute_rhodz(.TRUE., ps, mass)
             END DO
          ENDIF
          
          CALL check_conserve_detailed(it, AAM_dyn, &
               f_ps,f_dps,f_u,f_theta_rhodz,f_phis)
       
          CALL dissip(f_ps,f_mass,f_phis,f_geopot,f_theta_rhodz,f_u, f_dtheta_rhodz,f_du)
          
          CALL euler_scheme(.FALSE.)  ! update only u, theta
          IF (iflag_sponge > 0) THEN
             CALL sponge(f_u,f_du,f_theta_rhodz,f_dtheta_rhodz)
             CALL euler_scheme(.FALSE.)  ! update only u, theta
          END IF

          CALL check_conserve_detailed(it, AAM_dissip, &
               f_ps,f_dps,f_u,f_theta_rhodz,f_phis)
       END IF
       
       IF(MOD(it,itau_adv)==0) THEN
          CALL advect_tracer(f_hfluxt,f_wfluxt,f_u, f_q,f_rhodz, & ! update q and rhodz after RK step
               adv_over_out, f_masst,f_qmasst,f_massfluxt, f_qfluxt)  ! accumulate mass and fluxes if diagflux_on
          fluxt_zero=.TRUE. ! restart accumulation of hfluxt and qfluxt at next call to explicit_scheme / HEVI_scheme
          ! At this point advect_tracer has obtained the halos of u and rhodz,
          ! needed for correct computation of kinetic energy
          IF(diagflux_on) CALL diagflux_energy(adv_over_out, f_phis,f_rhodz,f_theta_rhodz,f_u, f_geopot,f_theta, f_hfluxt)

          IF (check_rhodz) THEN ! check that rhodz is consistent with ps
             DO ind=1,ndomain
                IF (.NOT. assigned_domain(ind)) CYCLE
                CALL swap_dimensions(ind)
                CALL swap_geometry(ind)
                rhodz=f_rhodz(ind); ps=f_ps(ind);
                CALL compute_rhodz(.FALSE., ps, rhodz)   
             END DO
          ENDIF
          IF(positive_theta) CALL copy_q_to_theta(f_theta_rhodz,f_rhodz,f_q)
       END IF
       
       IF (MOD(it,itau_physics)==0) THEN
          CALL check_conserve_detailed(it, AAM_dyn, &
            f_ps,f_dps,f_u,f_theta_rhodz,f_phis)
          CALL physics(it,f_phis, f_ps, f_theta_rhodz, f_u, f_wflux, f_q)        
          CALL check_conserve_detailed(it, AAM_phys, &
               f_ps,f_dps,f_u,f_theta_rhodz,f_phis)
          !$OMP MASTER
          IF (first_physic) CALL SYSTEM_CLOCK(start_clock)
          !$OMP END MASTER   
          first_physic=.FALSE.
       END IF

       IF (MOD(it,itau_check_conserv)==0) THEN
          CALL check_conserve_detailed(it, AAM_dyn, &
               f_ps,f_dps,f_u,f_theta_rhodz,f_phis)
          CALL check_conserve(f_ps,f_dps,f_u,f_theta_rhodz,f_phis,it) 
       ENDIF

       IF (mod(it,itau_out)==0 ) THEN
          CALL transfert_request(f_u,req_e1_vect)
          CALL write_output_fields_basic(.FALSE.,f_phis, f_ps, f_mass, f_geopot, f_theta_rhodz, f_u, f_W, f_q)
       ENDIF

    END DO
    
!    CALL write_etat0(itau0+itaumax,f_ps, f_phis,f_theta_rhodz,f_u,f_q) 
    CALL write_etat0(itau0+itaumax,f_ps, f_phis,f_theta_rhodz,f_u,f_q, f_geopot, f_W) 
    
    CALL check_conserve_detailed(it, AAM_dyn, &
         f_ps,f_dps,f_u,f_theta_rhodz,f_phis)
    CALL check_conserve(f_ps,f_dps,f_u,f_theta_rhodz,f_phis,it)  
    
    !$OMP MASTER
    CALL SYSTEM_CLOCK(stop_clock)
    
    IF (mpi_rank==0) THEN 
       PRINT *,"Time elapsed : ",(stop_clock-start_clock)*1./rate_clock 
    ENDIF
    !$OMP END MASTER 
    
    ! CONTAINS
  END SUBROUTINE timeloop

  SUBROUTINE print_iteration(it,itau0,itaumax,start_clock,rate_clock)
    INTEGER :: it, itau0, itaumax, start_clock, stop_clock, rate_clock, throughput
    REAL :: per_step,total, elapsed
    WRITE(*,'(A,I7,A,F14.1)') "It No :",it,"   t :",dt*it
    IF(MOD(it,10)==0) THEN
       CALL SYSTEM_CLOCK(stop_clock)
       elapsed = (stop_clock-start_clock)*1./rate_clock
       per_step = elapsed/(it-itau0)
       throughput = dt/per_step
       total = per_step*(itaumax-itau0)
       WRITE(*,'(A,I5,A,F6.2,A,I6)') 'Time spent (s):',INT(elapsed), &
            '  -- ms/step : ', 1000*per_step, &
            '  -- Throughput :', throughput
       WRITE(*,'(A,I5,A,I5)') 'Whole job (min) :', INT(total/60.), &
            '  -- Completion in (min) : ', INT((total-elapsed)/60.)
    END IF
  END SUBROUTINE print_iteration

  SUBROUTINE copy_theta_to_q(f_theta_rhodz,f_rhodz,f_q)
    TYPE(t_field),POINTER :: f_theta_rhodz(:),f_rhodz(:), f_q(:)
    REAL(rstd), POINTER :: theta_rhodz(:,:,:), rhodz(:,:), q(:,:,:)
    INTEGER :: ind, iq
    DO ind=1, ndomain
       IF (.NOT. assigned_domain(ind)) CYCLE
       CALL swap_dimensions(ind)
       CALL swap_geometry(ind)
       theta_rhodz=f_theta_rhodz(ind)
       rhodz=f_rhodz(ind)
       q=f_q(ind)
       DO iq=1, nqdyn
          q(:,:,iq)  = theta_rhodz(:,:,iq)/rhodz(:,:)
       END DO
    END DO
  END SUBROUTINE copy_theta_to_q

  SUBROUTINE copy_q_to_theta(f_theta_rhodz,f_rhodz,f_q)
    TYPE(t_field),POINTER :: f_theta_rhodz(:),f_rhodz(:), f_q(:)
    REAL(rstd), POINTER :: theta_rhodz(:,:,:), rhodz(:,:), q(:,:,:)
    INTEGER :: ind, iq
    DO ind=1, ndomain
       IF (.NOT. assigned_domain(ind)) CYCLE
       CALL swap_dimensions(ind)
       CALL swap_geometry(ind)
       theta_rhodz=f_theta_rhodz(ind)
       rhodz=f_rhodz(ind)
       q=f_q(ind)
       DO iq=1,nqdyn
          theta_rhodz(:,:,iq) = rhodz(:,:)*q(:,:,iq)
       END DO
    END DO
  END SUBROUTINE copy_q_to_theta

END MODULE timeloop_gcm_mod