MODULE physics_mod
  USE icosa
  USE field_mod
  USE physics_interface_mod
  USE omp_para
  IMPLICIT NONE
  PRIVATE

  INTEGER, PARAMETER :: phys_none=0, phys_column=1, &
       phys_HS94=3, phys_LB2012=4, &
       phys_DCMIP=11, phys_DCMIP2016=12, phys_plugin=13, &
       phys_lmdz_generic=21, phys_external=22
  INTEGER :: phys_type
  TYPE(t_field),POINTER,SAVE :: f_dulon(:), f_dulat(:)
  TYPE(t_field),POINTER,SAVE :: f_ulon(:), f_ulat(:)
  TYPE(t_field),POINTER,SAVE :: f_p(:), f_pk(:)
  TYPE(t_field),POINTER,SAVE :: f_temp(:)
  TYPE(t_field),POINTER,SAVE :: f_du_phys(:)

  CHARACTER(LEN=255),SAVE :: physics_type
!$OMP THREADPRIVATE(physics_type)

  INTERFACE
     SUBROUTINE plugin_init_physics
     END SUBROUTINE plugin_init_physics

     SUBROUTINE plugin_physics
     END SUBROUTINE plugin_physics
  END INTERFACE

  PROCEDURE(plugin_init_physics), POINTER :: init_physics_plugin => NULL()
  PROCEDURE(plugin_physics), POINTER      :: physics_plugin => NULL()

  PUBLIC :: physics, init_physics, init_physics_plugin, physics_plugin, zero_du_phys

CONTAINS

  SUBROUTINE init_physics
    USE mpipara
    USE abort_mod
    USE etat0_mod
    USE etat0_heldsz_mod, ONLY : init_phys_held_suarez => init_physics
    USE etat0_venus_mod, ONLY : init_phys_venus=>init_physics
    USE physics_dcmip_mod, ONLY : init_physics_dcmip=>init_physics
    USE physics_dcmip2016_mod, ONLY : init_physics_dcmip2016=>init_physics
    USE physics_lmdz_generic_mod, ONLY : init_physics_lmdz_generic=>init_physics
    USE physics_external_mod, ONLY : init_physics_external=>init_physics
    LOGICAL :: done
    physics_inout%dt_phys = dt*itau_physics
!$OMP PARALLEL
    CALL allocate_field(f_du_phys,field_u,type_real,llm, name='du_phys')

    physics_type='none'
    CALL getin("physics",physics_type)
    ! below, flag done is set to .FALSE. if the CALL to init_XXX must be done outside any OMP PARALLEL region
    done=.TRUE.
    phys_type=phys_column
    SELECT CASE(TRIM(physics_type))
    CASE ('none')
       IF(is_mpi_root) PRINT*,"NO PHYSICAL PACKAGE USED"
       phys_type = phys_none
    CASE ('held_suarez')
       phys_type = phys_HS94
       CALL init_phys_held_suarez
    CASE ('phys_lmdz_generic')
       phys_type=phys_lmdz_generic
       done = .FALSE.
    CASE ('phys_external')
       phys_type=phys_external
       done = .FALSE.
    END SELECT

    IF(phys_type == phys_column) THEN
       CALL allocate_field(f_dulon,field_t,type_real,llm, name='dulon')
       CALL allocate_field(f_dulat,field_t,type_real,llm, name='dulat')
       CALL allocate_field(f_temp,field_t,type_real,llm, name='temp')
       CALL allocate_field(f_ulon,field_t,type_real,llm, name='ulon')
       CALL allocate_field(f_ulat,field_t,type_real,llm, name='ulat')
       CALL allocate_field(f_p,field_t,type_real,llm+1, name='p')
       CALL allocate_field(f_pk,field_t,type_real,llm, name='pk')
       CALL init_pack_before ! Compute physics_inout%ngrid and offsets used by pack/unpack
       CALL init_pack_after  ! Defines Ai, lon, lat in physics_inout
       
       SELECT CASE(TRIM(physics_type))
       CASE ('dcmip')
          phys_type = phys_DCMIP
          CALL init_physics_dcmip
       CASE ('dcmip2016')
          phys_type = phys_DCMIP2016
          CALL init_physics_dcmip2016
       CASE ('Lebonnois2012')
          phys_type = phys_LB2012
          CALL init_phys_venus       
       CASE ('plugin')
          IF(.NOT.ASSOCIATED(init_physics_plugin)) THEN
             PRINT *, 'FATAL : physics = plugin selected by init_physics_plugin not set by driver program. Aborting.'
             STOP
          END IF
          phys_type = phys_plugin
          CALL init_physics_plugin
       CASE DEFAULT
          IF(is_mpi_root) PRINT*, 'init_physics : Bad selector for variable physics <',&
               TRIM(physics_type), '> options are <none>, <held_suarez>, <Lebonnois2012>,', &
               '<dcmip>, <dcmip2016>, <plugin>, <phys_lmdz_generic>, <phys_external>'
          STOP
       END SELECT
       
    END IF
!$OMP END PARALLEL

    IF(.NOT. done) THEN
       SELECT CASE(phys_type)
       CASE(phys_external) 
          CALL init_physics_external
       CASE(phys_lmdz_generic)
          CALL init_physics_lmdz_generic
       END SELECT
    END IF

    IF(is_mpi_root) PRINT *, 'phys_type = ',phys_type

  END SUBROUTINE init_physics

  SUBROUTINE zero_du_phys()
    REAL(rstd), DIMENSION(:,:), POINTER :: du
    INTEGER :: ind
    DO ind=1,ndomain
       IF (.NOT. assigned_domain(ind)) CYCLE
       CALL swap_dimensions(ind)
       CALL swap_geometry(ind)
       du=f_du_phys(ind)
       du(:,ll_begin:ll_end) = 0.
    END DO
  END SUBROUTINE zero_du_phys

  SUBROUTINE add_du_phys(coef, f_u)
    REAL(rstd), INTENT(IN) :: coef  ! -1 before physics, +1 after physics
    TYPE(t_field),POINTER :: f_u(:) ! velocity field before/after call to physics
    REAL(rstd), DIMENSION(:,:), POINTER :: u, du
    INTEGER :: ind
    DO ind=1,ndomain
       IF (.NOT. assigned_domain(ind)) CYCLE
       CALL swap_dimensions(ind)
       CALL swap_geometry(ind)
       du=f_du_phys(ind)
       u=f_u(ind)
       du(:,ll_begin:ll_end) = du(:,ll_begin:ll_end) + coef*u(:,ll_begin:ll_end)
    END DO
  END SUBROUTINE add_du_phys

  SUBROUTINE physics(it,f_phis, f_geopot, f_ps, f_theta_rhodz, f_ue, f_wflux, f_q)
    USE physics_lmdz_generic_mod, ONLY : physics_lmdz_generic => physics
    USE physics_external_mod, ONLY : physics_external => physics
    USE physics_dcmip_mod, ONLY : write_physics_dcmip => write_physics
    USE physics_dcmip2016_mod, ONLY : write_physics_dcmip2016 => write_physics
    USE etat0_heldsz_mod
    INTEGER, INTENT(IN)   :: it
    TYPE(t_field),POINTER :: f_phis(:)
    TYPE(t_field),POINTER :: f_geopot(:)
    TYPE(t_field),POINTER :: f_ps(:)
    TYPE(t_field),POINTER :: f_theta_rhodz(:)
    TYPE(t_field),POINTER :: f_ue(:)
    TYPE(t_field),POINTER :: f_wflux(:)
    TYPE(t_field),POINTER :: f_q(:)

    LOGICAL:: firstcall,lastcall
    INTEGER :: ind
    TYPE(t_physics_inout) :: args

    IF(MOD(it,itau_physics)==0 .AND. phys_type/=phys_none) THEN

       ! as a result of the the two calls to add_du_phys,
       ! du_phys increases by u(after physics) - u (before physics)
       CALL add_du_phys(-1., f_ue)
       
       SELECT CASE(phys_type)
       CASE(phys_HS94)
          CALL held_suarez(f_ps,f_theta_rhodz,f_ue) 
       CASE (phys_lmdz_generic)
         CALL physics_lmdz_generic(it ,f_phis, f_ps, f_theta_rhodz, f_ue, f_wflux, f_q)
       CASE (phys_external)
         CALL physics_external(it ,f_phis, f_ps, f_theta_rhodz, f_ue, f_wflux, f_q)
       CASE DEFAULT
          CALL transfert_request(f_ue, req_e1_vect)
          CALL physics_column(it/itau_physics, f_phis, f_geopot, f_ps, f_theta_rhodz, f_ue, f_q)
       END SELECT

       CALL transfert_request(f_theta_rhodz,req_i0)
       CALL transfert_request(f_ue,req_e0_vect)
       CALL transfert_request(f_q,req_i0)

       CALL add_du_phys(1., f_ue)
    END IF

    IF (mod(it,itau_out)==0 ) THEN
       CALL write_physics_tendencies
       CALL zero_du_phys
       SELECT CASE(phys_type)
       CASE (phys_DCMIP)
          CALL write_physics_dcmip
       CASE (phys_DCMIP2016)
          CALL write_physics_dcmip2016
       END SELECT
    END IF
    
  END SUBROUTINE physics

  SUBROUTINE write_physics_tendencies
    USE observable_mod, ONLY : f_buf_ulon, f_buf_ulat
    USE wind_to_lonlat_mod
    USE output_field_mod
    CALL transfert_request(f_du_phys,req_e1_vect)
    CALL un2ulonlat(f_du_phys, f_buf_ulon, f_buf_ulat, (1./(dt*itau_out)))
    CALL output_field("dulon_phys",f_buf_ulon)
    CALL output_field("dulat_phys",f_buf_ulat)
  END SUBROUTINE write_physics_tendencies
    
  SUBROUTINE physics_column(it, f_phis, f_geopot, f_ps, f_theta_rhodz, f_ue, f_q)
    USE physics_dcmip_mod, ONLY : full_physics_dcmip => full_physics
    USE physics_dcmip2016_mod, ONLY : full_physics_dcmip2016 => full_physics
    USE etat0_venus_mod, ONLY : full_physics_venus=>full_physics
    USE theta2theta_rhodz_mod
    USE mpipara
    USE checksum_mod
    TYPE(t_field),POINTER :: f_phis(:)
    TYPE(t_field),POINTER :: f_geopot(:)
    TYPE(t_field),POINTER :: f_ps(:)
    TYPE(t_field),POINTER :: f_theta_rhodz(:)
    TYPE(t_field),POINTER :: f_ue(:)
    TYPE(t_field),POINTER :: f_q(:)
    REAL(rstd),POINTER :: phis(:)
    REAL(rstd),POINTER :: geopot(:,:)
    REAL(rstd),POINTER :: ps(:)
    REAL(rstd),POINTER :: temp(:,:)
    REAL(rstd),POINTER :: ue(:,:)
    REAL(rstd),POINTER :: dulon(:,:)
    REAL(rstd),POINTER :: dulat(:,:)
    REAL(rstd),POINTER :: q(:,:,:)
    REAL(rstd),POINTER :: p(:,:)
    REAL(rstd),POINTER :: pk(:,:)
    REAL(rstd),POINTER :: ulon(:,:)
    REAL(rstd),POINTER :: ulat(:,:)
    INTEGER :: it, ind

    physics_inout%it = it

    CALL theta_rhodz2temperature(f_ps,f_theta_rhodz,f_temp)
    
    DO ind=1,ndomain
       IF (.NOT. assigned_domain(ind)) CYCLE
       CALL swap_dimensions(ind)
       CALL swap_geometry(ind)
       phis=f_phis(ind)
       geopot=f_geopot(ind)
       ps=f_ps(ind)
       temp=f_temp(ind)
       ue=f_ue(ind)
       q=f_q(ind)
       p=f_p(ind)
       pk=f_pk(ind)
       ulon=f_ulon(ind)
       ulat=f_ulat(ind)
       CALL pack_physics(pack_info(ind), phis, geopot, ps, temp, ue, q, p, pk, ulon, ulat)
    END DO

    SELECT CASE(phys_type)
    CASE (phys_DCMIP)
       IF (is_omp_level_master) CALL full_physics_dcmip
    CASE (phys_DCMIP2016)
       IF (is_omp_level_master) CALL full_physics_dcmip2016
    CASE(phys_LB2012)
       IF (is_omp_level_master) CALL full_physics_venus
    CASE (phys_plugin)
       IF (is_omp_level_master) CALL physics_plugin
    CASE DEFAULT
       IF(is_master) PRINT *,'Internal error : illegal value of phys_type', phys_type
       STOP
    END SELECT

    DO ind=1,ndomain
       IF (.NOT. assigned_domain(ind)) CYCLE
       CALL swap_dimensions(ind)
       CALL swap_geometry(ind)
       ps=f_ps(ind)
       temp=f_temp(ind)
       q=f_q(ind)
       dulon=f_dulon(ind)
       dulat=f_dulat(ind)
       ! unpack_physics updates temp and q ; a transfer is required before updating u
       CALL unpack_physics(pack_info(ind), ps, temp, q, dulon, dulat)
    END DO
    
    CALL temperature2theta_rhodz(f_ps,f_temp,f_theta_rhodz)

    ! Transfer dulon, dulat
    CALL transfert_request(f_dulon,req_i0)
    CALL transfert_request(f_dulat,req_i0)

    DO ind=1,ndomain
       IF (.NOT. assigned_domain(ind)) CYCLE
       CALL swap_dimensions(ind)
       CALL swap_geometry(ind)
       ue=f_ue(ind)
       dulon=f_dulon(ind)
       dulat=f_dulat(ind)
       CALL compute_update_velocity(dulon, dulat, ue)
    END DO

  END SUBROUTINE physics_column

  SUBROUTINE pack_physics(info, phis, geopot, ps, temp, ue, q, p, pk, ulon, ulat )
    USE wind_to_lonlat_mod
    USE pression_mod
    USE theta2theta_rhodz_mod
    USE exner_mod
    TYPE(t_pack_info) :: info
    REAL(rstd) :: phis(iim*jjm)
    REAL(rstd) :: geopot(iim*jjm,llm+1)
    REAL(rstd) :: ps(iim*jjm)
    REAL(rstd) :: temp(iim*jjm,llm)
    REAL(rstd) :: pks(iim*jjm)
    REAL(rstd) :: pk(iim*jjm,llm)
    REAL(rstd) :: ue(3*iim*jjm,llm)
    REAL(rstd) :: q(iim*jjm,llm,nqtot)

    REAL(rstd) :: p(iim*jjm,llm+1)
    REAL(rstd) :: uc(iim*jjm,llm,3)
    REAL(rstd) :: ulon(iim*jjm,llm)
    REAL(rstd) :: ulat(iim*jjm,llm)

!$OMP BARRIER
    CALL compute_pression(ps,p,0) ! pressure at interfaces
!$OMP BARRIER
    CALL compute_exner(ps,p,pks,pk,0) 
!$OMP BARRIER
    CALL compute_wind_centered(ue,uc)
    CALL compute_wind_centered_lonlat(uc, ulon, ulat)
!$OMP BARRIER
    IF (is_omp_level_master) THEN
      CALL pack_domain(info, phis, physics_inout%phis)
      CALL pack_domain(info, geopot, physics_inout%geopot)
      CALL pack_domain(info, p, physics_inout%p)
      CALL pack_domain(info, pk, physics_inout%pk)
      CALL pack_domain(info, Temp, physics_inout%Temp)
      CALL pack_domain(info, ulon, physics_inout%ulon)
      CALL pack_domain(info, ulat, physics_inout%ulat)
      CALL pack_domain(info, q, physics_inout%q)
    ENDIF
!$OMP BARRIER
  END SUBROUTINE pack_physics

  SUBROUTINE unpack_physics(info, ps,temp, q, dulon, dulat)
    USE theta2theta_rhodz_mod
    TYPE(t_pack_info) :: info
    REAL(rstd) :: ps(iim*jjm)
    REAL(rstd) :: temp(iim*jjm,llm)
    REAL(rstd) :: q(iim*jjm,llm,nqtot)
    REAL(rstd) :: dulon(iim*jjm,llm)
    REAL(rstd) :: dulat(iim*jjm,llm)

    REAL(rstd) :: dq(iim*jjm,llm,nqtot)
    REAL(rstd) :: dTemp(iim*jjm,llm)

!$OMP BARRIER
    IF (is_omp_level_master) THEN
      CALL unpack_domain(info, dulon, physics_inout%dulon)
      CALL unpack_domain(info, dulat, physics_inout%dulat)
      CALL unpack_domain(info, dq, physics_inout%dq)
      CALL unpack_domain(info, Temp, physics_inout%Temp)
      CALL unpack_domain(info, dTemp, physics_inout%dTemp)
      q = q + physics_inout%dt_phys * dq
      Temp = Temp + physics_inout%dt_phys * dTemp
    ENDIF
!$OMP BARRIER

!    CALL compute_temperature2theta_rhodz(ps,Temp,theta_rhodz,0)
  END SUBROUTINE unpack_physics

  SUBROUTINE compute_update_velocity(dulon, dulat, ue)
    USE wind_from_lonlat_mod
    REAL(rstd) :: dulon(iim*jjm,llm)
    REAL(rstd) :: dulat(iim*jjm,llm)
    REAL(rstd) :: ue(3*iim*jjm,llm)
    REAL(rstd) :: duc(iim*jjm,llm,3)
    REAL(rstd) :: dt2, due
    INTEGER :: i,j,ij,l
    ! Reconstruct wind tendencies at edges and add to normal wind
    CALL compute_wind_centered_from_lonlat_compound(dulon,dulat,duc)
    dt2=.5*physics_inout%dt_phys
    DO l=ll_begin,ll_end
      DO j=jj_begin,jj_end
        DO i=ii_begin,ii_end
          ij=(j-1)*iim+i
          due = sum( (duc(ij,l,:) + duc(ij+t_right,l,:))*ep_e(ij+u_right,:) )
          ue(ij+u_right,l) = ue(ij+u_right,l) + dt2*due

          due = sum( (duc(ij,l,:) + duc(ij+t_lup,l,:))*ep_e(ij+u_lup,:) )
          ue(ij+u_lup,l)=ue(ij+u_lup,l) + dt2*due

          due = sum( (duc(ij,l,:) + duc(ij+t_ldown,l,:))*ep_e(ij+u_ldown,:) )
          ue(ij+u_ldown,l)=ue(ij+u_ldown,l) + dt2*due
        ENDDO
      ENDDO
    ENDDO
  END SUBROUTINE compute_update_velocity

END MODULE physics_mod
