! $Id$
MODULE regr_pr_time_av_m

  USE print_control_mod, ONLY: lunout
  USE write_field_phy
  USE mod_phys_lmdz_transfert_para, ONLY: bcast
  IMPLICIT NONE

!-------------------------------------------------------------------------------
! Purpose: Regrid pressure with an averaging method. Operations done:
!  * on the root process: read a NetCDF 3D or 4D field at a given day.
!  * pack the fields to the LMDZ "horizontal "physics" grid and scatter
!    to all threads of all processes;
!  * in all the threads of all the processes, regrid the fields in pressure
!    to the LMDZ vertical grid.
!  * input files fields are stretched in the viscinity (+/- 5 kms) of the mean
!    tropopause (geometrical mean of LMDZ and input fields tropopauses) to force
!    the resulting field tropopause height to the one of LMDZ. To switch this
!    feature on, the following arguments must be present:
!     - "lat_in":   input file latitudes.
!     - "pcen_in":  input file cells center pressure levels.
!     - "ptrop_ou": target grid (LMDZ) tropopause pressure.
!    Note that the ozone quantity conservation is not ensured.
!-------------------------------------------------------------------------------
! Initial routine: regr_pr_av_m module (L. Guez).
! Present version: David Cugnet ; corresponding additions:
!    - time interpolation
!    - 3D compliant
!    - vertical stretching of the field to allow tropopause and ground pressure
!      matching between input field and current lmdz field.
!-------------------------------------------------------------------------------
! Remarks:
!  * 3D fields are zonal means, with dimensions (latitude, pressure, julian day)
!  * 4D fields have the dimensions:  (longitude, latitude, pressure, julian day)
!  * All the fields are already on the horizontal grid (rlatu) or (rlatu,rlonv),
!    except that the latitudes are in ascending order in the input file.
!  * We assume that all the input fields have the same coordinates.
!  * The target vertical LMDZ grid is the grid of layer boundaries.
!  * Regridding in pressure is conservative, second order.
!  * All the fields are regridded as a single multi-dimensional array, so it
!    saves CPU time to call this procedure once for several NetCDF variables
!    rather than several times, each time for a single NetCDF variable.
!  * The input file pressure at tropopause can be (in decreasing priority):
!    1) read from the file if "tropopause_air_pressure" field is available.
!    2) computed from the input file ozone field using:
!       - o3 concentration at tropopause if "tro3_at_tropopause" is available.
!       - a default value (100ppbv) if not.
!  * If available, the field "ps" (input file pressure at surface) is used.
!    If not, the current LMDZ ground pressure is taken instead.
!  * The O3 threshold for tropopause is defined using 3 coefficients:
!    o3t(ppbV)=co1+co2*SIN(PI*(month-2)/6)*TANH(lat_deg/c)
!     => co1 and co2 are in ppmV, and co3 in degrees.
!     => co3 allow a smooth transition between north and south hemispheres.
!-------------------------------------------------------------------------------
! * Fields with suffix "m"/"p" are at the closest records earlier/later than the
!   middle of the julian day "julien", on the global "dynamics" horizontal grid.
! * Fields(i,j,k,l) are at longitude-latitude "rlonv(i)-rlatu(j)", pressure
!   interval "pint_in(k:k+1)]" and variable "nam(l)".
! * In the 2D file case, the values are the same for all longitudes.
! * The following fields are on the global "dynamics" grid, as read from files:
  REAL,    SAVE, ALLOCATABLE :: v1m(:,:,:,:)       !--- Previous ozone fields
  REAL,    SAVE, ALLOCATABLE :: v1p(:,:,:,:)       !--- Next ozone fields
  REAL,    SAVE, ALLOCATABLE :: psm(:,:), psp(:,:) !--- Surface pressure
  REAL,    SAVE, ALLOCATABLE :: ptm(:,:), ptp(:,:) !--- Tropopause pressure
  REAL,    SAVE, ALLOCATABLE :: otm(:,:), otp(:,:) !--- Tropopause o3 mix. ratio
  INTEGER, SAVE :: irec                            !--- Current time index
!      * for daily   input files: current julian day number
!      * for monthly input files: julien is in [time_in(irec),time_in(irec+1)]
  LOGICAL, SAVE :: linterp                         !--- Interpolation in time
  LOGICAL, SAVE :: lPrSurf                         !--- Surface pressure flag
  LOGICAL, SAVE :: lPrTrop                         !--- Tropopause pressure flag
  LOGICAL, SAVE :: lO3Trop                         !--- Tropopause ozone flag
  LOGICAL, SAVE :: lfirst=.TRUE.                   !--- First call flag
!$OMP THREADPRIVATE(lfirst)
  REAL, PARAMETER :: co3(3)=[91.,28.,20.]          !--- Coeffs for o3 threshold
  REAL, PARAMETER :: prtrop=5.E+3  !--- Value lower than the tropopause pressure
  REAL, PARAMETER :: delta =5.     !--- Dist. around tropopause for strain (kms)

CONTAINS

!-------------------------------------------------------------------------------
!
SUBROUTINE regr_pr_time_av(fID, nam, julien, pint_in, pint_ou, v3, time_in,    &
                                     lat_in, pcen_in, ptrop_ou)
!
!-------------------------------------------------------------------------------
  USE dimphy,         ONLY: klon
  USE netcdf95,       ONLY: NF95_INQ_VARID, NF95_INQUIRE_VARIABLE, handle_err
  USE netcdf,         ONLY: NF90_INQ_VARID, NF90_GET_VAR, NF90_NOERR
  USE assert_m,       ONLY: assert
  USE assert_eq_m,    ONLY: assert_eq
  USE comvert_mod,    ONLY: scaleheight
  USE interpolation,  ONLY: locate
  USE regr_conserv_m, ONLY: regr_conserv
  USE regr_lint_m,    ONLY: regr_lint
  USE slopes_m,       ONLY: slopes
  USE mod_phys_lmdz_mpi_data,       ONLY: is_mpi_root
  USE mod_grid_phy_lmdz,            ONLY: nbp_lon, nbp_lat, nbp_lev
  USE mod_phys_lmdz_transfert_para, ONLY: scatter2d, scatter
  USE phys_cal_mod,                 ONLY: calend, year_len, days_elapsed, jH_cur
!-------------------------------------------------------------------------------
! Arguments:
  INTEGER, INTENT(IN)  :: fID           !--- NetCDF file ID
  CHARACTER(LEN=13), INTENT(IN) :: nam(:)     !--- NetCDF variables names
  REAL,    INTENT(IN)  :: julien        !--- Days since Jan 1st
  REAL,    INTENT(IN)  :: pint_in(:)    !--- Interfaces file Pres, Pa, ascending
  REAL,    INTENT(IN)  :: pint_ou(:,:)  !--- Interfaces LMDZ Pres, Pa (klon,llm+1)
  REAL,    INTENT(OUT) :: v3(:,:,:)     !--- Regridded field (klon,llm,SIZE(nam))
  REAL,    INTENT(IN), OPTIONAL :: time_in(:) !--- Records times, in days
                                              !    since Jan 1 of current year
  REAL,    INTENT(IN), OPTIONAL :: lat_in(:)  !--- Input/output latitudes vector
  REAL,    INTENT(IN), OPTIONAL :: pcen_in(:) !--- Mid-layers file pressure
  REAL,    INTENT(IN), OPTIONAL :: ptrop_ou(:)!--- Output tropopause pressure (klon)
!-------------------------------------------------------------------------------
! Local variables:
  include "clesphys.h"
  include "YOMCST.h"
  CHARACTER(LEN=80)  :: sub
  CHARACTER(LEN=320) :: msg
  INTEGER :: vID, ncerr, n_var, nlev_in,ntim_in, ndim, i, ibot, itop, itrp,itrp0
  LOGICAL :: lAdjTro                          !--- Need to adjust tropopause
  REAL    :: y_frac                           !--- Elapsed year fraction
  REAL    :: alpha, beta, al                  !--- For strectching/interpolation
  REAL    :: SigT_in, SigT_ou, SigT_mn        !--- Tropopause: in/out/mean
  REAL    :: SigA_bot, SigA_top               !--- Strained domain bounds
  REAL    :: Sig_in (SIZE(pint_in))           !--- Input field sigma levels
  REAL    :: phi    (SIZE(pint_in))           !--- Stretching exponent anomaly
  REAl    :: Pint_st(SIZE(pint_in))           !--- Stretched pressure levels
  REAL    :: v1(nbp_lon,nbp_lat,SIZE(pint_in)-1,SIZE(nam))
  REAL    :: v2(klon,           SIZE(pint_in)-1,SIZE(nam))
! v1: Field read/interpol at time "julien" on the global "dynamics" horiz. grid.
! v2: Field scattered to the partial "physics" horizontal grid.
!     In the 2D file case, the values are the same for all longitudes.
!     "v2(i,    k, l)" is at longitude-latitude  "xlon(i)-xlat(i)".
! Both are for pressure interval "press_in_edg(k:k+1)]" and variable "nam(l)"
  REAL, DIMENSION(nbp_lon, nbp_lat) ::  ps1, pt1, ot1
  REAL, DIMENSION(klon)             ::  ps2, pt2, ot2
!-------------------------------------------------------------------------------
  sub="regr_pr_time_av"
  nlev_in=SIZE(pint_in)-1; ntim_in=SIZE(time_in)
  CALL assert(SIZE(v3,1)==klon,                 TRIM(sub)//" v3 klon")
  CALL assert(SIZE(v3,2)==nbp_lev,              TRIM(sub)//" v3 nbp_lev")
  n_var = assert_eq(SIZE(nam), SIZE(v3,3),      TRIM(sub)//" v3 n_var")
  CALL assert(SHAPE(pint_ou)==[klon, nbp_lev+1],TRIM(sub)//" pint_ou")
  IF(PRESENT(lat_in))   CALL assert(SIZE(lat_in  )==klon,TRIM(sub)//" lat_in klon")
  IF(PRESENT(ptrop_ou)) CALL assert(SIZE(ptrop_ou)==klon,TRIM(sub)//" ptrop_ou klon")
  IF(PRESENT(pcen_in))  CALL assert(SIZE(pcen_in)==nlev_in,TRIM(sub)//" pcen_in")
  lAdjTro=PRESENT(ptrop_ou)
  IF(lAdjTro.AND.(.NOT.PRESENT(lat_in).OR..NOT.PRESENT(pcen_in))) &
  CALL abort_physic(sub, 'Missing lat_in and/or pcen_in (adjust_tropopause=T)', 1)

  !$OMP MASTER
  IF(is_mpi_root) THEN

    !=== CHECK WHICH FIELDS ARE AVAILABLE IN THE INPUT FILE
    IF(lfirst) THEN
      lPrSurf=NF90_INQ_VARID(fID,"ps"                     ,vID)==NF90_NOERR
      lPrTrop=NF90_INQ_VARID(fID,"tropopause_air_pressure",vID)==NF90_NOERR
      lO3Trop=NF90_INQ_VARID(fID,"tro3_at_tropopause"     ,vID)==NF90_NOERR
      linterp=PRESENT(time_in); IF(linterp) linterp=ntim_in==14
      IF(linterp) THEN
        ALLOCATE(v1m(nbp_lon,nbp_lat,nlev_in,n_var))
        ALLOCATE(v1p(nbp_lon,nbp_lat,nlev_in,n_var))
        ALLOCATE(psm(nbp_lon,nbp_lat),psp(nbp_lon,nbp_lat))
        ALLOCATE(ptm(nbp_lon,nbp_lat),ptp(nbp_lon,nbp_lat))
        IF(lO3Trop) ALLOCATE(otm(nbp_lon,nbp_lat),otp(nbp_lon,nbp_lat))
      END IF
      !--- INITIAL INDEX: LOCATE A LAYER WELL ABOVE TROPOPAUSE (50hPa)
      IF(lAdjTro) itrp0=locate(pcen_in,prtrop)
      IF(lPrSurf) WRITE(lunout,*)TRIM(sub)//': Using GROUND PRESSURE from input O3 forcing file.'
      IF(linterp) WRITE(lunout,*)TRIM(sub)//': Monthly O3 files => ONLINE TIME INTERPOLATION.'
      IF(lAdjTro) THEN
        WRITE(lunout,*)TRIM(sub)//': o3 forcing file tropopause location uses:'
        IF(lPrTrop) THEN
          WRITE(lunout,*)'    PRESSURE AT TROPOPAUSE from file.'
        ELSE IF(lO3Trop) THEN
          WRITE(lunout,*)'    O3 CONCENTRATION AT TROPOPAUSE from file.'
        ELSE
          WRITE(lunout,*)'    PARAMETRIZED O3 concentration at tropopause.'
        END IF
      END IF
    END IF

    !=== UPDATE (ALWAYS FOR DAILY FILES, EACH MONTH FOR MONTHLY FILES)
    CALL update_fields()


    !=== TIME INTERPOLATION FOR MONTHLY INPUT FILES
    IF(linterp) THEN
      WRITE(lunout,'(3(a,f7.3))')TRIM(sub)//': Interpolating O3 at julian day '&
      &,julien,' from fields at times ',time_in(irec),' and ', time_in(irec+1)
      al=(time_in(irec+1)-julien)/(time_in(irec+1)-time_in(irec))
      v1=al*v1m+(1.-al)*v1p
      IF(lPrSurf) ps1=al*psm+(1.-al)*psp
      IF(lPrTrop) pt1=al*ptm+(1.-al)*ptp
      IF(lO3Trop) ot1=al*otm+(1.-al)*otp
    END IF
  END IF
  !$OMP END MASTER
  IF(lfirst) THEN
    lfirst=.FALSE. ;     CALL bcast(lfirst) ; CALL bcast(lPrSurf)
    CALL bcast(lPrTrop); CALL bcast(lO3Trop); CALL bcast(linterp)
    IF(lAdjTro) CALL bcast(itrp0)
  END IF
  CALL scatter2d(v1,v2)
  !--- No "ps" in input file => assumed to be equal to current LMDZ ground press
  IF(lPrSurf) THEN; CALL scatter2d(ps1,ps2); ELSE; ps2=Pint_ou(:,1); END IF
  IF(lPrTrop) CALL scatter2d(pt1,pt2)
  IF(lO3Trop) CALL scatter2d(ot1,ot2)

  !--- REGRID IN PRESSURE ; 3rd index inverted because "paprs" is decreasing
  IF(.NOT.lAdjTro) THEN
    DO i=1,klon
      CALL regr_conserv(1,v2(i,:,:) , Pint_in , Pint_ou(i,nbp_lev+1:1:-1),    &
                          v3(i,nbp_lev:1:-1,:), slopes(1,v2(i,:,:),pint_in))
    END DO
  ELSE
    y_frac=(REAL(days_elapsed)+jH_cur)/year_len
    DO i=1,klon
      SigT_in = get_SigTrop(i,itrp)        !--- input  (file) tropopause
      SigT_ou = ptrop_ou(i)/pint_ou(i,1)   !--- output (lmdz) tropopause
      SigT_mn = SQRT(SigT_in*SigT_ou)      !--- mean tropopause>strained domain

      !--- DEFINE THE VERTICAL LAYER IN WHICH THE PROFILE IS STRECHED
      beta=EXP(delta/scaleheight); Sig_in(:) = [pint_in(1:nlev_in)/ps2(i),1.]
      SigA_bot = SigT_mn*beta ; ibot=locate(Sig_in(:),SigA_bot)
      SigA_top = SigT_mn/beta ; itop=locate(Sig_in(:),SigA_top)+1

      !--- HAT FUNCTION phi (/=0 in [SigA_bot,SigA_top] only)
      phi(:)=0.
      phi(itop:itrp)=(Sig_in(itop:itrp)-SigA_top)/(SigT_in-SigA_top)
      phi(itrp:ibot)=(SigA_bot-Sig_in(itrp:ibot))/(SigA_bot-SigT_in)

      !--- STRAINED INPUT logP PROFILE ; alpha: MAX. STRETCH. EXPONENT INCREMENT
      alpha = LOG(SigT_ou/SigT_in)/LOG(SigT_in)
      Pint_st(:) = pint_ou(i,1) * Sig_in(:)**(1+alpha*phi(:))

      !--- REGRID INPUT PROFILE ON STRAINED VERTICAL LEVELS
      CALL regr_conserv(1,v2(i,:,:) , Pint_st,  Pint_ou(i,nbp_lev+1:1:-1),     &
                          v3(i,nbp_lev:1:-1,:), slopes(1,v2(i,:,:),Pint_st))
    END DO
  END IF


CONTAINS


!-------------------------------------------------------------------------------
!
SUBROUTINE update_fields()
!
!-------------------------------------------------------------------------------
  IF(.NOT.linterp) THEN                 !=== DAILY FILES: NO TIME INTERPOLATION
    WRITE(lunout,*)TRIM(sub)//': Updating Ozone forcing field: read from file.'
    irec=MIN(INT(julien)+1,ntim_in)     !--- irec is just the julian day number
    !--- MIN -> Security in the unlikely case of roundup errors.
    CALL get_3Dfields(v1)               !--- Read ozone field(s)
    IF(lAdjTro) THEN                    !--- Additional files for fields strain
      IF(lPrSurf) CALL get_2Dfield(ps1,"ps")
      IF(lPrTrop) CALL get_2Dfield(pt1,"tropopause_air_pressure")
      IF(lO3Trop) CALL get_2Dfield(ot1,"tro3_at_tropopause")
    END IF
  ELSE                                  !=== MONTHLY FILES: GET 2 NEAREST RECS
    IF(.NOT.lfirst.AND.julien<time_in(irec+1)) RETURN
    WRITE(lunout,*)TRIM(sub)//': Refreshing adjacent Ozone forcing fields.'
    IF(lfirst) THEN                     !=== READ EARLIEST TIME FIELDS
      irec=locate(time_in,julien)       !--- Need to locate surrounding times
      CALL get_3Dfields(v1m)            !--- Read ozone field(s)
      IF(lAdjTro) THEN                  !--- Additional files for fields strain
        IF(lPrSurf) CALL get_2Dfield(psm,"ps")
        IF(lPrTrop) CALL get_2Dfield(ptm,"tropopause_air_pressure")
        IF(lO3Trop) CALL get_2Dfield(otm,"tro3_at_tropopause")
      END IF
    ELSE                                !=== SHIFT FIELDS
      v1m=v1p                           !--- Ozone fields
      IF(lAdjTro) THEN                  !--- Additional files for fields strain
        IF(lPrSurf) psm=psp             !--- Surface pressure
        IF(lPrTrop) ptm=ptp             !--- Tropopause pressure
        IF(lO3Trop) otm=otp             !--- Tropopause ozone
      END IF
    END IF
    irec=irec+1
    CALL get_3Dfields(v1p)              !--- Read ozone field(s)
    IF(lAdjTro) THEN                    !--- Additional files for fields strain
      IF(lPrSurf) CALL get_2Dfield(psp,"ps")
      IF(lPrTrop) CALL get_2Dfield(ptp,"tropopause_air_pressure")
      IF(lO3Trop) CALL get_2Dfield(otp,"tro3_at_tropopause")
    END IF
    IF(lfirst) irec=irec-1
  END IF

END SUBROUTINE update_fields
!
!-------------------------------------------------------------------------------


!-------------------------------------------------------------------------------
!
SUBROUTINE get_2Dfield(v,var)
!
!-------------------------------------------------------------------------------
! Purpose: Shortcut to get the "irec"th record of the surface field named "var"
!          from the input file.
! Remark: In case the field is zonal, it is duplicated along first dimension.
!-------------------------------------------------------------------------------
! Arguments:
  REAL,             INTENT(INOUT) :: v(:,:)
  CHARACTER(LEN=*), INTENT(IN)    :: var
!-------------------------------------------------------------------------------
  CALL NF95_INQ_VARID(fID, TRIM(var), vID)
  CALL NF95_INQUIRE_VARIABLE(fID, vID, ndims=ndim)
  IF(ndim==2) ncerr=NF90_GET_VAR(fID,vID,v(1,:), start=[  1,irec])
  IF(ndim==3) ncerr=NF90_GET_VAR(fID,vID,v(:,:), start=[1,1,irec])
  CALL handle_err(TRIM(sub)//" NF90_GET_VAR "//TRIM(var),ncerr,fID)

  !--- Flip latitudes: ascending in input file, descending in "rlatu".
  IF(ndim==3) THEN
    v(1,:) = v(1,nbp_lat:1:-1)
    v(2:,:)= SPREAD(v(1,:),DIM=1,ncopies=nbp_lon-1)  !--- Duplication
  ELSE
    v(:,:) = v(:,nbp_lat:1:-1)
  END IF

END SUBROUTINE get_2Dfield
!
!-------------------------------------------------------------------------------


!-------------------------------------------------------------------------------
!
SUBROUTINE get_3Dfields(v)
!
!-------------------------------------------------------------------------------
! Purpose: Shortcut to get the "irec"th record of the 3D fields named "nam"
!          from the input file. Fields are stacked on fourth dimension.
! Remark: In case the field is zonal, it is duplicated along first dimension.
!-------------------------------------------------------------------------------
! Arguments:
  REAL, INTENT(INOUT) :: v(:,:,:,:)
!-------------------------------------------------------------------------------
  DO i=1,SIZE(nam)
    CALL NF95_INQ_VARID(fID, TRIM(nam(i)), vID)
    CALL NF95_INQUIRE_VARIABLE(fID, vID, ndims=ndim)
    IF(ndim==3) ncerr=NF90_GET_VAR(fID,vID,v(1,:,:,i), start=[  1,1,irec])
    IF(ndim==4) ncerr=NF90_GET_VAR(fID,vID,v(:,:,:,i), start=[1,1,1,irec])
    CALL handle_err(TRIM(sub)//" NF90_GET_VAR "//TRIM(nam(i)),ncerr,fID)
  END DO

  !--- Flip latitudes: ascending in input file, descending in "rlatu".
  IF(ndim==3) THEN
    v(1,:,:,:) = v(1,nbp_lat:1:-1,:,:)
    v(2:,:,:,:)= SPREAD(v(1,:,:,:),DIM=1,ncopies=nbp_lon-1)  !--- Duplication
  ELSE
    v(:,:,:,:) = v(:,nbp_lat:1:-1,:,:)
  END IF

END SUBROUTINE get_3Dfields
!
!-------------------------------------------------------------------------------


!-------------------------------------------------------------------------------
!
FUNCTION get_SigTrop(ih,it)
!
!-------------------------------------------------------------------------------
! Arguments:
  INTEGER, INTENT(IN)  :: ih
  INTEGER, INTENT(OUT) :: it
  REAL                 :: get_Sigtrop
!-------------------------------------------------------------------------------
! Local variables:
  INTEGER :: ii                                !--- Idx of layer containing o3t
  REAL    :: o3t                               !--- Ozone concent. at tropopause
  REAL    :: prt                               !--- Air pressure   at tropopause
  REAL    :: al                                !--- Interpolation coefficient
  REAL    :: coef                              !--- Coef: North/South transition
!-------------------------------------------------------------------------------
  !--- DETERMINE PRESSURE AT TROPOPAUSE AND DIVIDE IT BY GROUND PRESSURE
  IF(lPrTrop) THEN                             !--- PrTrop KNOWN FROM FILE
    get_SigTrop=pt2(ih)/ps2(ih)
    it=locate(pint_in(:),pt2(ih))
  ELSE                                         !--- O3 THRESHOLD
    coef = TANH(lat_in(i)/co3(3))              !--- Latitude dependant ponderat.
    coef = SIN (2.*RPI*(y_frac-1./6.)) * coef  !--- Time-dependant ponderation
    o3t  = 1.E-9 * (co3(1)+co3(2)*coef)        !--- Value from parametrization
    IF(lO3Trop) o3t=ot2(ih)                    !--- Value from file
    !--- Starts from 50hPa and go down.
    it=itrp0; DO WHILE(v2(ih,it+1,1)>=o3t); it=it+1; END DO
    al=(v2(ih,it,1)-o3t)/(v2(ih,it,1)-v2(ih,it+1,1))
    get_SigTrop = ( pcen_in(it)**(1.-al) * pcen_in(it+1)**al )/ps2(ih)
  END IF

END FUNCTION get_SigTrop
!
!-------------------------------------------------------------------------------


END SUBROUTINE regr_pr_time_av
!
!-------------------------------------------------------------------------------

END MODULE regr_pr_time_av_m