MODULE time_field_mod
  USE prec,       ONLY : rstd
  USE field_mod,  ONLY : t_field
  IMPLICIT NONE
  PRIVATE
  LOGICAL,PARAMETER :: xios_remap=.TRUE. 
  SAVE

  TYPE t_time_interpolation
     ! immutable members
     INTEGER :: interval ! interval t1-t0, in time step units
     ! mutable members
     LOGICAL :: swap, &  ! if set, swap fields for t0 and t1 (start/end of interval)
     &          read_t0, & ! if set, read or compute data at t0
     &          read_t1    ! if set, read or compute data at t1
     INTEGER :: itau0, itau, itau1 ! time step corresponding to times t0, t, t1
     REAL(rstd) :: coef0, coef1
  END type t_time_interpolation
  
  TYPE t_time_field
     CHARACTER(30) :: name, pname, suffix ! XIOS names of field and of 3D pressure used for interpolation
     
     TYPE(t_field), POINTER :: f_t(:), & ! value of field at time t in [t0,t1]
          &                    f_t0(:),  & ! value at time t0
          &                    f_t1(:),   &  ! value at time t1
          &                    f_tmp(:),   &! temporary field for reading tracers  
          &                    f_levels(:)  ! temporary field on read level  
     LOGICAL :: is_ps ! surface pressure ps plays a special role when reading 3D input files
     LOGICAL :: is_q !
     INTEGER :: ndim
     REAL(rstd),ALLOCATABLE :: levels(:) ! vertical axis avlue of the read field
     INTEGER    :: nlevels ! nb levels of read field 
     LOGICAL    :: do_vertical_remap = .TRUE. ! do the vertical remapping (True) or let XIOS doing it (false)
  END type t_time_field

  TYPE(t_field), POINTER :: f_p(:)  ! 3D pressure field reconstructed from ps and hybrid coordinate
  TYPE(t_field), POINTER :: f_ps(:) ! ps

  PUBLIC :: t_time_interpolation, init_time_interpolation, update_time_interpolation, &
       &    t_time_field, init_time_field, allocate_time_field, &
       &    read_time_field, swap_time_field, interpolate_time_field

CONTAINS

  SUBROUTINE init_time_field
    USE icosa
    USE field_mod,  ONLY : field_T, type_real, allocate_field
    USE grid_param, ONLY : llm
    USE xios_mod
    USE omp_para
    
    CALL allocate_field(f_p, field_t, type_real, llm, name='pressure_time_field')
    CALL allocate_field(f_ps, field_t, type_real, name='pressure_surface_time_field')

  END SUBROUTINE init_time_field

  SUBROUTINE init_time_interpolation(interval, interp)
    INTEGER,                    INTENT(IN)  :: interval
    TYPE(t_time_interpolation), INTENT(OUT) :: interp
    interp%interval = interval
  END SUBROUTINE init_time_interpolation


  SUBROUTINE allocate_time_field(do_vertical_remap, field, field_type, data_type, dim3, dim4, name,suffix, pname)
    USE icosa
    USE field_mod,  ONLY : allocate_field
    USE domain_mod, ONLY : ndomain, assigned_domain
    USE xios_mod
    USE omp_para
    ! value of argument pname must reflect the contents of XIOS XML configuration files
    LOGICAL, INTENT(IN)                 :: do_vertical_remap
    TYPE(t_time_field),     INTENT(OUT) :: field
    INTEGER,                INTENT(IN)  :: field_type
    INTEGER,                INTENT(IN)  :: data_type
    INTEGER,      OPTIONAL, INTENT(IN)  :: dim3, dim4
    CHARACTER(*), OPTIONAL, INTENT(IN)  :: name, suffix, pname
    INTEGER :: ind

    !$OMP BARRIER
    IF (is_omp_master) THEN 
      field%do_vertical_remap = do_vertical_remap
      field%name  = TRIM(name) ! note that name is in fact a compulsory argument
      field%suffix =""
      IF (PRESENT(suffix)) field%suffix = TRIM(suffix) 
      field%is_ps = PRESENT(pname)
      IF(field%is_ps) field%pname = TRIM(pname)
      field%is_q = .FALSE.
      IF (PRESENT(dim4) .AND. dim4 == nqtot) field%is_q = .TRUE.
 
      CALL xios_get_axis_attr("axis_nudged",n_glo=field%nlevels)
      ALLOCATE(field%levels(field%nlevels))

      IF (is_omp_master) CALL xios_get_axis_attr("axis_nudged",value=field%levels)
!      CALL bcast_omp(field%levels)
    
      IF (PRESENT(dim4)) THEN 
        field%ndim=4
      ELSE IF (PRESENT(dim3)) THEN 
        field%ndim=3
      ELSE
        field%ndim=2
      ENDIF
    ENDIF
    !$OMP BARRIER
    
    IF (field%do_vertical_remap .AND. PRESENT(dim3)) THEN
      CALL allocate_field(field%f_t0, field_type, data_type, field%nlevels, dim4, name=TRIM(name)//'_t0')
      CALL allocate_field(field%f_t1, field_type, data_type, field%nlevels, dim4, name=TRIM(name)//'_t1')
      CALL allocate_field(field%f_tmp, field_type, data_type, field%nlevels, name=TRIM(name)//'_tmp')
      CALL allocate_field(field%f_levels,  field_type, data_type, field%nlevels, dim4, name=TRIM(name)//'_level')
      CALL allocate_field(field%f_t,  field_type, data_type, dim3, dim4, name=TRIM(name)//'_t')
   ELSE 
      CALL allocate_field(field%f_t0, field_type, data_type, dim3, dim4, name=TRIM(name)//'_t0')
      CALL allocate_field(field%f_t1, field_type, data_type, dim3, dim4, name=TRIM(name)//'_t1')
      CALL allocate_field(field%f_tmp, field_type, data_type, dim3, name=TRIM(name)//'_tmp')
      CALL allocate_field(field%f_levels,  field_type, data_type, field%nlevels, dim4, name=TRIM(name)//'_level')
      CALL allocate_field(field%f_t,  field_type, data_type, dim3, dim4, name=TRIM(name)//'_t')
    ENDIF
    
    DO ind=1, ndomain
       IF(.NOT. assigned_domain(ind)) CYCLE
       field%f_t(ind)%rval4d = 0. 
       field%f_t0(ind)%rval4d = 0. 
       field%f_t1(ind)%rval4d = 0. 
       field%f_tmp(ind)%rval4d = 0. 
       field%f_levels(ind)%rval4d = 0. 
    END DO
!$OMP BARRIER
  END SUBROUTINE allocate_time_field
  
  SUBROUTINE update_time_interpolation(time, interp)
    ! must be called once before calling read_time_field
    ! sets flags read_t0, swap, read_t1 that trigger the corresponding actions in read_time_field
    ! the same flags can be used to for data that is computed rather than read from XIOS
    USE time_mod,   ONLY : dt, itau0
    REAL(rstd),                 INTENT(IN)    :: time
    TYPE(t_time_interpolation), INTENT(INOUT) :: interp

    INTEGER :: itau
    REAL(rstd) :: t0,t1

    itau=time/dt
    interp%itau = itau
    
    IF (itau==itau0+1) THEN
       ! first iteration : get data at both ends of time interval, do not swap t0, t1
       interp%read_t0 = .TRUE.
       interp%read_t1 = .TRUE.
       interp%swap    = .FALSE.
       interp%itau0   = itau
       interp%itau1   = itau + interp%interval
    ELSE
       IF (MOD(itau,interp%interval)==1) THEN
          ! we have reached t1, swap data and read/compute new data for t1
          interp%read_t0 = .FALSE.
          interp%read_t1 = .TRUE.
          interp%swap    = .TRUE.
          interp%itau0   = interp%itau1
          interp%itau1   = interp%itau0 + interp%interval
       ELSE
          ! in general do nothing
          interp%read_t0 = .FALSE.
          interp%read_t1 = .FALSE.
          interp%swap    = .FALSE.
       END IF
    END IF

    ! if itau==itau0, coef0, coef1 = 1., 0.
    interp%coef0 = REAL( interp%itau1 - interp%itau ) / REAL( interp%itau1 - interp%itau0 )
    interp%coef1 = REAL( interp%itau - interp%itau0 ) / REAL( interp%itau1 - interp%itau0 )
  END SUBROUTINE update_time_interpolation

  SUBROUTINE swap_time_field(field)
    TYPE(t_time_field), INTENT(INOUT) :: field
    TYPE(t_field), POINTER :: junk(:)
    ! t1 becomes t0 and t0 memory is reused to read/compute the new t1
    !$OMP BARRIER
    !$OMP MASTER
    junk       => field%f_t0
    field%f_t0 => field%f_t1
    field%f_t1 => junk
    !$OMP END MASTER
    !$OMP BARRIER
  END SUBROUTINE swap_time_field
  
  SUBROUTINE read_time_field(interp, field)
    ! surface pressure ps must always be read before 3D fields
    ! in the special case that we read ps, we compute the 3D pressure field
    ! and provide it to XIOS for later use in vertical interpolation of 3D fields
    ! this special role of ps is marked by the flag is_ps
    USE layout_mod,    ONLY : field_T, field_U, field_Z
    USE transfert_mod, ONLY : transfert_request, req_i0, req_i1
    USE pression_mod,  ONLY : pression_mid
    USE xios_mod,      ONLY : xios_read_field, xios_write_field, xios_is_valid_field, xios_is_active_field
    USE tracer_icosa_mod, ONLY : tracers
    USE icosa
    USE write_field_mod
    USE omp_para
    USE vertical_remap_mod

    TYPE(t_time_interpolation), INTENT(IN)    :: interp
    TYPE(t_time_field),         INTENT(INOUT) :: field

    REAL(rstd) :: t0,t1
    INTEGER :: iq,l,ind
    LOGICAL :: test

    IF(interp%read_t0) THEN
       ! we get data for t0 only once, from XIOS field XXX0
       IF (field%is_q) THEN ! for tracers
         DO iq=1,nqtot
!           CALL xios_read_field(TRIM(tracers(iq)%name)//TRIM(field%suffix)//"0", field%f_tmp)
           IF (is_omp_master) THEN
             test=xios_is_valid_field(TRIM(tracers(iq)%name)//TRIM(field%suffix)//"0")
             IF (test) test=xios_is_active_field(TRIM(tracers(iq)%name)//TRIM(field%suffix)//"0")
            ENDIF
           CALL bcast_omp(test)
           IF (test) THEN
             CALL xios_read_field(TRIM(tracers(iq)%name)//TRIM(field%suffix)//"0", field%f_tmp)
           ELSE
             DO ind=1,ndomain
               IF (is_omp_first_level) field%f_tmp(ind)%rval4d(:,:,:)=0
             ENDDO
           ENDIF
           
           DO ind=1,ndomain
             IF(.NOT. assigned_domain(ind)) CYCLE
             CALL swap_dimensions(ind)
             CALL swap_geometry(ind) 
             IF (is_omp_first_level) field%f_t0(ind)%rval4d(:,:,iq) = field%f_tmp(ind)%rval4d(:,:,1)
           ENDDO
         ENDDO
       ELSE
         CALL xios_read_field(TRIM(field%name)//"0", field%f_t0)
       ENDIF
       !$OMP BARRIER
       
       SELECT CASE(field%f_t0(1)%field_type)
       CASE(field_T)
         IF (.NOT. field%do_vertical_remap) THEN
           CALL transfert_request(field%f_t0, req_i0)
           CALL transfert_request(field%f_t0, req_i1)
         ENDIF
       CASE DEFAULT
          STOP 'read_time_field : unsupported field_type'
       END SELECT


       IF (.NOT. field%do_vertical_remap) THEN
         IF(field%is_ps) THEN
           DO ind=1,ndomain
             IF(.NOT. assigned_domain(ind)) CYCLE
             IF (is_omp_first_level) f_ps(ind)%rval4d = field%f_t0(ind)%rval4d
           ENDDO
           !$OMP BARRIER
           CALL pression_mid(field%f_t0, f_p)
           CALL xios_write_field(TRIM(field%pname)//"0",f_p)
         ENDIF
      END IF
    END IF

    IF(interp%swap) CALL swap_time_field(field)

    IF(interp%read_t1) THEN
       IF (field%is_q) THEN ! for tracers
         DO iq=1,nqtot
           IF (is_omp_master) THEN
             test=xios_is_valid_field(TRIM(tracers(iq)%name)//TRIM(field%suffix))
             IF (test) test=xios_is_active_field(TRIM(tracers(iq)%name)//TRIM(field%suffix))
            ENDIF
           CALL bcast_omp(test)
           IF (test) THEN
             CALL xios_read_field(TRIM(tracers(iq)%name)//TRIM(field%suffix), field%f_tmp)
           ELSE
              DO ind=1,ndomain
                IF(.NOT. assigned_domain(ind)) CYCLE
                IF (is_omp_first_level) field%f_tmp(ind)%rval4d(:,:,:)=0
              ENDDO
              !$OMP BARRIER
           ENDIF
 
           DO ind=1,ndomain
             IF(.NOT. assigned_domain(ind)) CYCLE
             CALL swap_dimensions(ind)
             CALL swap_geometry(ind) 
             IF (is_omp_first_level) field%f_t1(ind)%rval4d(:,:,iq) = field%f_tmp(ind)%rval4d(:,:,1)
           ENDDO
           !$OMP BARRIER
         ENDDO
       ELSE
         CALL xios_read_field(TRIM(field%name), field%f_t1)
       ENDIF
       
       SELECT CASE(field%f_t1(1)%field_type)
       CASE(field_T)
         IF (.NOT. field%do_vertical_remap) THEN
           CALL transfert_request(field%f_t1, req_i0)
           CALL transfert_request(field%f_t1, req_i1)
         ENDIF
       CASE DEFAULT
          STOP 'read_time_field : unsupported field_type'
       END SELECT

       IF (.NOT. field%do_vertical_remap) THEN
         IF(field%is_ps) THEN
           CALL pression_mid(field%f_t1, f_p)
           CALL xios_write_field(field%pname, f_p)
         ENDIF
       END IF

    ENDIF

  END SUBROUTINE read_time_field

 
  SUBROUTINE interpolate_time_field(interp, field, f_ps)
    USE field_mod,  ONLY : t_field
    USE domain_mod, ONLY : ndomain, assigned_domain
    USE vertical_remap_mod
    IMPLICIT NONE
    TYPE(t_time_interpolation), INTENT(IN)      :: interp
    TYPE(t_time_field),         INTENT(INOUT)   :: field
    TYPE(t_field), POINTER,OPTIONAL,INTENT(IN)  :: f_ps(:)
    INTEGER :: ind, n

!$OMP BARRIER
     DO ind=1, ndomain
       IF (field%do_vertical_remap .AND. field%ndim>2) THEN
         n = SIZE(field%f_levels(ind)%rval4d)
         CALL compute_interpolate_time_field(interp%coef0, interp%coef1, n, &
                      field%f_t0(ind)%rval4d, field%f_t1(ind)%rval4d, field%f_levels(ind)%rval4d) 
       ELSE
         n = SIZE(field%f_t(ind)%rval4d)
         CALL compute_interpolate_time_field(interp%coef0, interp%coef1, n, &
                      field%f_t0(ind)%rval4d, field%f_t1(ind)%rval4d, field%f_t(ind)%rval4d) 
       ENDIF 
     END DO

     IF (field%do_vertical_remap) THEN
       IF (field%ndim==2) THEN
         ! no vertical interpolation ; do nothing
       ELSE 
         CALL vertical_remap(field%levels,field%f_levels, f_ps, field%f_t)
       ENDIF
     ENDIF
!     CALL writeField(TRIM(field%name)//"_time",  field%f_t)
    
  END SUBROUTINE interpolate_time_field

  SUBROUTINE compute_interpolate_time_field(coef0, coef1, n, f0, f1, f)
    INTEGER,    INTENT(IN)  :: n
    REAL(rstd), INTENT(IN)  :: coef0, coef1, f0(n), f1(n)
    REAL(rstd), INTENT(OUT) :: f(n)
    INTEGER :: ij
    !$OMP DO SCHEDULE(STATIC)
    DO ij=1,n
       f(ij) = coef0*f0(ij) + coef1*f1(ij)
    END DO
  END SUBROUTINE compute_interpolate_time_field
  
END MODULE time_field_mod
