MODULE timeloop_gcm_mod

  PRIVATE
  
  PUBLIC :: timeloop

  INTEGER, PARAMETER :: euler=1, rk4=2, mlf=3

CONTAINS
  
  SUBROUTINE timeloop
  USE icosa
  USE dissip_gcm_mod
  USE caldyn_mod
  USE etat0_mod
  USE guided_mod
  USE advect_tracer_mod
  USE physics_mod
  USE mpipara
  
  IMPLICIT NONE
  TYPE(t_field),POINTER :: f_phis(:)
!  TYPE(t_field),POINTER :: f_theta(:)
  TYPE(t_field),POINTER :: f_q(:)
  TYPE(t_field),POINTER :: f_dtheta(:), f_rhodz(:)
  TYPE(t_field),POINTER :: f_ps(:),f_psm1(:), f_psm2(:)
  TYPE(t_field),POINTER :: f_u(:),f_um1(:),f_um2(:)
  TYPE(t_field),POINTER :: f_theta_rhodz(:),f_theta_rhodzm1(:),f_theta_rhodzm2(:)
  TYPE(t_field),POINTER :: f_dps(:),f_dpsm1(:), f_dpsm2(:)
  TYPE(t_field),POINTER :: f_du(:),f_dum1(:),f_dum2(:)
  TYPE(t_field),POINTER :: f_dtheta_rhodz(:),f_dtheta_rhodzm1(:),f_dtheta_rhodzm2(:)
  TYPE(t_field),POINTER :: f_hflux(:), f_wflux(:), f_hfluxt(:), f_wfluxt(:)

  REAL(rstd),POINTER :: phis(:)
  REAL(rstd),POINTER :: q(:,:,:)
  REAL(rstd),POINTER :: ps(:) ,psm1(:), psm2(:)
  REAL(rstd),POINTER :: u(:,:) , um1(:,:), um2(:,:)
  REAL(rstd),POINTER :: rhodz(:,:), theta_rhodz(:,:) , theta_rhodzm1(:,:), theta_rhodzm2(:,:)
  REAL(rstd),POINTER :: dps(:), dpsm1(:), dpsm2(:)
  REAL(rstd),POINTER :: du(:,:), dum1(:,:), dum2(:,:)
  REAL(rstd),POINTER :: dtheta_rhodz(:,:),dtheta_rhodzm1(:,:),dtheta_rhodzm2(:,:)
  REAL(rstd),POINTER :: hflux(:,:),wflux(:,:),hfluxt(:,:),wfluxt(:,:)

!  REAL(rstd) :: dt, run_length
  INTEGER :: ind
  INTEGER :: it,i,j,n, nb_stage, stage, matsuno_period, scheme
  CHARACTER(len=255) :: scheme_name
  LOGICAL :: fluxt_zero(ndomain) ! set to .TRUE. to start accumulating fluxes in time
!  INTEGER :: itaumax
!  REAL(rstd) ::write_period
!  INTEGER    :: itau_out
  
!  dt=90.
!  CALL getin('dt',dt)
  
!  itaumax=100
!  CALL getin('itaumax',itaumax)

!  run_length=dt*itaumax
!  CALL getin('run_length',run_length)
!  itaumax=run_length/dt
!  PRINT *,'itaumax=',itaumax
!  dt=dt/scale_factor

! Trends
  CALL allocate_field(f_dps,field_t,type_real)
  CALL allocate_field(f_du,field_u,type_real,llm)
  CALL allocate_field(f_dtheta_rhodz,field_t,type_real,llm)
! Model state at current time step (RK/MLF/Euler)
  CALL allocate_field(f_phis,field_t,type_real)
  CALL allocate_field(f_ps,field_t,type_real)
  CALL allocate_field(f_u,field_u,type_real,llm)
  CALL allocate_field(f_theta_rhodz,field_t,type_real,llm)
! Model state at previous time step (RK/MLF)
  CALL allocate_field(f_psm1,field_t,type_real)
  CALL allocate_field(f_um1,field_u,type_real,llm)
  CALL allocate_field(f_theta_rhodzm1,field_t,type_real,llm)
! Tracers
  CALL allocate_field(f_q,field_t,type_real,llm,nqtot)
  CALL allocate_field(f_rhodz,field_t,type_real,llm)
! Mass fluxes
  CALL allocate_field(f_hflux,field_u,type_real,llm)    ! instantaneous mass fluxes
  CALL allocate_field(f_wflux,field_t,type_real,llm+1)  ! computed by caldyn
  CALL allocate_field(f_hfluxt,field_u,type_real,llm)   ! mass "fluxes" accumulated in time
  CALL allocate_field(f_wfluxt,field_t,type_real,llm+1)

  scheme_name='runge_kutta'
  CALL getin('scheme',scheme_name)

  SELECT CASE (TRIM(scheme_name))
  CASE('euler')
     scheme=euler
     nb_stage=1
  CASE ('runge_kutta')
     scheme=rk4
     nb_stage=4
  CASE ('leapfrog_matsuno')
     scheme=mlf
     matsuno_period=5
     CALL getin('matsuno_period',matsuno_period)
     nb_stage=matsuno_period+1
     ! Model state 2 time steps ago (MLF)
     CALL allocate_field(f_psm2,field_t,type_real)
     CALL allocate_field(f_theta_rhodzm2,field_t,type_real,llm)
     CALL allocate_field(f_um2,field_u,type_real,llm)
  CASE default
     PRINT*,'Bad selector for variable scheme : <', TRIM(scheme_name),             &
          ' > options are <euler>, <runge_kutta>, <leapfrog_matsuno>'
     STOP
  END SELECT
  
!  write_period=0
!  CALL getin('write_period',write_period)
!  write_period=write_period/scale_factor
!  itau_out=FLOOR(.5+write_period/dt)
!  PRINT *, 'Output frequency (scaled) set to ',write_period, ' : itau_out = ',itau_out
 
! Trends at previous time steps needed only by Adams-Bashforth
!  CALL allocate_field(f_dpsm1,field_t,type_real)
!  CALL allocate_field(f_dpsm2,field_t,type_real)
!  CALL allocate_field(f_dum1,field_u,type_real,llm)
!  CALL allocate_field(f_dum2,field_u,type_real,llm)
!  CALL allocate_field(f_dtheta_rhodzm1,field_t,type_real,llm)
!  CALL allocate_field(f_dtheta_rhodzm2,field_t,type_real,llm)

!  CALL allocate_field(f_theta,field_t,type_real,llm)
!  CALL allocate_field(f_dtheta,field_t,type_real,llm)

  CALL init_dissip
  CALL init_caldyn
  CALL init_guided
  CALL init_advect_tracer
  CALL init_physics
  
  CALL etat0(f_ps,f_phis,f_theta_rhodz,f_u, f_q)
  CALL writefield("phis",f_phis,once=.TRUE.)
  CALL transfert_request(f_q,req_i1) 

  DO ind=1,ndomain
     CALL swap_dimensions(ind)
     CALL swap_geometry(ind)
     rhodz=f_rhodz(ind); ps=f_ps(ind)
     CALL compute_rhodz(.TRUE., ps,rhodz) ! save rhodz for transport scheme before dynamics update ps
  END DO
  fluxt_zero=.TRUE.

  ! check that rhodz is consistent with ps
  CALL transfert_request(f_rhodz,req_i1)
  CALL transfert_request(f_ps,req_i1)
  DO ind=1,ndomain
     CALL swap_dimensions(ind)
     CALL swap_geometry(ind)
     rhodz=f_rhodz(ind); ps=f_ps(ind)
     CALL compute_rhodz(.FALSE., ps, rhodz)   
  END DO
  
  DO it=0,itaumax

    IF (is_mpi_root) PRINT *,"It No :",It,"   t :",dt*It
    IF (mod(it,itau_out)==0 ) THEN
      CALL writefield("q",f_q)
      CALL update_time_counter(dt*it)
    ENDIF
    
    CALL guided(it*dt,f_ps,f_theta_rhodz,f_u,f_q)

    DO stage=1,nb_stage
       CALL caldyn((stage==1) .AND. (MOD(it,itau_out)==0), &
            f_phis,f_ps,f_theta_rhodz,f_u, f_q, &
            f_hflux, f_wflux, f_dps, f_dtheta_rhodz, f_du)
       SELECT CASE (scheme)
       CASE(euler)
          CALL euler_scheme(.TRUE.)
       CASE (rk4)
          CALL rk_scheme(stage)
       CASE (mlf)
          CALL  leapfrog_matsuno_scheme(stage)
          
          !      CASE ('leapfrog')
          !        CALL leapfrog_scheme
          !
          !      CASE ('adam_bashforth')
          !        CALL dissip(f_u,f_du,f_ps,f_phis, f_theta_rhodz,f_dtheta_rhodz)
          !        CALL adam_bashforth_scheme
       CASE DEFAULT
          STOP
       END SELECT
    END DO

    CALL dissip(f_u,f_du,f_ps,f_phis, f_theta_rhodz,f_dtheta_rhodz)
    CALL euler_scheme(.FALSE.)

    IF(MOD(it+1,itau_adv)==0) THEN
       CALL transfert_request(f_wfluxt,req_i1) ! FIXME
!       CALL transfert_request(f_hfluxt,req_e1) ! FIXME

       CALL advect_tracer(f_hfluxt,f_wfluxt,f_u, f_q,f_rhodz)  ! update q and rhodz after RK step
       fluxt_zero=.TRUE.

       ! FIXME : check that rhodz is consistent with ps
       CALL transfert_request(f_rhodz,req_i1)
       CALL transfert_request(f_ps,req_i1)
       CALL transfert_request(f_dps,req_i1) ! FIXME
       CALL transfert_request(f_wflux,req_i1) ! FIXME
       DO ind=1,ndomain
          CALL swap_dimensions(ind)
          CALL swap_geometry(ind)
          rhodz=f_rhodz(ind); ps=f_ps(ind); dps=f_dps(ind); 
          wflux=f_wflux(ind); wfluxt=f_wfluxt(ind)
          CALL compute_rhodz(.FALSE., ps, rhodz)   
       END DO

    END IF

!    CALL physics(it,f_phis, f_ps, f_theta_rhodz, f_u, f_q)
    ENDDO
 
  CONTAINS

    SUBROUTINE Euler_scheme(with_dps)
    IMPLICIT NONE
    LOGICAL :: with_dps
    INTEGER :: ind
    DO ind=1,ndomain
       CALL swap_dimensions(ind)
       CALL swap_geometry(ind)
       IF(with_dps) THEN
          ps=f_ps(ind) ; dps=f_dps(ind) ; 
          ps(:)=ps(:)+dt*dps(:)
          hflux=f_hflux(ind);     hfluxt=f_hfluxt(ind)
          wflux=f_wflux(ind);     wfluxt=f_wfluxt(ind)
          CALL accumulate_fluxes(hflux,wflux,hfluxt,wfluxt,dt,fluxt_zero(ind))
       END IF
       u=f_u(ind) ; theta_rhodz=f_theta_rhodz(ind)
       du=f_du(ind) ; dtheta_rhodz=f_dtheta_rhodz(ind)
       u(:,:)=u(:,:)+dt*du(:,:)
       theta_rhodz(:,:)=theta_rhodz(:,:)+dt*dtheta_rhodz(:,:)
    ENDDO

    END SUBROUTINE Euler_scheme

    SUBROUTINE RK_scheme(stage)
      IMPLICIT NONE
      INTEGER :: ind, stage
      REAL(rstd), DIMENSION(4), PARAMETER :: coef = (/ .25, 1./3., .5, 1. /)
      REAL(rstd) :: tau

      tau = dt*coef(stage)
      
      DO ind=1,ndomain
         CALL swap_dimensions(ind)
         CALL swap_geometry(ind)
         ps=f_ps(ind)   ; u=f_u(ind)   ; theta_rhodz=f_theta_rhodz(ind)
         psm1=f_psm1(ind) ; um1=f_um1(ind) ; theta_rhodzm1=f_theta_rhodzm1(ind)
         dps=f_dps(ind) ; du=f_du(ind) ; dtheta_rhodz=f_dtheta_rhodz(ind)
         
         IF (stage==1) THEN ! first stage : save model state in XXm1
            psm1(:)=ps(:) ; um1(:,:)=u(:,:) ; theta_rhodzm1(:,:)=theta_rhodz(:,:)
         END IF
         ! updates are of the form x1 := x0 + tau*f(x1)
         ps(:)=psm1(:)+tau*dps(:)
         u(:,:)=um1(:,:)+tau*du(:,:)
         theta_rhodz(:,:)=theta_rhodzm1(:,:)+tau*dtheta_rhodz(:,:)
         IF(stage==nb_stage) THEN ! accumulate mass fluxes at last stage
            hflux=f_hflux(ind);     hfluxt=f_hfluxt(ind)
            wflux=f_wflux(ind);     wfluxt=f_wfluxt(ind)
            CALL accumulate_fluxes(hflux,wflux, hfluxt,wfluxt, dt,fluxt_zero(ind))
         END IF
      END DO
    END SUBROUTINE RK_scheme

    SUBROUTINE leapfrog_scheme
    IMPLICIT NONE
      INTEGER :: ind

      DO ind=1,ndomain
        CALL swap_dimensions(ind)
        CALL swap_geometry(ind)
        ps=f_ps(ind)   ; u=f_u(ind)   ; theta_rhodz=f_theta_rhodz(ind)
        psm1=f_psm1(ind) ; um1=f_um1(ind) ; theta_rhodzm1=f_theta_rhodzm1(ind)
        psm2=f_psm2(ind) ; um2=f_um2(ind) ; theta_rhodzm2=f_theta_rhodzm2(ind)
        dps=f_dps(ind) ; du=f_du(ind) ; dtheta_rhodz=f_dtheta_rhodz(ind)

        IF (it==0) THEN
          psm2(:)=ps(:) ; theta_rhodzm2(:,:)=theta_rhodz(:,:) ; um2(:,:)=u(:,:) 

          ps(:)=ps(:)+dt*dps(:)
          u(:,:)=u(:,:)+dt*du(:,:)
          theta_rhodz(:,:)=theta_rhodz(:,:)+dt*dtheta_rhodz(:,:)

          psm1(:)=ps(:) ; theta_rhodzm1(:,:)=theta_rhodz(:,:) ; um1(:,:)=u(:,:) 
        ELSE
        
          ps(:)=psm2(:)+2*dt*dps(:)
          u(:,:)=um2(:,:)+2*dt*du(:,:)
          theta_rhodz(:,:)=theta_rhodzm2(:,:)+2*dt*dtheta_rhodz(:,:)

          psm2(:)=psm1(:) ; theta_rhodzm2(:,:)=theta_rhodzm1(:,:) ; um2(:,:)=um1(:,:) 
          psm1(:)=ps(:) ; theta_rhodzm1(:,:)=theta_rhodz(:,:) ; um1(:,:)=u(:,:) 
        ENDIF
          
      ENDDO
    END SUBROUTINE leapfrog_scheme  
 
    SUBROUTINE leapfrog_matsuno_scheme(stage)
    IMPLICIT NONE
    INTEGER :: ind, stage
    REAL :: tau
    tau = dt/nb_stage
      DO ind=1,ndomain
        CALL swap_dimensions(ind)
        CALL swap_geometry(ind)

        ps=f_ps(ind)   ; u=f_u(ind)   ; theta_rhodz=f_theta_rhodz(ind)
        psm1=f_psm1(ind) ; um1=f_um1(ind) ; theta_rhodzm1=f_theta_rhodzm1(ind)
        psm2=f_psm2(ind) ; um2=f_um2(ind) ; theta_rhodzm2=f_theta_rhodzm2(ind)
        dps=f_dps(ind) ; du=f_du(ind) ; dtheta_rhodz=f_dtheta_rhodz(ind)

        
!        IF (MOD(it,matsuno_period+1)==0) THEN
        IF (stage==1) THEN
          psm1(:)=ps(:) ; um1(:,:)=u(:,:) ; theta_rhodzm1(:,:)=theta_rhodz(:,:)
          ps(:)=psm1(:)+tau*dps(:)
          u(:,:)=um1(:,:)+tau*du(:,:)
          theta_rhodz(:,:)=theta_rhodzm1(:,:)+tau*dtheta_rhodz(:,:)

!        ELSE IF (MOD(it,matsuno_period+1)==1) THEN
        ELSE IF (stage==2) THEN

          ps(:)=psm1(:)+tau*dps(:)
          u(:,:)=um1(:,:)+tau*du(:,:)
          theta_rhodz(:,:)=theta_rhodzm1(:,:)+tau*dtheta_rhodz(:,:)

          psm2(:)=psm1(:) ; theta_rhodzm2(:,:)=theta_rhodzm1(:,:) ; um2(:,:)=um1(:,:) 
          psm1(:)=ps(:) ; theta_rhodzm1(:,:)=theta_rhodz(:,:) ; um1(:,:)=u(:,:) 

        ELSE 

          ps(:)=psm2(:)+2*tau*dps(:)
          u(:,:)=um2(:,:)+2*tau*du(:,:)
          theta_rhodz(:,:)=theta_rhodzm2(:,:)+2*tau*dtheta_rhodz(:,:)

          psm2(:)=psm1(:) ; theta_rhodzm2(:,:)=theta_rhodzm1(:,:) ; um2(:,:)=um1(:,:) 
          psm1(:)=ps(:) ; theta_rhodzm1(:,:)=theta_rhodz(:,:) ; um1(:,:)=u(:,:) 

        ENDIF
      
      ENDDO
      
    END SUBROUTINE leapfrog_matsuno_scheme  
         
    SUBROUTINE adam_bashforth_scheme
    IMPLICIT NONE
      INTEGER :: ind

      DO ind=1,ndomain
        CALL swap_dimensions(ind)
        CALL swap_geometry(ind)
        ps=f_ps(ind)   ; u=f_u(ind)   ; theta_rhodz=f_theta_rhodz(ind)
        dps=f_dps(ind) ; du=f_du(ind) ; dtheta_rhodz=f_dtheta_rhodz(ind)
        dpsm1=f_dpsm1(ind) ; dum1=f_dum1(ind) ; dtheta_rhodzm1=f_dtheta_rhodzm1(ind)
        dpsm2=f_dpsm2(ind) ; dum2=f_dum2(ind) ; dtheta_rhodzm2=f_dtheta_rhodzm2(ind)

        IF (it==0) THEN
          dpsm1(:)=dps(:) ; dum1(:,:)=du(:,:) ; dtheta_rhodzm1(:,:)=dtheta_rhodz(:,:)
          dpsm2(:)=dpsm1(:) ; dum2(:,:)=dum1(:,:) ; dtheta_rhodzm2(:,:)=dtheta_rhodzm1(:,:)
        ENDIF
              
        ps(:)=ps(:)+dt*(23*dps(:)-16*dpsm1(:)+5*dpsm2(:))/12
        u(:,:)=u(:,:)+dt*(23*du(:,:)-16*dum1(:,:)+5*dum2(:,:))/12
        theta_rhodz(:,:)=theta_rhodz(:,:)+dt*(23*dtheta_rhodz(:,:)-16*dtheta_rhodzm1(:,:)+5*dtheta_rhodzm2(:,:))/12

        dpsm2(:)=dpsm1(:) ; dum2(:,:)=dum1(:,:) ; dtheta_rhodzm2(:,:)=dtheta_rhodzm1(:,:)
        dpsm1(:)=dps(:) ; dum1(:,:)=du(:,:) ; dtheta_rhodzm1(:,:)=dtheta_rhodz(:,:)

      ENDDO     
      
    END SUBROUTINE adam_bashforth_scheme

  END SUBROUTINE timeloop    

  SUBROUTINE compute_rhodz(comp, ps, rhodz)
    USE icosa
    USE disvert_mod
    LOGICAL, INTENT(IN) :: comp ! .TRUE. to compute, .FALSE. to check
    REAL(rstd), INTENT(IN) :: ps(iim*jjm)
    REAL(rstd), INTENT(OUT) :: rhodz(iim*jjm,llm)
    REAL(rstd) :: m, err
    INTEGER :: l,i,j,ij,dd
    err=0.
    IF(comp) THEN 
       dd=1
    ELSE
!       dd=-1
       dd=0
    END IF

    DO l = 1, llm
       DO j=jj_begin-dd,jj_end+dd
          DO i=ii_begin-dd,ii_end+dd
             ij=(j-1)*iim+i
             m = ( ap(l) - ap(l+1) + (bp(l)-bp(l+1))*ps(ij) )/g 
             IF(comp) THEN
                rhodz(ij,l) = m
             ELSE
                err = MAX(err,abs(m-rhodz(ij,l)))
             END IF
          ENDDO
       ENDDO
    ENDDO

    IF(.NOT. comp) THEN
       IF(err>1e-10) THEN
          PRINT *, 'Discrepancy between ps and rhodz detected', err
          STOP
       ELSE
!          PRINT *, 'No discrepancy between ps and rhodz detected'
       END IF
    END IF

  END SUBROUTINE compute_rhodz

  SUBROUTINE accumulate_fluxes(hflux,wflux, hfluxt,wfluxt, tau,fluxt_zero)
    USE icosa
    REAL(rstd), INTENT(IN)    :: hflux(3*iim*jjm,llm), wflux(iim*jjm,llm+1)
    REAL(rstd), INTENT(INOUT) :: hfluxt(3*iim*jjm,llm), wfluxt(iim*jjm,llm+1)
    REAL(rstd), INTENT(IN) :: tau
    LOGICAL, INTENT(INOUT) :: fluxt_zero
    IF(fluxt_zero) THEN
!       PRINT *, 'Accumulating fluxes (first)'
       fluxt_zero=.FALSE.
       hfluxt = tau*hflux
       wfluxt = tau*wflux
    ELSE
!       PRINT *, 'Accumulating fluxes (next)'
       hfluxt = hfluxt + tau*hflux
       wfluxt = wfluxt + tau*wflux
    END IF
  END SUBROUTINE accumulate_fluxes
  
END MODULE timeloop_gcm_mod