MODULE disvert_mod
  USE icosa
  USE abort_mod
  REAL(rstd), SAVE, POINTER :: ap(:)
!$OMP THREADPRIVATE(ap) 
  REAL(rstd), SAVE, POINTER :: bp(:)
!$OMP THREADPRIVATE(bp) 
  REAL(rstd), SAVE, POINTER :: presnivs(:)
!$OMP THREADPRIVATE(presnivs) 
  REAL(rstd), SAVE, POINTER :: mass_al(:), mass_bl(:), mass_ak(:), mass_bk(:), mass_dak(:), mass_dbk(:)
!$OMP THREADPRIVATE(mass_al, mass_bl, mass_ak, mass_bk, mass_dak, mass_dbk) 

  REAL(rstd) :: ptop ! pressure at top of atmosphere l=llm+1
!$OMP THREADPRIVATE(ptop) 
  ! vertical coordinate : Lagrangian or mass-based
  INTEGER, SAVE :: caldyn_eta
!$OMP THREADPRIVATE(caldyn_eta) 
  INTEGER, PARAMETER :: eta_mass=1, eta_lag=2
  LOGICAL,SAVE :: ap_bp_present
!$OMP THREADPRIVATE(ap_bp_present) 

CONTAINS

  SUBROUTINE init_disvert
  USE disvert_std_mod, ONLY: ap_std=>ap, bp_std=>bp, presnivs_std=>presnivs, init_disvert_std=>init_disvert
  USE disvert_strato_mod, ONLY: ap_strato=>ap, bp_strato=>bp, presnivs_strato=>presnivs, init_disvert_strato=>init_disvert, init_disvert_strato_custom
  USE disvert_apbp_mod, ONLY: ap_apbp=>ap, bp_apbp=>bp, presnivs_apbp=>presnivs, init_disvert_apbp=>init_disvert
  USE disvert_ncar_mod, ONLY: ap_ncar=>ap, bp_ncar=>bp, presnivs_ncar=>presnivs, init_disvert_ncar=>init_disvert
  USE disvert_ncarl30_mod, ONLY: ap_ncarl30=>ap, bp_ncarl30=>bp, presnivs_ncarl30=>presnivs, init_disvert_ncarl30=>init_disvert
  USE disvert_dcmip31_mod, ONLY: ap_dcmip31=>ap, bp_dcmip31=>bp, presnivs_dcmip31=>presnivs, init_disvert_dcmip31=>init_disvert
  USE disvert_dcmip200_mod, ONLY: ap_dcmip200=>ap, bp_dcmip200=>bp, presnivs_dcmip200=>presnivs, init_disvert_dcmip200=>init_disvert
  USE icosa
  USE mpipara
  IMPLICIT NONE
    CHARACTER(len=255) :: def
    INTEGER :: l

    def='eta_mass'
    CALL getin('caldyn_eta',def)
    SELECT CASE(TRIM(def))
    CASE('eta_mass')
       caldyn_eta=eta_mass
    CASE('eta_lag')
       caldyn_eta=eta_lag
    CASE DEFAULT
       IF (is_mpi_root) PRINT *,'Bad selector for variable caldyn_eta : <', TRIM(def),'> options are <eta_mass>, <eta_lag>'
       STOP
    END SELECT
    IF (is_mpi_root) PRINT *, 'caldyn_eta=',TRIM(def), caldyn_eta, eta_lag, eta_mass

    def='std'
    ap_bp_present=.TRUE.
    CALL getin("disvert",def)

    SELECT CASE (TRIM(def))
      CASE('none')
         IF(caldyn_eta/=eta_lag) THEN
            IF(is_mpi_root) PRINT *,'disvert_type=none must be used with caldyn_eta=eta_lag'
            STOP
         END IF
         ap_bp_present=.FALSE.
      CASE('std')
    
        CALL init_disvert_std
        ap=>ap_std
        bp=>bp_std
        presnivs=>presnivs_std
      
      CASE('strato')
        
        CALL init_disvert_strato
        ap=>ap_strato
        bp=>bp_strato
        presnivs=>presnivs_strato

      CASE('strato_custom')
        
        CALL init_disvert_strato_custom
        ap=>ap_strato
        bp=>bp_strato
        presnivs=>presnivs_strato

      CASE('read_apbp')
    
        CALL init_disvert_apbp
        ap=>ap_apbp
        bp=>bp_apbp
        presnivs=>presnivs_apbp
        
      CASE ('ncar')

        CALL init_disvert_ncar
        ap=>ap_ncar
        bp=>bp_ncar
        presnivs=>presnivs_ncar

      CASE ('dcmip31')

        CALL init_disvert_dcmip31
        ap=>ap_dcmip31
        bp=>bp_dcmip31
        presnivs=>presnivs_dcmip31

      CASE ('dcmip200')

        CALL init_disvert_dcmip200
        ap=>ap_dcmip200
        bp=>bp_dcmip200
        presnivs=>presnivs_dcmip200

      CASE ('ncarl30')

        CALL init_disvert_ncarl30
        ap=>ap_ncarl30
        bp=>bp_ncarl30
        presnivs=>presnivs_ncarl30
      CASE default
        IF (is_mpi_root) PRINT*,'Bad selector for variable disvert : <', TRIM(def),"> options are <std>, <ncar>, <ncarl30>" 
        STOP
        
    END SELECT

    IF(ap_bp_present) THEN
       ! Convert from pressure coordinate to mass coordinate
       ! pk = ap(k) + bp(k)*ps(ij) = ptop + g*( mass_ak(k) + mass_bk(k)*ms(ij) )
       ptop = ap(llm+1)
       ALLOCATE(mass_al(llm+1))
       ALLOCATE(mass_bl(llm+1))
       ALLOCATE(mass_ak(llm))
       ALLOCATE(mass_bk(llm))
       ALLOCATE(mass_dak(llm))
       ALLOCATE(mass_dbk(llm))

       ! FIXME : leave ps for the moment ; change ps to ms later
       DO l=1,llm+1
          !       mass_al(l) = (ap(l)-ptop)/g
          !       mass_bl(l) = bp(l)/g
          mass_al(l) = (ap(l)-ptop)
          mass_bl(l) = bp(l)
       END DO
       DO l=1,llm
          mass_ak(l) = .5*(mass_al(l)+mass_al(l+1))
          mass_bk(l) = .5*(mass_bl(l)+mass_bl(l+1))
          mass_dak(l) = mass_al(l)-mass_al(l+1)  ! >0
          mass_dbk(l) = mass_bl(l)-mass_bl(l+1)  ! >0
       END DO
    END IF

    !$acc enter data copyin(ap(:), bp(:), mass_dak(:), mass_dbk(:)) async
    
  END SUBROUTINE init_disvert  
  
  SUBROUTINE compute_rhodz(comp, ps, rhodz, ondevice)
    USE icosa
    USE omp_para
    LOGICAL, INTENT(IN) :: comp ! .TRUE. to compute, .FALSE. to check
    REAL(rstd), INTENT(IN) :: ps(iim*jjm)
    REAL(rstd), INTENT(INOUT) :: rhodz(iim*jjm,llm)
    LOGICAL, INTENT(IN), OPTIONAL :: ondevice ! .TRUE. compute on device, .FALSE. stay on host
    REAL(rstd) :: m, err
    INTEGER :: l,i,j,ij,dd
    LOGICAL :: ondevice_

    err=0.

    ! by default, do not compute on device
    ondevice_ = .FALSE.
    IF (PRESENT(ondevice)) ondevice_ = ondevice

    IF(comp) THEN
       dd=1
    ELSE
       dd=0
    END IF

    !$acc data present(ps(:), rhodz(:,:), ap(:), bp(:)) async if (ondevice_)

    IF(comp) THEN
      !$acc parallel loop collapse(3) async if (ondevice_)
      DO l = ll_begin, ll_end
        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
            rhodz(ij,l) = m
          ENDDO
        ENDDO
      ENDDO
    ELSE
      !$acc parallel loop reduction(max:err) collapse(3) async if (ondevice_)
      DO l = ll_begin, ll_end
        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
            err = MAX(err,abs(m-rhodz(ij,l)))
          ENDDO
        ENDDO
      ENDDO

      !$acc wait
      IF(err>1e-10) THEN
          PRINT *, 'Discrepancy between ps and rhodz detected', err
          STOP
      ENDIF
    ENDIF

    !$acc end data
  END SUBROUTINE compute_rhodz


  SUBROUTINE write_apbp
  USE icosa
  USE netcdf_mod
  IMPLICIT NONE
    REAL(rstd) :: val(llm)
    INTEGER :: status
    INTEGER :: lev,ilev
    INTEGER :: ncid,levid,ilevid,hyaiid,hybiid,hyamid,hybmid,P0id
    INTEGER :: l
    
    IF (no_io) RETURN
    
!$OMP BARRIER
!$OMP MASTER    
    status = NF90_CREATE('apbp.nc', NF90_CLOBBER, ncid)
    status = NF90_DEF_DIM(ncid,'lev',llm,lev)
    status = NF90_DEF_DIM(ncid,'ilev',llm+1,ilev)
    
    status = NF90_DEF_VAR(ncid,'lev',NF90_DOUBLE,(/ lev /),levid)
    status = NF90_PUT_ATT(ncid,levid,"long_name","hybrid level at midpoints (1000*(A+B))")
    status = NF90_PUT_ATT(ncid,levid,"units","Pa")
    status = NF90_PUT_ATT(ncid,levid,"positive","up")
    status = NF90_PUT_ATT(ncid,levid,"standard_name","atmosphere_hybrid_sigma_pressure_coordinate")
    status = NF90_PUT_ATT(ncid,levid,"formula_terms","a: hyam b: hybm p0: P0 ps: PS")

    status = NF90_DEF_VAR(ncid,'ilev',NF90_DOUBLE,(/ ilev /),ilevid)
    status = NF90_PUT_ATT(ncid,ilevid,"long_name","hybrid level at interfaces (1000*(A+B))")
    status = NF90_PUT_ATT(ncid,ilevid,"units","Pa")
    status = NF90_PUT_ATT(ncid,ilevid,"positive","up")
    status = NF90_PUT_ATT(ncid,ilevid,"standard_name","atmosphere_hybrid_sigma_pressure_coordinate")
    status = NF90_PUT_ATT(ncid,ilevid,"formula_terms","a: hyai b: hybi p0: P0 ps: PS")
   
    status = NF90_DEF_VAR(ncid,'hyai',NF90_DOUBLE,(/ ilev /),hyaiid)
    status = NF90_PUT_ATT(ncid,hyaiid,"long_name","hybrid A coefficient at layer interfaces")

    status = NF90_DEF_VAR(ncid,'hybi',NF90_DOUBLE,(/ ilev /),hybiid)
    status = NF90_PUT_ATT(ncid,hybiid,"long_name","hybrid B coefficient at layer interfaces")

    status = NF90_DEF_VAR(ncid,'hyam',NF90_DOUBLE,(/ lev /),hyamid)
    status = NF90_PUT_ATT(ncid,hyamid,"long_name","hybrid A coefficient at midpoint interfaces")

    status = NF90_DEF_VAR(ncid,'hybm',NF90_DOUBLE,(/ lev /),hybmid)
    status = NF90_PUT_ATT(ncid,hybmid,"long_name","hybrid B coefficient at midpoint interfaces")
    
    status = NF90_DEF_VAR(ncid,'P0',NF90_DOUBLE,varid=P0id)

    status = NF90_ENDDEF(ncid)   

    IF(ap_bp_present) THEN
    
    status=NF90_PUT_VAR(ncid,ilevid, ap(:)+bp(:)*Preff)
    
    DO l=1,llm
      val(l)= 0.5*(ap(l+1)+ap(l)+Preff*(bp(l)+bp(l+1)))
    ENDDO
    
    status=NF90_PUT_VAR(ncid,levid, val)

    status=NF90_PUT_VAR(ncid,hyaiid, ap(:)/Preff)
    status=NF90_PUT_VAR(ncid,hybiid, bp(:))
    
    DO l=1,llm
      val(l)= 0.5*(ap(l+1)+ap(l))
    ENDDO
    status=NF90_PUT_VAR(ncid,hyamid, val(:)/Preff)
   
    DO l=1,llm
      val(l)= 0.5*(bp(l+1)+bp(l))
    ENDDO
    status=NF90_PUT_VAR(ncid,hybmid, val(:))
 
    status=NF90_PUT_VAR(ncid,P0id, Preff)
   
    END IF

    status=NF90_CLOSE(ncid)
!$OMP END MASTER
!$OMP BARRIER
  END SUBROUTINE write_apbp

END MODULE disvert_mod