module vertutils_m

!  Routines involving the vertical coordinate

   public :: setsig, sig2sigh, setsig_lbl, initial, height

contains

   subroutine setsig ( sig, sigh, n )
      real, intent(out), dimension(:) :: sig, sigh
      integer, intent(in) :: n
      integer :: k
!     Calculate sigma levels using CSIRO model formula

!     Set half levels
      do k=1,n+1
         sigh(k) = (n+1-k)**2*(n-2+2.*k)/n**3
      end do
!     Set full levels
      do k=1,n
         sig(k) = 0.5*(sigh(k)+sigh(k+1))
      end do
 
   end subroutine setsig
   
   subroutine sig2sigh(sig, sigh)

!     Calculate half level values of sigma from the full levels.
!     For the CSIRO model the level is at the mid-point of the layer. 

      real, intent(in),  dimension(:) :: sig
      real, intent(out), dimension(:) :: sigh

      sigh(1) = 1.0
      do k=1,size(sig)-1
         sigh(k+1) = 2*sig(k) - sigh(k)
      end do
      sigh(size(sig)+1) = 0.0

   end subroutine sig2sigh

   subroutine setsig_lbl ( sig, sigh, n )
      real, intent(out), dimension(:) :: sig, sigh
      integer, intent(in) :: n
      integer :: k
      real, dimension(n+1) :: plev
!     Calculate sigma levels as used in the LBL calculations described in 
!     Schwarzkopf and Fels 1991. These are equally spaced in pressure below 
!     100 mb and log spaced above that. The surface pressure is assumed to
!     be 1013 and there are flux levels at both 1013 and 1000.
!     Use n=122 to match their levels with a top at 0.001 mb.

!     Equally spaced from 1 to 0.1
      plev(1) = 1013.
      do k=2,min(n+1,47)
         plev(k) = 1000 - 20.0*(k-2)
      end do

!     Log spaced above this with 15 levels per decade
      do k=48,n
         plev(k) = exp ( log(100.) - (k-47)*log(10.0)/15.0 )
      end do

      plev(n+1) = 0.0

      print*, "PLEV", plev
      sigh = plev / plev(1)

!     Set full levels
      do k=1,n
         sig(k) = 0.5*(sigh(k)+sigh(k+1))
      end do
 
   end subroutine setsig_lbl

subroutine initial(nl,sig,sigh,am,algf,alogg)

!  Calculate the A matrix needed to integrate hydrostatic equation

   use physparams
   implicit none
   integer, intent(in) :: nl
   real, intent(in), dimension(nl) :: sig
   real, intent(in), dimension(nl+1) :: sigh
   real, intent(out), dimension(nl,nl) :: am
   real, intent(out), dimension(nl) :: algf, alogg

   real, dimension(nl,nl) :: amh
   real, dimension(nl) :: algh, atm, atp, btm, btp
   integer :: i, j, k

!  Sigma level arrays
   do k=1,nl
      algf(k)=alog(sig(k))
      algh(k)=alog(sigh(k))
   end do
   do k=2,nl
      alogg(k) = 1.0 / (algf(k) - algf(k-1))
   end do

   do k=1,nl-1
      btm(k)=1.0/(algf(k)-algf(k+1))
      btp(k)=-btm(k)
      atm(k)=-algf(k+1)*btm(k)
      atp(k)=algf(k)*btm(k)
   end do

!****  temp/geopotl coupling for phi=am*t+phist
!***  based upon t=a+b.log(sigma) : energy conserving -
!****  changes made to dynamical terms (in DYNMNL).
   am = 0.
   amh = 0.

!  Define temporary full level geopotential matrix values
   am(1,1)=rdry*(atm(1)*(algh(1)-algf(1))+btm(1)*(algh(1)**2-algf(1)**2)*0.5)
   am(1,2)=rdry*(atp(1)*(algh(1)-algf(1))+btp(1)*(algh(1)**2-algf(1)**2)*0.5)
   am(2,1)=rdry*(atm(1)*(algh(1)-algf(2))+btm(1)*(algh(1)**2-algf(2)**2)*0.5)
   am(2,2)=rdry*(atp(1)*(algh(1)-algf(2))+btp(1)*(algh(1)**2-algf(2)**2)*0.5)
   do i=3,nl
      do j=1,i
	 if ( j == i-1 ) then
	    am(i,j) = am(i-1,j)    &
              + rdry*(atm(i-1)*(algf(i-1)-algf(i))+btm(i-1)*(algf(i-1)**2  &
              - algf(i)**2)*0.5)
	 else if ( j == i) then
	    am(i,j) = am(i-1,j)    &
	      + rdry*(atp(i-1)*(algf(i-1)-algf(i))+btp(i-1)*(algf(i-1)**2  &
              - algf(i)**2)*0.5)
	 else 
	    am(i,j) = am(i-1,j)
	 end if
      end do
   end do

!  The model redefines the full level am from a half level matrix. This
!  is needed for energy conservation. However for diagnostic purposes it's
!  more accurate to skip this part.

end subroutine initial


!-----------------------------------------------------------------
subroutine height(lon,lat2,nl,tg,qg,zg,pg,am,phi)

!  Calculate geopotential height on sigma levels using virtual temperature
   use physparams
   implicit none
   integer, intent(in) :: lon, lat2, nl
   real, intent(in), dimension(lon,lat2,nl) :: tg, qg
   real, intent(in), dimension(lon,lat2)    :: zg, pg
   real, intent(in), dimension(nl,nl) :: am
   real, intent(out), dimension(lon,lat2,nl) :: phi

   real, dimension(lon,nl) :: tv
   integer :: j, k, lg 

   do lg=1,lat2

!     Calculate virtual temperature.
      tv = tg(:,lg,:) * (epsil+qg(:,lg,:))/(epsil*(1.+qg(:,lg,:)))

!     Calculate height on sigma levels
      do k=1,nl
	 phi(:,lg,k) = zg(:,lg) * grav
	 do j=1,nl
	    phi(:,lg,k) = phi(:,lg,k)+am(k,j)*tv(:,j)
	 end do
      end do
   end do       ! Latitude loop

   phi = phi / grav

   return
end subroutine height

end module vertutils_m