module polar_avg 4,8
!----------------------------------------------------------------------- 
! 
! Purpose: 
!  These routines are used by the fv dycore to set the collocated 
!  pole points at the limits of the latitude dimension to the same 
!  value. 
!
! Methods: 
!  The repro_sum reproducible distributed sum is used for these 
!  calculations.
!
! Author: A. Mirin
! 
!-----------------------------------------------------------------------

!-----------------------------------------------------------------------
!- use statements ------------------------------------------------------
!-----------------------------------------------------------------------
   use shr_kind_mod,  only: r8 => shr_kind_r8
   use dycore,        only: dycore_is
   use dyn_grid,      only: get_dyn_grid_parm
   use phys_grid,     only: get_ncols_p, get_lat_all_p
   use ppgrid,        only: begchunk, endchunk, pcols
   use repro_sum_mod, only: repro_sum

!-----------------------------------------------------------------------
!- module boilerplate --------------------------------------------------
!-----------------------------------------------------------------------
   implicit none
   private
   save

!-----------------------------------------------------------------------
! Public interfaces ----------------------------------------------------
!-----------------------------------------------------------------------
   public :: &
      polar_average           ! support for LR dycore polar averaging


   interface polar_average 5
      module procedure polar_average2d, polar_average3d
   end interface

   CONTAINS
!
!========================================================================
!

   subroutine polar_average2d(field) 1,9
!----------------------------------------------------------------------- 
! Purpose: Set the collocated pole points at the limits of the latitude 
!          dimension to the same value. 
! Author: J. Edwards
!-----------------------------------------------------------------------
!
! Arguments
!
     real(r8), intent(inout) :: field(pcols,begchunk:endchunk)
!
! Local workspace
!
     integer :: i, c, ln, ls, ncols
     integer :: plat, plon
     integer, allocatable :: lats(:)

     real(r8) :: sum(2)
     real(r8), allocatable :: n_pole(:), s_pole(:)
!
!-----------------------------------------------------------------------
!
     if(.not. dycore_is('LR')) return

     plon = get_dyn_grid_parm('plon')
     plat = get_dyn_grid_parm('plat')
     allocate(lats(pcols), n_pole(plon), s_pole(plon))
     ln=0
     ls=0
     do c=begchunk,endchunk
        call get_lat_all_p(c,pcols,lats) 
	ncols = get_ncols_p(c)

	do i=1,ncols
           if(lats(i).eq.1) then
              ln=ln+1
              n_pole(ln) = field(i,c)
           else if(lats(i).eq.plat) then
              ls=ls+1
              s_pole(ls) = field(i,c)
           end if
	enddo
        
     end do
     
     call repro_sum(n_pole, sum(1:1), ln, plon, 1, gbl_count=plon)

     call repro_sum(s_pole, sum(2:2), ls, plon, 1, gbl_count=plon)


     ln=0
     ls=0
     do c=begchunk,endchunk
        call get_lat_all_p(c,pcols,lats) 
	ncols = get_ncols_p(c)

	do i=1,ncols
           if(lats(i).eq.1) then
              ln=ln+1
              field(i,c) = sum(1)/plon
           else if(lats(i).eq.plat) then
              ls=ls+1
              field(i,c) = sum(2)/plon
           end if
	enddo
        
     end do

     deallocate(lats, n_pole, s_pole)
   
   end subroutine polar_average2d

!
!========================================================================
!


   subroutine polar_average3d(nlev, field) 1,9
!----------------------------------------------------------------------- 
! Purpose: Set the collocated pole points at the limits of the latitude 
!          dimension to the same value. 
! Author: J. Edwards
!-----------------------------------------------------------------------
!
! Arguments
!
     integer, intent(in) :: nlev
     real(r8), intent(inout) :: field(pcols,nlev,begchunk:endchunk)
!
! Local workspace
!
     integer :: i, c, ln, ls, ncols, k
     integer :: plat, plon
     integer, allocatable :: lats(:)

     real(r8) :: sum(nlev,2)
     real(r8), allocatable :: n_pole(:,:), s_pole(:,:)
!
!-----------------------------------------------------------------------
!
     if(.not. dycore_is('LR')) return

     plon = get_dyn_grid_parm('plon')
     plat = get_dyn_grid_parm('plat')
     allocate(lats(pcols), n_pole(plon,nlev), s_pole(plon,nlev))
     ln=0
     ls=0
     do c=begchunk,endchunk
        call get_lat_all_p(c,pcols,lats) 
	ncols = get_ncols_p(c)

        do i=1,ncols
           if(lats(i).eq.1) then
              ln=ln+1
              do k=1,nlev
                 n_pole(ln,k) = field(i,k,c)
              end do
           else if(lats(i).eq.plat) then
              ls=ls+1
              do k=1,nlev                 
                 s_pole(ls,k) = field(i,k,c)
              end do
           end if
        enddo
     end do
     
     call repro_sum(n_pole, sum(:,1), ln, plon, nlev, gbl_count=plon)

     call repro_sum(s_pole, sum(:,2), ls, plon, nlev, gbl_count=plon)

     ln=0
     ls=0
     do c=begchunk,endchunk
        call get_lat_all_p(c,pcols,lats) 
	ncols = get_ncols_p(c)

	do i=1,ncols
           if(lats(i).eq.1) then
              ln=ln+1
              do k=1,nlev
                 field(i,k,c) = sum(k,1)/plon
              end do
           else if(lats(i).eq.plat) then
              ls=ls+1
              do k=1,nlev
                 field(i,k,c) = sum(k,2)/plon
              end do
           end if
	enddo
        
     end do
   
     deallocate(lats, n_pole, s_pole)

   end subroutine polar_average3d

end module polar_avg