module stpjcmod
!$$$ module documentation block
!           .      .    .                                       .
! module:   stpjcmod    module for weak constraint stp routines
!  pgrmmr:  kleist
!
! abstract: module for Jc step routines
!
! program history log:
!   2012-01-21  kleist - consolidation of Jc step routines into single module
!
! subroutines included:
!
! attributes:
!   language: f90
!   machine:
!
!$$$ end documentation block

use kinds, only: r_kind,i_kind,r_quad
use constants, only: zero,two,one,half,zero_quad,one_quad,two_quad
use gsi_bundlemod, only: gsi_bundle,gsi_bundlegetpointer

implicit none

PRIVATE
PUBLIC stplimq,stplimg,stplimp,stplimv,stpjcdfi,stpjcpdry

contains

subroutine stplimq(rval,sval,sges,outmin,outmax,nstep,itbin)
!$$$  subprogram documentation block
!                .      .    .                                       .
! subprogram:    stplimq     calculate penalty and stepsize for limit of q
!   prgmmr: derber           org: np23                date: 1996-11-19
!
! abstract: calculate stepsize contribution and penalty for limiting q
!
! program history log:
!   1996-11-19  derber
!   1998-07-10  weiyu yang
!   1999-08-24  derber, j., treadon, r., yang, w., first frozen mpp version
!   2004-03-15  kleist, d., derber, j., treadon, r., use negative q only
!   2004-06-02  kleist, add penalty for excess moisture
!   2004-07-29  treadon - add only to module use, add intent in/out
!   2004-11-22  derber - modify for openMP
!   2006-09-18  derber - modify output b1 and b3
!   2007-06-04  derber  - use quad precision to get reproducability over number of processors
!   2007-06-04  derber  - use quad precision to get reproducability over number of processors
!   2008-08-14  derber  - optimize
!   2010-05-13  todling - update to use gsi_bundle
!   2010-07-10  todling - merge w/ r8741 (trunk); qx(:)->qx (who made the change?)
!   2011-12-27  kleist - add bins for 4d capability (4densvar option)
!
!   input argument list:
!     rq       - search direction
!     sq       - increment in grid space
!     sges     - step size estimates (4)
!     nstep    - number of step size estimates if == 0 then just do outer loop
!     itbin    - observation bin number (time level)
!
!   output argument list:
!     outmin(1:nstep)  - current penalty for negative q sges(1:nstep)
!     outmax(1:nstep)  - current penalty for excess q sges(1:nstep)
!
! attributes:
!   language: f90
!   machine:  ibm RS/6000 SP
!
!$$$
  use gridmod, only: lat1,lon1,lat2,lon2,nsig
  use jfunc, only: factqmin,factqmax
  use guess_grids, only: ges_q,ges_qsat
  implicit none

! Declare passed variables
  integer(i_kind)                     ,intent(in   ) :: nstep,itbin
  real(r_kind),dimension(max(1,nstep)),intent(in   ) :: sges
  real(r_quad),dimension(max(1,nstep)),intent(  out) :: outmin,outmax
  type(gsi_bundle)                    ,intent(in   ) :: rval,sval

! Declare local variables
  integer(i_kind) i,j,k,kk,ier,istatus
  real(r_kind) q,qx
  real(r_kind),pointer,dimension(:,:,:) :: rq,sq

  outmin=zero_quad; outmax=zero_quad

  if (factqmin==zero .and. factqmax==zero) return

! Retrieve pointers
! Simply return if any pointer not found
  ier=0
  call gsi_bundlegetpointer(sval,'q',sq,istatus);ier=istatus+ier
  call gsi_bundlegetpointer(rval,'q',rq,istatus);ier=istatus+ier
  if(ier/=0)return

! Loop over interior of subdomain
  if(nstep > 0)then
     do k = 1,nsig
        do j = 2,lon1+1
           do i = 2,lat1+1

!             Values for q using stepsizes
              q  = ges_q(i,j,k,itbin) + sq(i,j,k)
              do kk=1,nstep
                 qx = q + sges(kk)*rq(i,j,k)
                 if(qx < zero)then
                    outmin(kk)=outmin(kk)+factqmin*qx*qx/(ges_qsat(i,j,k,itbin)*ges_qsat(i,j,k,itbin))
                 else
                    if(qx > ges_qsat(i,j,k,itbin))then
                       outmax(kk)=outmax(kk)+factqmax*(qx-ges_qsat(i,j,k,itbin))* &
                            (qx-ges_qsat(i,j,k,itbin))/(ges_qsat(i,j,k,itbin)*ges_qsat(i,j,k,itbin))
                    end if
                 end if
              end do
           end do
        end do
     end do
  else
     do k = 1,nsig
        do j = 2,lon1+1
           do i = 2,lat1+1

!             Values for q using stepsizes
              q  = ges_q(i,j,k,itbin)
              if(q < zero)then
                 outmin(1)=outmin(1)+factqmin*q*q/(ges_qsat(i,j,k,itbin)*ges_qsat(i,j,k,itbin))
              else
                 if(q > ges_qsat(i,j,k,itbin))then
                    outmax(1)=outmax(1)+factqmax*(q-ges_qsat(i,j,k,itbin))*(q-ges_qsat(i,j,k,itbin))/ &
                             (ges_qsat(i,j,k,itbin)*ges_qsat(i,j,k,itbin))
                 end if
              end if
           end do
        end do
     end do
  end if

  do kk=2,nstep
     outmin(kk)=outmin(kk)-outmin(1)
     outmax(kk)=outmax(kk)-outmax(1)
  end do
  return
end subroutine stplimq

subroutine stplimg(rval,sval,sges,out,nstep)
!$$$  subprogram documentation block
!                .      .    .                                       .
! subprogram:    stplimg     calculate penalty and stepsize for limit of q
!   prgmmr: derber           org: np23                date: 1996-11-19
!
! abstract: calculate stepsize contribution and penalty for limiting q
!
! program history log:
!   2011-02-23  zhu
!
!   input argument list:
!     rg       - search direction
!     sg       - increment in grid space
!     sges     - step size estimates (4)
!     nstep    - number of step size estimates if == 0 then just do outer loop
!
!   output argument list:
!     out(1:nstep)  - current penalty for negative gust sges(1:nstep)
!
! attributes:
!   language: f90
!   machine:  ibm RS/6000 SP
!
!$$$
  use gridmod, only: lat1,lon1,lat2,lon2,nsig
  use jfunc, only: factg,ggues
  implicit none

! Declare passed variables
  integer(i_kind)                     ,intent(in   ) :: nstep
  real(r_kind),dimension(max(1,nstep)),intent(in   ) :: sges
  real(r_quad),dimension(max(1,nstep)),intent(  out) :: out
  type(gsi_bundle)                    ,intent(in   ) :: rval,sval

! Declare local variables
  integer(i_kind) i,j,k,kk,ier,istatus
  real(r_kind) gust,gx
  real(r_kind),pointer,dimension(:,:) :: rg,sg

  out=zero_quad

  if (factg==zero) return

! Retrieve pointers
! Simply return if any pointer not found
  ier=0
  call gsi_bundlegetpointer(sval,'gust',sg,istatus);ier=istatus+ier
  call gsi_bundlegetpointer(rval,'gust',rg,istatus);ier=istatus+ier
  if(ier/=0)return

! Loop over interior of subdomain
  if(nstep > 0)then
     do j = 2,lon1+1
        do i = 2,lat1+1

!          Values for gust using stepsizes
           gust  = ggues(i,j) + sg(i,j)
           do kk=1,nstep
              gx = gust + sges(kk)*rg(i,j)
              if(gx < zero)then
                 out(kk)=out(kk)+factg*gx*gx/(ggues(i,j)*ggues(i,j))
              end if
           end do
        end do
     end do
  end if

  do kk=2,nstep
     out(kk)=out(kk)-out(1)
  end do
  return
end subroutine stplimg

subroutine stplimp(rval,sval,sges,out,nstep)
!$$$  subprogram documentation block
!                .      .    .                                       .
! subprogram:    stplimp     calculate penalty and stepsize for limit of q
!   prgmmr: derber           org: np23                date: 1996-11-19
!
! abstract: calculate stepsize contribution and penalty for limiting q
!
! program history log:
!   2011-02-23  zhu
!
!   input argument list:
!     rp       - search direction
!     sp       - increment in grid space
!     sges     - step size estimates (4)
!     nstep    - number of step size estimates if == 0 then just do outer loop
!
!   output argument list:
!     out(1:nstep)  - current penalty for negative pblh sges(1:nstep)
!
! attributes:
!   language: f90
!   machine:  ibm RS/6000 SP
!
!$$$
  use gridmod, only: lat1,lon1,lat2,lon2,nsig
  use jfunc, only: factp,pgues
  implicit none

! Declare passed variables
  integer(i_kind)                     ,intent(in   ) :: nstep
  real(r_kind),dimension(max(1,nstep)),intent(in   ) :: sges
  real(r_quad),dimension(max(1,nstep)),intent(  out) :: out
  type(gsi_bundle)                    ,intent(in   ) :: rval,sval

! Declare local variables
  integer(i_kind) i,j,k,kk,ier,istatus
  real(r_kind) pblh,px
  real(r_kind),pointer,dimension(:,:) :: rp,sp

  out=zero_quad

  if (factp==zero) return

! Retrieve pointers
! Simply return if any pointer not found
  ier=0
  call gsi_bundlegetpointer(sval,'pblh',sp,istatus);ier=istatus+ier
  call gsi_bundlegetpointer(rval,'pblh',rp,istatus);ier=istatus+ier
  if(ier/=0)return

! Loop over interior of subdomain
  if(nstep > 0)then
     do j = 2,lon1+1
        do i = 2,lat1+1

!          Values for pblh using stepsizes
           pblh  = pgues(i,j) + sp(i,j)
           do kk=1,nstep
              px = pblh + sges(kk)*rp(i,j)
              if(px < zero)then
                 out(kk)=out(kk)+factp*px*px/(pgues(i,j)*pgues(i,j))
              end if
           end do
        end do
     end do
  end if

  do kk=2,nstep
     out(kk)=out(kk)-out(1)
  end do
  return
end subroutine stplimp

subroutine stplimv(rval,sval,sges,out,nstep)
!$$$  subprogram documentation block
!                .      .    .                                       .
! subprogram:    stplimv     calculate penalty and stepsize for limit of q
!   prgmmr: derber           org: np23                date: 1996-11-19
!
! abstract: calculate stepsize contribution and penalty for limiting q
!
! program history log:
!   2011-02-23  zhu
!
!   input argument list:
!     rg       - search direction
!     sg       - increment in grid space
!     sges     - step size estimates (4)
!     nstep    - number of step size estimates if == 0 then just do outer loop
!
!   output argument list:
!     out(1:nstep)  - current penalty for negative vis sges(1:nstep)
!
! attributes:
!   language: f90
!   machine:  ibm RS/6000 SP
!
!$$$
  use gridmod, only: lat1,lon1,lat2,lon2,nsig
  use jfunc, only: factv,vgues
  implicit none

! Declare passed variables
  integer(i_kind)                     ,intent(in   ) :: nstep
  real(r_kind),dimension(max(1,nstep)),intent(in   ) :: sges
  real(r_quad),dimension(max(1,nstep)),intent(  out) :: out
  type(gsi_bundle)                    ,intent(in   ) :: rval,sval

! Declare local variables
  integer(i_kind) i,j,k,kk,ier,istatus
  real(r_kind) vis,vx
  real(r_kind),pointer,dimension(:,:) :: rg,sg

  out=zero_quad

  if (factv==zero) return

! Retrieve pointers
! Simply return if any pointer not found
  ier=0
  call gsi_bundlegetpointer(sval,'vis',sg,istatus);ier=istatus+ier
  call gsi_bundlegetpointer(rval,'vis',rg,istatus);ier=istatus+ier
  if(ier/=0)return

! Loop over interior of subdomain
  if(nstep > 0)then
     do j = 2,lon1+1
        do i = 2,lat1+1

!          Values for vis using stepsizes
           vis  = vgues(i,j) + sg(i,j)
           do kk=1,nstep
              vx = vis + sges(kk)*rg(i,j)
              if(vx < zero)then
                 out(kk)=out(kk)+factv*vx*vx/(vgues(i,j)*vgues(i,j))
              end if
           end do
        end do
     end do
  end if

  do kk=2,nstep
     out(kk)=out(kk)-out(1)
  end do
  return
end subroutine stplimv
 
subroutine stpjcpdry(rval,sval,pen,b,c)
!$$$  subprogram documentation block
!                .      .    .                                       .
! subprogram:    stpjcpdry   penalty and stp size for mean dry ps conservation
!   prgmmr: kleist           org: np23                date: 2009-07-07
!
! abstract: calculate stepsize contribution and penalty for limiting changes to global
!           mean dry ps increment
!
! program history log:
!   2009-07-07  kleist
!   2010-05-13  todling - update to use gsi_bundle
!   2010-05-25  derber  - modify to decrease number of communications
!   2010-08-18  hu      - add qpvals= to mpl_allreduce call
!   2011-11-01  eliu    - add handling for ql & qi increments and search directions
!
!   input argument list:
!     rq       - q search direction
!     rc       - cloud water search direction
!     rp       - surface pressure search direction
!     sq       - q increment
!     sc       - cloud water increment
!     sp       - increment in grid space
!
!   output argument list:
!     pen      - current penalty for mean dry pressure constraint
!     b        - contribution to numerator
!     c        - contribution to denomenator
!
! attributes:
!   language: f90
!   machine:  ibm RS/6000 SP
!
!$$$
  use mpimod, only: mype
  use gridmod, only: lat2,lon2,nsig,wgtlats,nlon,istart
  use guess_grids, only:  ges_prsi,ntguessig
  use mpl_allreducemod, only: mpl_allreduce
  use jcmod, only: bamp_jcpdry
  use gsi_bundlemod, only: assignment(=)
  use gsi_metguess_mod,  only: gsi_metguess_get
  implicit none

! Declare passed variables
  type(gsi_bundle),intent(in   ) :: sval
  type(gsi_bundle),intent(in   ) :: rval
  real(r_quad)    ,intent(  out) :: pen,b,c

! Declare local variables
  real(r_quad),dimension(2):: dmass
  real(r_quad) :: rcon,con
  integer(i_kind) i,j,k,it,mm1,ii,ier,icw,iql,iqi,istatus
  real(r_kind),pointer,dimension(:,:,:) :: rq,sq,rc,sc,rql,rqi,sql,sqi
  real(r_kind),pointer,dimension(:,:)   :: rp,sp

  pen=zero_quad ; b=zero_quad ; c=zero_quad
  it=ntguessig

! Retrieve pointers
! Simply return if any pointer not found
  ier=0; icw=0; iql=0; iqi=0
  call gsi_bundlegetpointer(sval,'q' ,sq, istatus);ier=istatus+ier
  call gsi_bundlegetpointer(sval,'cw',sc, istatus);icw=istatus+icw
  call gsi_bundlegetpointer(sval,'ql',sql,istatus);iql=istatus+iql
  call gsi_bundlegetpointer(sval,'qi',sqi,istatus);iqi=istatus+iqi
  call gsi_bundlegetpointer(sval,'ps',sp, istatus);ier=istatus+ier
  call gsi_bundlegetpointer(rval,'q' ,rq, istatus);ier=istatus+ier
  call gsi_bundlegetpointer(rval,'cw',rc, istatus);icw=istatus+icw
  call gsi_bundlegetpointer(rval,'ql',rql,istatus);iql=istatus+iql
  call gsi_bundlegetpointer(rval,'qi',rqi,istatus);iqi=istatus+iqi
  call gsi_bundlegetpointer(rval,'ps',rp, istatus);ier=istatus+ier
  if(ier+icw*(iql+iqi)/=0)then
    if (mype==0) write(6,*)'stpjcpdry: checking ier+icw*(iql+iqi)=', ier+icw*(iql+iqi)
    return
  end if
 
  dmass=zero_quad
  rcon=one_quad/(two_quad*float(nlon))
  mm1=mype+1

! Calculate mean surface pressure contribution in subdomain
  do j=2,lon2-1
    do i=2,lat2-1
      ii=istart(mm1)+i-2
      con=wgtlats(ii)*rcon
      dmass(1)=dmass(1)+sp(i,j)*con
      dmass(2)=dmass(2)+rp(i,j)*con
    end do
  end do
! Remove water to get incremental dry ps
  do k=1,nsig
     do j=2,lon2-1
        do i=2,lat2-1
           ii=istart(mm1)+i-2
           con=(ges_prsi(i,j,k,it)-ges_prsi(i,j,k+1,it))*wgtlats(ii)*rcon
           if(icw==0)then
              dmass(1)=dmass(1) - (sq(i,j,k)+sc(i,j,k))*con
              dmass(2)=dmass(2) - (rq(i,j,k)+rc(i,j,k))*con
           else
              dmass(1)=dmass(1) - (sq(i,j,k)+sql(i,j,k)+sqi(i,j,k))*con
              dmass(2)=dmass(2) - (rq(i,j,k)+rql(i,j,k)+rqi(i,j,k))*con
           endif
        end do
     end do
  end do

  call mpl_allreduce(2,qpvals=dmass)

  if (mype==0) then

!    Now penalize non-zero global mean dry ps increment
!    Notice there will only be a contribution from PE=0

     pen = bamp_jcpdry*dmass(1)*dmass(1)
     b  = -bamp_jcpdry*dmass(2)*dmass(1)
     c  =  bamp_jcpdry*dmass(2)*dmass(2)
  end if

  return
end subroutine stpjcpdry

subroutine stpjcdfi(rval,sval,pen,b,c)
!$$$  subprogram documentation block
!                .      .    .                                       .
! subprogram:    stpjcdfi   penalty and stp size for weak constraint DFI
!   prgmmr: kleist           org: np23                date: 2012-01-19
!
! abstract: calculate stepsize contribution and penalty for Jc DFI
!
! program history log:
!   2012-01-19  kleist - adaptation of evaljcdfi
!
!   input argument list:
!     rval     -
!     sval     -
!
!   output argument list:
!     pen      - current penalty for Jc DFI
!     b        - contribution to numerator
!     c        - contribution to denomenator
!
! attributes:
!   language: f90
!   machine:  ibm RS/6000 SP
!
!$$$
  use jcmod, only: wgtdfi,alphajc
  use gsi_bundlemod, only : self_add,self_mul,assignment(=)
  use gsi_4dvar, only: nobs_bins
  use mpimod, only: mype
  use state_vectors, only : allocate_state,deallocate_state
  use state_vectors, only : dot_product
  implicit none

! Declare passed variables
  type(gsi_bundle),dimension(nobs_bins),intent(in) :: sval
  type(gsi_bundle),dimension(nobs_bins),intent(in) :: rval

  real(r_quad),intent(out) :: pen,b,c

! Declare local variables
  integer(i_kind) :: jj,idfi
  real(r_quad) :: pjc,rjc
  type(gsi_bundle) :: sfilter,afilter
  type(gsi_bundle) :: rfilter,bfilter

  pen=zero_quad ; b=zero_quad ; c=zero_quad

!************************************************************************************

  idfi = (nobs_bins-1)/2+1
  call allocate_state(sfilter)
  call allocate_state(afilter)
  call allocate_state(rfilter)
  call allocate_state(bfilter)

! Compute filtered state
  sfilter=zero
  rfilter=zero
  do jj=1,nobs_bins
     call self_add(sfilter,wgtdfi(jj),sval(jj))
     call self_add(rfilter,wgtdfi(jj),rval(jj))
  enddo

! Compute difference from filtered state
  call self_add(sfilter,-one,sval(idfi))
  call self_add(rfilter,-one,rval(idfi))

! Apply Jc multiplicative factor
  call self_mul(sfilter,alphajc)
  call self_mul(rfilter,alphajc)

! Convert to energy norm/apply grid factors
  call enorm_state(sfilter,pjc,afilter)
  call enorm_state(rfilter,rjc,bfilter)

! Penalty, b, c
  pen =  dot_product(afilter,afilter)
  b   = -dot_product(afilter,bfilter)
  c   =  dot_product(bfilter,bfilter)

  if (mype==0) write(6,*)'Jc DFI penalty = ',pen

  call deallocate_state(sfilter)
  call deallocate_state(afilter)
  call deallocate_state(rfilter)
  call deallocate_state(bfilter)

  return
end subroutine stpjcdfi

end module stpjcmod


