MODULE TRAJECTORY_MOD

!     Purpose.
!     --------
!       Manage trajectory for TL and AD runs in a unified way.

!     Author.
!     -------
!        Y. Tremolet *ECMWF*

!     Modifications.
!     --------------
!        Original : 2000-08-04
!        Modified : 02-05-06 C. Fischer - ltrajsave=false when deallocate
!        M.Hamrud      01-Oct-2003 CY28 Cleaning
!        Modified : 03-12-19 C. Temperton - case LSLAG=.F. but LSPRT=.T.
!        Y.Tremolet    18-Mar-2004 Add lprttraj option for diagnostics
!        L.Isaksen: 2004-11-11 Upperair grid point trajectory option added
!        G. Desroziers and K. Yessad (sept 2005):
!         - split option LTRAJHR into LTRAJHR_ALTI and LTRAJHR_SURF.
!         - adapt option LTRAJHR to METEO-FRANCE configurations.
!        M.Hamrud      01-Jul-2006  Revised surface fields
!     ------------------------------------------------------------------

USE PARKIND1  ,ONLY : JPIM     ,JPRB
USE YOMHOOK   ,ONLY : LHOOK,   DR_HOOK

USE YOMCT0, ONLY: LIFSTRAJ, LIFSMIN, LSLAG, LSPRT, NSTART, NSTOP, &
 & LARPEGEF_TRAJHR,LARPEGEF_TRAJBG
USE YOMCT0B,ONLY: LECMWF
USE YOPHNC, ONLY: LETRAJP
USE YOMNMIA,ONLY: NITNMI
USE YOMDYN, ONLY: TSTEP
USE YOMVAR, ONLY: MUPTRA, NUPTRA
USE YOMLUN, ONLY: NULOUT
USE YOM_YGFL , ONLY : YGFL,YGFLC
USE SURFACE_FIELDS_MIX, ONLY : NSTRAJGRIB
USE YOMGRB, ONLY:  NGRBLNSP  
USE grib_header_mix
USE IOSTREAM_MIX , ONLY : SETUP_IOSTREAM, SETUP_IOREQUEST,  &
 & CLOSE_IOSTREAM, TYPE_IOSTREAM , TYPE_IOREQUEST, IO_INQUIRE, &
 & CLOSE_IOREQUEST

USE YOMTRAJ, ONLY: LTRAJSAVE, LTRAJSLAG, LTRAJPHYS, LTRAJCST,&
 & LTRAJGP, LTRAJHR, LREADGPTRAJ, TSTEP_TRAJ, NSMAX_TRAJ,&
 & NSMAX_BACKGR00, NSMAX_BACKGR01, NSMAX_BACKGR02, NSMAX_BACKGR03, &
 & LTRAJHR_ALTI,LTRAJHR_SURF, &
 & NGP5, NTRAJ_CST, LTRAJALLOC, MSTART,&
 & NSTEPTRAJ, MSTEPTRAJW, MSTEPTRAJR, LPRTTRAJ, LTRAJRESET, &
 & mtype_surf_traj, mtraj_grib, mtype_main3_traj, mtype_main2_traj, &
 & main_grib, mbackgr_grib
USE TRAJ_MAIN_MOD
USE TRAJ_SURFACE_MOD
USE TRAJ_PHYSICS_MOD
USE TRAJ_SEMILAG_MOD

IMPLICIT NONE

PRIVATE
PUBLIC LTRAJSAVE, LTRAJSLAG, LTRAJPHYS, LTRAJCST, &
 & LTRAJGP, LTRAJHR, LREADGPTRAJ, TSTEP_TRAJ, NSMAX_TRAJ, &
 & LTRAJHR_ALTI,LTRAJHR_SURF, &
 & NSMAX_BACKGR00, NSMAX_BACKGR01, NSMAX_BACKGR02, NSMAX_BACKGR03, &
 & ALLOCATE_TRAJECTORY, DEALLOCATE_TRAJECTORY, &
 & READ_TRAJECTORY, &
 & STORE_TRAJ_MAIN, GET_TRAJ_SPEC, GET_TRAJ_GRID, &
 & STORE_TRAJ_SFC, STORE_TRAJ_CST, GET_TRAJ_SFC, FORCE_SFC_TRAJ,&
 & STORE_TRAJ_PHYS, GET_TRAJ_PHYS, &
 & STORE_TRAJ_SLAG, GET_TRAJ_SLAG, &
 & NGP5, NTRAJ_CST, LPRTTRAJ, LTRAJRESET

#include "abor1.intfb.h"

!-----------------------------------------------------------------------
CONTAINS
!-----------------------------------------------------------------------

SUBROUTINE ALLOCATE_TRAJECTORY

REAL(KIND=JPRB) :: ZHOOK_HANDLE
IF (LHOOK) CALL DR_HOOK('TRAJECTORY_MOD:ALLOCATE_TRAJECTORY',0,ZHOOK_HANDLE)
IF (LTRAJALLOC) THEN
  IF (LECMWF) THEN
    ! ky: not sure that a test on LECMWF is allowed here.
    !     This ABOR1 should not be called for METEO-FRANCE configurations.
    CALL ABOR1('Trajectory already allocated')
  ELSE
    WRITE(NULOUT,*) ' Warning in routine TRAJECTORY_MOD:ALLOCATE_TRAJECTORY: '
    WRITE(NULOUT,*) '  Trajectory already allocated, execution continuing'
  ENDIF
ELSE
! Initialize controls for the trajectory
! --------------------------------------
  CALL INIT_TRAJECTORY

  LTRAJSAVE=.true.

  IF (LIFSTRAJ) THEN
    LTRAJCST =.false.
    LTRAJSLAG=.false.
    LTRAJPHYS=.false.
  ELSE
    LTRAJCST =.true.
    LTRAJSLAG=LSLAG.OR.LSPRT
    LTRAJPHYS=LETRAJP
  ENDIF

! Allocate trajectory
! -------------------
  CALL ALLOCATE_TRAJ_MAIN
  CALL ALLOCATE_TRAJ_SFC
  IF (LTRAJSLAG) CALL ALLOCATE_TRAJ_SLAG
  IF (LTRAJPHYS) CALL ALLOCATE_TRAJ_PHYS

  LTRAJALLOC=.true.
  IF (LPRTTRAJ) WRITE(NULOUT,*)'GREPTRAJ Trajectory storage allocated'
ENDIF

IF (LHOOK) CALL DR_HOOK('TRAJECTORY_MOD:ALLOCATE_TRAJECTORY',1,ZHOOK_HANDLE)
END SUBROUTINE ALLOCATE_TRAJECTORY

!-----------------------------------------------------------------------

SUBROUTINE DEALLOCATE_TRAJECTORY

REAL(KIND=JPRB) :: ZHOOK_HANDLE
IF (LHOOK) CALL DR_HOOK('TRAJECTORY_MOD:DEALLOCATE_TRAJECTORY',0,ZHOOK_HANDLE)
CALL DEALLOCATE_TRAJ_MAIN
CALL DEALLOCATE_TRAJ_SFC
CALL DEALLOCATE_TRAJ_PHYS
CALL DEALLOCATE_TRAJ_SLAG
IF (ALLOCATED(MSTEPTRAJW)) DEALLOCATE(MSTEPTRAJW)
IF (ALLOCATED(MSTEPTRAJR)) DEALLOCATE(MSTEPTRAJR)
IF (ALLOCATED(MTYPE_SURF_TRAJ)) DEALLOCATE(MTYPE_SURF_TRAJ)
IF (ALLOCATED(MTYPE_MAIN3_TRAJ)) DEALLOCATE(MTYPE_MAIN3_TRAJ)
IF (ALLOCATED(MTYPE_MAIN2_TRAJ)) DEALLOCATE(MTYPE_MAIN2_TRAJ)
LTRAJALLOC=.FALSE.

LTRAJSAVE = .FALSE.
IF (LPRTTRAJ) WRITE(NULOUT,*)'GREPTRAJ Trajectory storage deallocated'

IF (LHOOK) CALL DR_HOOK('TRAJECTORY_MOD:DEALLOCATE_TRAJECTORY',1,ZHOOK_HANDLE)
END SUBROUTINE DEALLOCATE_TRAJECTORY

!-----------------------------------------------------------------------

SUBROUTINE READ_TRAJECTORY(kstep,LD_LASTRAJ,LDREADGPTRAJ,LDBACKGR)

INTEGER(KIND=JPIM), INTENT(IN) :: kstep
LOGICAL, INTENT(IN) :: LD_LASTRAJ,LDREADGPTRAJ
LOGICAL, OPTIONAL,INTENT(IN) :: LDBACKGR

INTEGER(KIND=JPIM) :: istep
LOGICAL :: LLBACKGR
REAL(KIND=JPRB) :: ZHOOK_HANDLE

IF (LHOOK) CALL DR_HOOK('TRAJECTORY_MOD:READ_TRAJECTORY',0,ZHOOK_HANDLE)

LLBACKGR = .FALSE.
IF(PRESENT(LDBACKGR)) LLBACKGR=LDBACKGR
istep=MSTEPTRAJW(kstep)

IF (istep>0) THEN
  ! LDREADGPTRAJ=.T.:  Read both upperair q (and O3) plus surface grid point fields 
  ! LDREADGPTRAJ=.F.:  Only read surface grid point fields 

  IF (.NOT.LLBACKGR.AND.LTRAJHR_SURF) THEN
    CALL read_traj_sfc(istep,LD_LASTRAJ)
  ENDIF

! Read main trajectory either in grid point space or spectral space
! Read the background state for the initial time step (kstep=1),
! except for the first minimization

  IF (LTRAJHR_ALTI) THEN
    IF (LDREADGPTRAJ) THEN
      CALL read_traj_grid(istep,LD_LASTRAJ,LDBACKGR=LLBACKGR)
    ENDIF
    CALL read_traj_spec(istep,LD_LASTRAJ,LDBACKGR=LLBACKGR)
  ENDIF
ENDIF
IF (LPRTTRAJ) WRITE(NULOUT,*)'GREPTRAJ READ_TRAJECTORY: istep=', &
 & istep,LD_LASTRAJ,LLBACKGR

IF (LHOOK) CALL DR_HOOK('TRAJECTORY_MOD:READ_TRAJECTORY',1,ZHOOK_HANDLE)
END SUBROUTINE READ_TRAJECTORY

!-----------------------------------------------------------------------

SUBROUTINE INIT_TRAJECTORY

INTEGER(KIND=JPIM) :: itsteplr,itstephr,itime,i,jupdate,ifirst,ilast
INTEGER(KIND=JPIM) :: JGFL, IFLDS3D
CHARACTER(LEN=9) :: CL_CFILE
CHARACTER(LEN=11) :: clupd_file
CHARACTER(LEN=2) :: clupdate
TYPE(TYPE_IOSTREAM) :: YL_IOSTREAM
TYPE(TYPE_IOREQUEST) :: YL_IOREQUEST
REAL(KIND=JPRB) :: ZHOOK_HANDLE

! Initialize controls for the trajectory
! --------------------------------------

IF (LHOOK) CALL DR_HOOK('TRAJECTORY_MOD:INIT_TRAJECTORY',0,ZHOOK_HANDLE)
ALLOCATE(MSTEPTRAJW(nstart-nitnmi:nstop))
ALLOCATE(MSTEPTRAJR(nstart-nitnmi:nstop))

itstephr=nint(tstep)
itsteplr=nint(tstep_traj)
nsteptraj=0
DO i=nstart-nitnmi,nstart-1
  nsteptraj=nsteptraj+1
  MSTEPTRAJW(i)=nsteptraj
  MSTEPTRAJR(i)=nsteptraj
ENDDO
DO i=nstart,nstop
  itime=(i-nstart)*itstephr
  IF (mod(itime,MAX(itsteplr,1))==0) THEN
    nsteptraj=nsteptraj+1
    MSTEPTRAJW(i)=nsteptraj
  ELSE
    MSTEPTRAJW(i)=-1
  ENDIF
  MSTEPTRAJR(i)=nsteptraj
ENDDO
MSTART=MSTEPTRAJW(nstart)

write(NULOUT,*)'GREPTRAJ: MSTART=',MSTART

! Initialize GRIB parameters
! --------------------------
IF (LTRAJHR) THEN

  IF (LTRAJHR_SURF.AND.(.NOT.LARPEGEF_TRAJHR)) THEN


    !   Headers for grid point fields
    !   ----------------------------
    !   It would be better to use INQGRIB... YT
    CL_CFILE='lowres_gg'
    CALL SETUP_IOSTREAM(YL_IOSTREAM,'CIO',TRIM(CL_CFILE),CDMODE='r',KIOPROCS=(/-1/))
    CALL SETUP_IOREQUEST(YL_IOREQUEST,'GRIDPOINT_FIELDS',LDGRIB=.TRUE.,&
     & LDINTONLY=.TRUE.,KCHUNKSIZE=1,LDISTRIBUTE=.FALSE.)
    CALL IO_INQUIRE(YL_IOSTREAM,YL_IOREQUEST,KGRIB_HANDLE=mtraj_grib)
    CALL CLOSE_IOREQUEST(YL_IOREQUEST)
    CALL CLOSE_IOSTREAM(YL_IOSTREAM)
!    CALL init_grib_header(CL_CFILE,.true.,traj_grib)

    ALLOCATE(mtype_surf_traj(NGP5))
    mtype_surf_traj(:) = NSTRAJGRIB(1:NGP5)
  ENDIF

! BEN - LTRAJHR_ALTI=.true. by default when LECMWF=.true.
  IF (LTRAJHR_ALTI) THEN

    IF (.NOT.LARPEGEF_TRAJHR) THEN
      CL_CFILE='main_gg'
      CALL SETUP_IOSTREAM(YL_IOSTREAM,'CIO',TRIM(CL_CFILE),CDMODE='r',KIOPROCS=(/-1/))
      CALL SETUP_IOREQUEST(YL_IOREQUEST,'GRIDPOINT_FIELDS',LDGRIB=.TRUE.,&
       & LDINTONLY=.TRUE.,KCHUNKSIZE=1,LDISTRIBUTE=.FALSE.)
      CALL IO_INQUIRE(YL_IOSTREAM,YL_IOREQUEST,KGRIB_HANDLE=main_grib)
      CALL CLOSE_IOREQUEST(YL_IOREQUEST)
      CALL CLOSE_IOSTREAM(YL_IOSTREAM)
!      CALL init_grib_header(CL_CFILE,.true.,main_grib)
    ENDIF

    !==================================================================
    ! ky: has this part to go under LTRAJHR_ALTI?
    IF (((LIFSTRAJ .AND. NUPTRA==0) .OR. LIFSMIN).AND.(.NOT.LARPEGEF_TRAJBG)) THEN
      ! The background state is required at all resolutions used from 
      ! nuptra=0 to nuptra=muptra-1. 
      ! This is controlled by the file links in the script ifstraj.

      IF (LIFSTRAJ) THEN
        ifirst=0
        ilast=MUPTRA-1
      ELSE
        ifirst=nuptra
        ilast=nuptra
      ENDIF
      ALLOCATE(mbackgr_grib(ifirst:ilast))
      do jupdate=ifirst,ilast
        write(clupdate,'(I2.2)')jupdate
        clupd_file='backgr_gg'//clupdate
        CALL SETUP_IOSTREAM(YL_IOSTREAM,'CIO',TRIM(clupd_file),CDMODE='r',KIOPROCS=(/-1/))
        CALL SETUP_IOREQUEST(YL_IOREQUEST,'GRIDPOINT_FIELDS',LDGRIB=.TRUE.,&
         & LDINTONLY=.TRUE.,KCHUNKSIZE=1,LDISTRIBUTE=.FALSE.)
        CALL IO_INQUIRE(YL_IOSTREAM,YL_IOREQUEST,KGRIB_HANDLE=mbackgr_grib(jupdate))
        CALL CLOSE_IOREQUEST(YL_IOREQUEST)
        CALL CLOSE_IOSTREAM(YL_IOSTREAM)
!        CALL init_grib_header(clupd_file,.true.,backgr_grib(jupdate))
      enddo
    ENDIF
    !==================================================================

    IF (.NOT.LARPEGEF_TRAJHR) THEN

      ALLOCATE(mtype_main2_traj(1)) 
      mtype_main2_traj(1)=NGRBLNSP  

      IFLDS3D=0
      DO JGFL=1,YGFL%NUMFLDS
        IF ( YGFLC(JGFL)%LTRAJIO ) THEN
          IFLDS3D=IFLDS3D+1
        ENDIF
      ENDDO

      WRITE(NULOUT,*) 'TRAJECTORY: NUMBER OF GRIB CODES', IFLDS3D
      ALLOCATE(mtype_main3_traj(IFLDS3D))

      IFLDS3D=0
      DO JGFL=1,YGFL%NUMFLDS
        IF ( YGFLC(JGFL)%LTRAJIO ) THEN
          IFLDS3D= IFLDS3D+1
          mtype_main3_traj(IFLDS3D) = YGFLC(JGFL)%IGRBCODE
          WRITE(NULOUT,*) 'TRAJECTORY: GRIB CODES', mtype_main3_traj(IFLDS3D)
        ENDIF
      ENDDO

    ENDIF

  ENDIF 
ENDIF
IF (LPRTTRAJ) WRITE(NULOUT,*)'GREPTRAJ End INIT_TRAJECTORY'

IF (LHOOK) CALL DR_HOOK('TRAJECTORY_MOD:INIT_TRAJECTORY',1,ZHOOK_HANDLE)
END SUBROUTINE INIT_TRAJECTORY

!-----------------------------------------------------------------------

END MODULE TRAJECTORY_MOD
