MODULE physics_interface_mod

  USE prec

  PRIVATE

  TYPE t_physics_inout
     ! Input, time-independent
     INTEGER :: ngrid
     REAL(rstd) :: dt_phys
     REAL(rstd), DIMENSION(:), POINTER :: Ai, lon, lat, phis
     ! Input, time-dependent
     REAL(rstd), DIMENSION(:,:), POINTER :: geopot, p, pk, Temp, ulon, ulat
     REAL(rstd), DIMENSION(:,:,:), POINTER :: q
     ! Output arrays
     REAL(rstd), DIMENSION(:,:), POINTER :: dTemp, dulon, dulat
     REAL(rstd), DIMENSION(:,:,:), POINTER :: dq
  END TYPE t_physics_inout

! physics_inout is used to exchange information with physics
! Field ngrid is initialized by physics.f90/init_physics. Its other fields
! must be defined by XX/init_physics (where XX =  e.g. physics_dcmip.f90) 
! by either pointing to internal data of the physics package
! or by a specific allocation
! size : (ngrid), (ngrid,llm) except p(ngrid,llm+1), (ngrid,llm,nqtot) 

  TYPE(t_physics_inout), SAVE :: physics_inout
!$OMP THREADPRIVATE(physics_inout) 
 
! pack_info contains indices used by pack/unpack routines
! to pack together the data of all the domains managed by the MPI process
! It is initialized by physics.f90/init_physics 

  TYPE t_pack_info
     INTEGER :: ngrid, & ! number of non-halo points in that domain
          nseg ! number of segments (contigous parts) in that domain
     ! size and start of each segment : ij domain index, k packed index
     INTEGER, ALLOCATABLE :: n(:), ij(:), k(:)
  END TYPE t_pack_info

  TYPE(t_pack_info), ALLOCATABLE, SAVE :: pack_info(:)
!$OMP THREADPRIVATE(pack_info)


  INTERFACE pack_field
     MODULE PROCEDURE pack_2D
     MODULE PROCEDURE pack_3D
     MODULE PROCEDURE pack_4D
  END INTERFACE pack_field

  INTERFACE unpack_field
     MODULE PROCEDURE unpack_2D
     MODULE PROCEDURE unpack_3D
     MODULE PROCEDURE unpack_4D
  END INTERFACE unpack_field

  INTERFACE pack_domain
     MODULE PROCEDURE pack_domain_2D
     MODULE PROCEDURE pack_domain_3D
     MODULE PROCEDURE pack_domain_4D
  END INTERFACE pack_domain

  INTERFACE unpack_domain
     MODULE PROCEDURE unpack_domain_2D
     MODULE PROCEDURE unpack_domain_3D
     MODULE PROCEDURE unpack_domain_4D
  END INTERFACE unpack_domain

  PUBLIC :: nb_extra_physics_2D, nb_extra_physics_3D, &
       t_physics_inout, physics_inout, &
       t_pack_info, pack_info, init_pack_before, init_pack_after, &
       pack_domain, pack_field, unpack_domain, unpack_field, &
       garbage_3D

CONTAINS
  
    SUBROUTINE init_pack_before
    USE icosa
    IMPLICIT NONE
    INTEGER :: ind, offset, ngrid

    offset=0
    ALLOCATE(pack_info(ndomain))
    DO ind=1,ndomain
       IF (.NOT. assigned_domain(ind)) CYCLE
       CALL swap_dimensions(ind)
       CALL swap_geometry(ind)
       CALL count_segments(domain(ind)%own, pack_info(ind))
       pack_info(ind)%k = pack_info(ind)%k + offset
       offset = offset + pack_info(ind)%ngrid
    END DO
    physics_inout%ngrid = offset

    ngrid=offset
    ! Input
    ALLOCATE(physics_inout%Ai(ngrid))
    ALLOCATE(physics_inout%lon(ngrid))
    ALLOCATE(physics_inout%lat(ngrid))
    ALLOCATE(physics_inout%phis(ngrid))
    ALLOCATE(physics_inout%p(ngrid,llm+1))
    ALLOCATE(physics_inout%geopot(ngrid,llm+1))
    ALLOCATE(physics_inout%pk(ngrid,llm))
    ALLOCATE(physics_inout%Temp(ngrid,llm))
    ALLOCATE(physics_inout%ulon(ngrid,llm))
    ALLOCATE(physics_inout%ulat(ngrid,llm))
    ALLOCATE(physics_inout%q(ngrid,llm,nqtot))
    ! Output (tendencies)
    ALLOCATE(physics_inout%dTemp(ngrid,llm))
    ALLOCATE(physics_inout%dulon(ngrid,llm))
    ALLOCATE(physics_inout%dulat(ngrid,llm))
    ALLOCATE(physics_inout%dq(ngrid,llm,nqtot))

  END SUBROUTINE init_pack_before

  SUBROUTINE count_segments(own, info)
    USE icosa
    IMPLICIT NONE
    LOGICAL, DIMENSION(:,:) :: own
    TYPE(t_pack_info) :: info

    INTEGER, DIMENSION(jjm) :: n
    INTEGER :: ngrid, nseg, i, j, jj, k
    INTEGER, PARAMETER :: method=4
    SELECT CASE(method)
    CASE(1) ! Copy all points, including halo (works)
       info%nseg=1
       info%ngrid=iim*jjm
       ALLOCATE(info%n(1))
       ALLOCATE(info%ij(1))
       ALLOCATE(info%k(1))
       info%n(1)=iim*jjm
       info%ij(1)=1
       info%k(1)=1
    CASE(2) ! Copy all points, including halo, one at a time (works, slow ?)
       info%nseg=iim*jjm
       info%ngrid=iim*jjm
       ALLOCATE(info%n(iim*jjm))
       ALLOCATE(info%ij(iim*jjm))
       ALLOCATE(info%k(iim*jjm))
       DO jj=1,iim*jjm
          info%n(jj) =1
          info%ij(jj)=jj
          info%k(jj) =jj
       END DO
    CASE(3) ! Copy non-halo points only, one at a time (works, slow ?)
       n=0
       n(jj_begin:jj_end)=COUNT(own(ii_begin:ii_end,jj_begin:jj_end),1)
       ngrid=SUM(n)
       info%ngrid=ngrid
       info%nseg=ngrid
       ALLOCATE(info%n(ngrid))
       ALLOCATE(info%ij(ngrid))
       ALLOCATE(info%k(ngrid))
       jj=1
       DO j=1,jjm
          DO i=1,iim
             IF(own(i,j)) THEN
                info%n(jj)=1
                info%k(jj)=jj
                info%ij(jj) = iim*(j-1)+i
                jj=jj+1
             END IF
          END DO
       END DO
       
    CASE DEFAULT ! Copy non-halo points only, as contiguous segments (works)
       n=0
       n(jj_begin:jj_end)=COUNT(own(ii_begin:ii_end,jj_begin:jj_end),1)
       ngrid=SUM(n)
       info%ngrid=ngrid
       nseg=COUNT(n>0)
       info%nseg=nseg
       ALLOCATE(info%n(nseg))
       ALLOCATE(info%ij(nseg))
       ALLOCATE(info%k(nseg))
       info%n(:)=0
       info%ij(:)=0
       info%k(:)=0
       
       jj=1
       k=1
       DO j=jj_begin,jj_end
          IF(n(j)>0) THEN
             ! find first .TRUE. value in own(:,j)
             DO i=ii_begin,ii_end
                IF(own(i,j)) THEN
                   info%n(jj)=n(j)
                   info%k(jj)=k
                   info%ij(jj) = iim*(j-1)+i
                   IF(COUNT(own(i:i+n(j)-1,j)) /= n(j)) STOP
                   EXIT
                END IF
             END DO
             k = k + n(j)
             jj=jj+1
          END IF
       END DO

       IF(k-1/=ngrid) THEN
          PRINT *, 'Total number of grid points inconsistent', k-1, ngrid
          STOP
       END IF
       IF(jj-1/=nseg) THEN
          PRINT *, 'Number of segments inconsistent', jj-1, nseg
          STOP
       END IF

    END SELECT
    
    PRINT *, 'count_segments', info%nseg, info%ngrid, SUM(info%n), COUNT(own), iim*jjm
  END SUBROUTINE count_segments

  SUBROUTINE init_pack_after
    USE icosa
    IMPLICIT NONE
    INTEGER :: ind, offset
    DO ind=1,ndomain
       IF (.NOT. assigned_domain(ind)) CYCLE
       CALL swap_dimensions(ind)
       CALL swap_geometry(ind)
       CALL pack_domain_2D(pack_info(ind), Ai, physics_inout%Ai)
       CALL pack_domain_2D(pack_info(ind), lon_i, physics_inout%lon)
       CALL pack_domain_2D(pack_info(ind), lat_i, physics_inout%lat)
    END DO
  END SUBROUTINE init_pack_after

!-------------------------------- Pack / Unpack 2D ---------------------------

  SUBROUTINE pack_2D(f_2D, packed)
    USE icosa
    IMPLICIT NONE
    TYPE(t_field),POINTER :: f_2D(:)
    REAL(rstd)            :: packed(:)
    REAL(rstd), POINTER   :: loc(:)
    INTEGER :: ind
    DO ind=1,ndomain
       IF (.NOT. assigned_domain(ind)) CYCLE
       loc = f_2D(ind)
       CALL pack_domain_2D(pack_info(ind), loc, packed)
    END DO
  END SUBROUTINE pack_2D

  SUBROUTINE unpack_2D(f_2D, packed) 
    USE icosa
    IMPLICIT NONE
    TYPE(t_field),POINTER :: f_2D(:)
    REAL(rstd)            :: packed(:)
    REAL(rstd), POINTER   :: loc(:)
    INTEGER :: ind
    DO ind=1,ndomain
       IF (.NOT. assigned_domain(ind)) CYCLE
       loc = f_2D(ind)
       CALL unpack_domain_2D(pack_info(ind), loc, packed)
    END DO
  END SUBROUTINE unpack_2D

  SUBROUTINE pack_domain_2D(info, loc, glob)
    USE icosa
    IMPLICIT NONE
    TYPE(t_pack_info) :: info
    REAL(rstd), DIMENSION(:) :: loc, glob
    INTEGER :: jj,n,k,ij
    DO jj=1, info%nseg
       n = info%n(jj)-1
       ij = info%ij(jj)
       k = info%k(jj)
       glob(k:k+n) = loc(ij:ij+n)
    END DO
  END SUBROUTINE pack_domain_2D

  SUBROUTINE unpack_domain_2D(info, loc, glob)
    IMPLICIT NONE
    TYPE(t_pack_info) :: info
    REAL(rstd), DIMENSION(:) :: loc, glob
    INTEGER :: jj,n,k,ij
    DO jj=1, info%nseg
       n = info%n(jj)-1
       ij = info%ij(jj)
       k = info%k(jj)
       loc(ij:ij+n) = glob(k:k+n)
    END DO
  END SUBROUTINE unpack_domain_2D

!-------------------------------- Pack / Unpack 3D ---------------------------

  SUBROUTINE pack_3D(f_3D, packed)
    USE icosa
    IMPLICIT NONE
    TYPE(t_field),POINTER :: f_3D(:)
    REAL(rstd)            :: packed(:,:)
    REAL(rstd), POINTER   :: loc(:,:)
    INTEGER :: ind
    DO ind=1,ndomain
       IF (.NOT. assigned_domain(ind)) CYCLE
       loc = f_3D(ind)
       CALL pack_domain_3D(pack_info(ind), loc, packed)
    END DO
  END SUBROUTINE pack_3D

  SUBROUTINE unpack_3D(f_3D, packed) 
    USE icosa
    IMPLICIT NONE
    TYPE(t_field),POINTER :: f_3D(:)
    REAL(rstd)            :: packed(:,:)
    REAL(rstd), POINTER   :: loc(:,:)
    INTEGER :: ind
    DO ind=1,ndomain
       IF (.NOT. assigned_domain(ind)) CYCLE
       loc = f_3D(ind)
       CALL unpack_domain_3D(pack_info(ind), loc, packed)
    END DO
  END SUBROUTINE unpack_3D

  SUBROUTINE pack_domain_3D(info, loc, glob)
    IMPLICIT NONE
    TYPE(t_pack_info) :: info
    REAL(rstd), DIMENSION(:,:) :: loc, glob
    INTEGER :: jj,n,k,ij
    DO jj=1, info%nseg
       n = info%n(jj)-1
       ij = info%ij(jj)
       k = info%k(jj)
       glob(k:k+n,:) = loc(ij:ij+n,:)
    END DO
  END SUBROUTINE pack_domain_3D

  SUBROUTINE unpack_domain_3D(info, loc, glob)
    IMPLICIT NONE
    TYPE(t_pack_info) :: info
    REAL(rstd), DIMENSION(:,:) :: loc, glob
    INTEGER :: jj,n,k,ij
    DO jj=1, info%nseg
       n = info%n(jj)-1
       ij = info%ij(jj)
       k = info%k(jj)
       loc(ij:ij+n,:) = glob(k:k+n,:)
    END DO    
  END SUBROUTINE unpack_domain_3D

  SUBROUTINE garbage_3D(loc,own)
    USE icosa
    IMPLICIT NONE
    LOGICAL :: own(iim,jjm)
    REAL(rstd) :: loc(iim*jjm,llm)
    INTEGER :: i,j,ij
    ! write garbage in non-owned points
    DO j=1,jjm
       DO i=1,iim
          IF(.NOT.own(i,j)) THEN
             ij=iim*(j-1)+i
             loc(ij,:)=-1e30
          END IF
       END DO
    END DO   
  END SUBROUTINE garbage_3D

!-------------------------------- Pack / Unpack 4D ---------------------------

  SUBROUTINE pack_4D(f_4D, packed)
    USE icosa
    IMPLICIT NONE
    TYPE(t_field),POINTER :: f_4D(:)
    REAL(rstd)            :: packed(:,:,:)
    REAL(rstd), POINTER   :: loc(:,:,:)
    INTEGER :: ind
    DO ind=1,ndomain
       IF (.NOT. assigned_domain(ind)) CYCLE
       loc = f_4D(ind)
       CALL pack_domain_4D(pack_info(ind), loc, packed)
    END DO
  END SUBROUTINE pack_4D

  SUBROUTINE unpack_4D(f_4D, packed) 
    USE icosa
    IMPLICIT NONE
    TYPE(t_field),POINTER :: f_4D(:)
    REAL(rstd)            :: packed(:,:,:)
    REAL(rstd), POINTER   :: loc(:,:,:)
    INTEGER :: ind
    DO ind=1,ndomain
       IF (.NOT. assigned_domain(ind)) CYCLE
       loc = f_4D(ind)
       CALL unpack_domain_4D(pack_info(ind), loc, packed)
    END DO
  END SUBROUTINE unpack_4D

  SUBROUTINE pack_domain_4D(info, loc, glob)
    IMPLICIT NONE
    TYPE(t_pack_info) :: info
    REAL(rstd), DIMENSION(:,:,:) :: loc, glob
    INTEGER :: jj,n,k,ij
    DO jj=1, info%nseg
       n = info%n(jj)-1
       ij = info%ij(jj)
       k = info%k(jj)
       glob(k:k+n,:,:) = loc(ij:ij+n,:,:)
    END DO
  END SUBROUTINE pack_domain_4D

  SUBROUTINE unpack_domain_4D(info, loc, glob)
    IMPLICIT NONE
    TYPE(t_pack_info) :: info
    REAL(rstd), DIMENSION(:,:,:) :: loc, glob
    INTEGER :: jj,n,k,ij
    DO jj=1, info%nseg
       n = info%n(jj)-1
       ij = info%ij(jj)
       k = info%k(jj)
       loc(ij:ij+n,:,:) = glob(k:k+n,:,:)
    END DO
  END SUBROUTINE unpack_domain_4D

END MODULE physics_interface_mod