!============================================================
! <next few lines under version control, D O  N O T  E D I T>
! $Date$
! $Author$
! $Revision$
! $Id$
!============================================================

!========================================================================
! QG+1 PV inversion module
!========================================================================

!========================================================================
MODULE pvinv_mod

    USE netcdf
    USE spectral_mod

CONTAINS

!========================================================================
SUBROUTINE pvinv_main()

    implicit none

    real, dimension(2*kmax,2*lmax)      :: itbxy,ittxy
    real, dimension(2*kmax,2*lmax,pmax) :: ipvxy

    real, dimension(2*kmax,2*lmax)      :: thbB, thbT
    real, dimension(2*kmax,2*lmax,pmax) :: pvb

    complex, dimension(2*kmax,2*lmax)      :: phib,phit,ub,ut,vb,vt,tbsp,ttsp
    complex, dimension(2*kmax,2*lmax,pmax) :: phi,u,v,theta,w,pvsp

    complex, dimension(2*kmax,2*lmax)      :: phi0b,phi0t,phi1b,phi1t
    complex, dimension(2*kmax,2*lmax,pmax) :: F1,G1,P1
    complex, dimension(2*kmax,2*lmax,pmax) :: phi0,phi1
    complex, dimension(2*kmax,2*lmax,pmax) :: wcont

    real, dimension(2*kmax,2*lmax,pmax) :: epv,vort,div

    integer :: k, itim, ntims

    character(len=64), parameter :: inpfile = 'pvInput.nc'
    character(len=64), parameter :: outfile = 'pvOutput.nc'
    character(len=64), parameter :: bsefile = 'sqgBasic.nc'

    real, dimension(2*kmax,2*lmax,2),    parameter :: blank2 = 0.0
    real, dimension(2*kmax,2*lmax,pmax), parameter :: blank3 = 0.0

    ntims = get_ntimes(inpfile)

    call write_diag(outfile,0,blank2,blank2,blank2,blank2,blank3,blank3,blank3,blank3,blank3,blank3)

    if ( ibase ) call init(bsefile,1,thbB,thbT,pvb)

    do itim = 1, ntims

    call init(inpfile,itim,itbxy,ittxy,ipvxy)

    if ( ibase ) then
        itbxy = itbxy - thbB
        ittxy = ittxy - thbT
        ipvxy = ipvxy - pvb
    endif

    ! max values
    if (verbose .gt. 0)  then
    print*,'max |theta(z=0)| = ',maxval(abs(itbxy))
    print*,'max |theta(z=1)| = ',maxval(abs(ittxy))
    print*,'max |pv| = ',maxval(abs(ipvxy))
    endif

    ! map into spectral space at the same resolution
    call xy_to_sp(cmplx(itbxy,0.),tbsp,2*kmax,2*lmax,kmax,lmax)
    call xy_to_sp(cmplx(ittxy,0.),ttsp,2*kmax,2*lmax,kmax,lmax)
    !$omp parallel do schedule(dynamic) private(k)
    do k=1,pmax
        call xy_to_sp(cmplx(ipvxy(:,:,k),0.),pvsp(:,:,k),2*kmax,2*lmax,kmax,lmax)
    enddo
    !$omp end parallel do

    if (verbose .gt. 0) print*,'inverting for leading order geopotential...'
    call invert(tbsp,ttsp,pvsp,phi0b,phi0t,phi0)

    phi = phi0; phib = phi0b; phit = phi0t
    if (verbose .gt. 1) print*,'max nondimensional leading-order phi = ',maxval(abs(phi))

    if (order .eq. 1 .or. order .eq. 2 .and. Ross .ne. 0.) then

    if (verbose .gt. 0) print*,'inverting for O(R) fields...'
    call invert_R(tbsp,ttsp,phi0b,phi0t,phi0,F1,G1,P1,phi1,phi1b,phi1t,pvsp)

    if (order .eq. 1) then
        phi = Ross*phi1; phib = Ross*phi1b; phit = Ross*phi1t
    else
        phi = phi + Ross*phi1
        phib = phib + Ross*phi1b
        phit = phit + Ross*phi1t
    endif

    endif

    if (verbose .gt. 0) print*,'calculating u and v...'
    call uvwtp(tbsp,ttsp,phi0,phi0b,phi0t,F1,G1,P1,phi1,u,ub,ut,v,vb,vt,w,theta)
    if (verbose .gt. 1) print*,'max |u| = ',maxval(abs(u))
    if (verbose .gt. 1) print*,'max |v| = ',maxval(abs(v))

    ! check w by vertically integrating the continuity equation.
    call w_check(u,v,w,wcont)

    ! check Ertel PV
    call epv_pseudoheight(u,v,theta,ub,vb,tbsp,ut,vt,ttsp,epv)
    !call epv_pseudoheight_new(u,v,theta,tbsp,ttsp,phi0,phi0b,phi0t,epv)
    do k = 1,pmax
    !    if (verbose .gt. 1) print*,'pv check at level ',k,': ',maxval(abs(epv(:,:,k)-ipvxy(:,:,k)))
        if (verbose .gt. 1) print*,'pv check at level ',k,': ',maxval(ipvxy(:,:,k))
    enddo

    ! compute vorticity and divergence
    call vortdiv(u,v,vort,div)
    if (verbose .gt. 1) print *,'max surface vorticity, divergence:',maxval(vort(:,:,2)),maxval(div(:,:,2))

    call dump(u,v,theta,phi,w,pvsp,ub,vb,ut,vt,tbsp,ttsp,phib,phit,itim,outfile)

enddo

END SUBROUTINE pvinv_main

SUBROUTINE invert(tbsp,ttsp,pvsp,phi0b,phi0t,phi0)
! Invert PV for streamfunction; compute spectral derivatives on
! the transform (advection) grid.

    implicit none

    complex, dimension(2*kmax,2*lmax),      intent(in)  :: tbsp,ttsp
    complex, dimension(2*kmax,2*lmax,pmax), intent(in)  :: pvsp
    complex, dimension(2*kmax,2*lmax),      intent(out) :: phi0b,phi0t
    complex, dimension(2*kmax,2*lmax,pmax), intent(out) :: phi0

    complex, dimension(2*kmax,2*lmax)        :: pxx,pyy,pzz,errsp
    complex, dimension(2*kmax,2*lmax,pmax)   :: bpvsp,lap
    complex, dimension(2*kmax,2*lmax,pmax-1) :: pz

    real :: ak,bl,dz
    integer :: j,k,l,kk,ll

    real,    dimension(pmax) :: facz,faczo,zl,zhl
    complex, dimension(pmax) :: psi
    real,    dimension(2*kmax,2*lmax)       :: errxy
    complex, dimension(2*kmax,2*lmax), save :: dx,dy,Id

    dz = ZH/real(pmax)

    ! derivative operators
    call d_setup(dx,dy,Id)

    ! factors for tanh vertical grid
    call setup_tanh(facz,faczo,zl,zhl)

    ! initialize arrays to zero:
    phi0b=0.;phi0t=0.;phi0=0.

    bpvsp = pvsp
    !$omp parallel do schedule(dynamic) private(k,l,ak,bl)
    do k=1,2*kmax; do l=1,2*lmax

        !     add boundary theta to a copy of spectral pv:
         bpvsp(k,l,1) = bpvsp(k,l,1) + (tbsp(k,l)/(facz(1)*dz))
         bpvsp(k,l,pmax) = bpvsp(k,l,pmax) - (ttsp(k,l)/(facz(pmax)*dz))

         !     get wavenumbers
         call get_waves(k,l,ak,bl)

         !     spectral inversion (zero mean)
         if (k .eq. 1 .and. l .eq. 1) then
            phi0b(k,l) = 0.; phi0t(k,l) = 0.; phi0(k,l,:) = 0.
         else
            call matinv(bpvsp(k,l,:),ak,bl,psi,1)
            phi0b(k,l) = psi(1) - (0.5*dz*facz(1)*tbsp(k,l))
            phi0t(k,l) = psi(pmax) + (0.5*dz*facz(pmax)*ttsp(k,l))
            phi0(k,l,:) = psi(:)
         endif

    enddo; enddo
    !$omp end parallel do

    ! BEGIN leading-order PV check
    ! phi_xx + phi_yy + phi_zz = q

    call laplacian(phi0,tbsp,ttsp,lap)

    ! loop and print our errors
    do k = 1,pmax

        errsp = pvsp(:,:,k) - lap(:,:,k)
        call sp_to_xy(errsp,errxy,kmax,lmax,2*kmax,2*lmax)
        if (verbose .gt. 1)  print*,'max leading-order solution error at level ',k,' : ',maxval(abs(errxy))

        call sp_to_xy(lap(:,:,k),errxy,kmax,lmax,2*kmax,2*lmax)
        if (verbose .gt. 1) print*,'leading-order lap level ',k,' : ',maxval(abs(errxy))

    enddo
    ! END leading-order PV check

    return
END SUBROUTINE invert

SUBROUTINE get_waves(k,l,ak,bl)
! compute x and y wavenumbers for use in spectral calculations.

    implicit none

    integer, intent(in)  :: k,l
    real,    intent(out) :: ak, bl

    ak = facx*real(k - 1); bl = facy*real(l - 1)

    ! other spectral quadrants
    if (k .ge. kmax .and. l .le. lmax) then
        ak = -1.*facx*real(2*kmax - k + 1)
    elseif (l .ge. lmax .and. k .le. kmax) then
        bl = -1.*facy*real(2*lmax -l + 1)
    elseif (k .ge. kmax .and. l .ge. lmax) then
        ak = -1.*facx*real(2*kmax - k + 1)
        bl = -1.*facy*real(2*lmax - l + 1)
    endif

    return
END SUBROUTINE get_waves

SUBROUTINE d_setup(dx,dy,Id)
! Set up matrices for derivatives and integrals.

    implicit none

    complex, dimension(2*kmax,2*lmax), intent(out) :: dx,dy,Id
    real    :: m
    integer :: k,l

    ! dx
    dx(1,:) = 0.; dx(kmax+1,:) = 0.
    do k=2,kmax
        dx(k,:) = facx*real(k-1)*cmplx(0.,1.)
        dx(2*kmax-k+2,:) = -1.*dx(k,:)
    enddo

    ! dy
    dy(:,1) = 0.; dy(:,lmax+1) = 0.
    do l=2,lmax
        dy(:,l) = facy*real(l-1)*cmplx(0.,1.)
        dy(:,2*lmax-l+2) = -1.*dy(:,l)
    enddo

    ! Id
    Id = 1.0

    return
END SUBROUTINE d_setup

SUBROUTINE d_s2s(temp_in,temp_out,dflag,dn)
! small array to small array.
! dflag =  n: n derivatives. dflag = -n: n integrations.

    implicit none

    complex, dimension(:,:), intent(in)  :: temp_in,dn
    integer,                 intent(in)  :: dflag
    complex, dimension(:,:), intent(out) :: temp_out

    integer :: k,l

    temp_out = 0.

    !$omp parallel do schedule(dynamic) private(k,l)
    do k = 1,2*kmax; do l = 1,2*lmax
        if (dn(k,l) .ne. 0) temp_out(k,l) = (dn(k,l)**dflag)*temp_in(k,l)
    enddo; enddo
    !$omp end parallel do

    return
END SUBROUTINE d_s2s

SUBROUTINE d_s2b(temp_in,temp_out,dflag,dn)
! small array to big array.
! dflag =  n: n derivatives. dflag = -n: n integrations.

    implicit none

    complex, dimension(:,:), intent(in)  :: temp_in,dn
    integer,                 intent(in)  :: dflag
    complex, dimension(:,:), intent(out) :: temp_out

    integer :: k,l,kk,ll

    temp_out = 0.

    !$omp parallel do schedule(dynamic) private(k,l)
    do k = 1,kmax; do l = 1,lmax

        ! waves: 0 <= k <= +/-kmax; 0 <= l <= +/-lmax:
        temp_out(k,l) = (dn(k,l)**dflag)*temp_in(k,l)

        kk = kmax + k; ll = lmax + l

        ! +/-kmax <= k <= -1; 0 <= l <= +/-lmax:
        if (k .gt. 1) temp_out(k2+k,l) = (dn(kk,l)**dflag)*temp_in(kk,l)

        ! 0 <= k <= +/-kmax; +/-lmax <= l <= -1:
        if (l .gt. 1) temp_out(k,l2+l) = (dn(k,ll)**dflag)*temp_in(k,ll)

        ! +/-kmax <= k <= -1; +/-lmax <= l <= -1:
        if ((k .gt. 1) .and. (l .gt. 1)) temp_out(k2+k,l2+l) = (dn(kk,ll)**dflag)*temp_in(kk,ll)

    enddo; enddo
    !$omp end parallel do

    return
END SUBROUTINE d_s2b

SUBROUTINE matinv(pv,ak,bl,psi,idn)
! Solve for psi(z) given pv(z) for a given (k,l) in ** uniform ** N^2:

    implicit none

    complex, dimension(:), intent(in)  :: pv
    real,                  intent(in)  :: ak,bl
    integer,               intent(in)  :: idn
    complex, dimension(:), intent(out) :: psi

    complex, dimension(pmax) :: f
    real,    dimension(pmax) :: e

    real    :: A,B,C,dz
    integer :: k

    psi = 0.

    ! set up coeficients A, B, and C (note they are set for uniform n^2!!!):
    dz = ZH/real(pmax)
    A = -1./(dz*dz); B = -1.*((ak*ak) + (bl*bl) + (2./(dz*dz)))
    C = -1./(dz*dz)

    ! idn determines inversion BCs (1.=Neumann, -1.=Dirichlet)
    !      idn = 1

    ! first pass:
    e(1) = A / (B - (C*real(idn)))
    f(1) = pv(1) / (B-(C*real(idn)))

    !$omp parallel do schedule(dynamic) private(k)
    do k = 2, pmax
        e(k) = A / (B - (C*e(k-1)))
        f(k) = (pv(k) + (C*f(k-1))) / (B - (C*e(k-1)))
        !print*,'deep check=',j,f(j),pmax
    enddo
    !$omp end parallel do

    !print*,'f = ',f(pmax)
    !second pass
    if(e(pmax).eq.idn) then
        psi(pmax) = 0.0
    else
        psi(pmax) = f(pmax) / (1. - (real(idn)*e(pmax)))
        !print*,'deep check=',f(pmax),real(idn),e(pmax)
    endif

    do k = 1, pmax-1
        psi(pmax-k) = (e(pmax-k)*psi(pmax+1-k)) + f(pmax-k)
    enddo

    return
END SUBROUTINE matinv
!========================================================================

!========================================================================
SUBROUTINE xy_to_sp(xy,sp,mx,ny,km,lm)
! Map an (x,y) array onto a _smaller_ spectral array.
! Input: xy(mx,ny) --- a grid point array.
! Output: sp(2*km,2*lm) --- a spectral array.

    implicit none

    integer,                       intent(in)  :: km,lm,mx,ny
    complex, dimension(:,:),       intent(in)  :: xy
    complex, dimension(2*km,2*lm), intent(out) :: sp

    complex, dimension(mx,ny) :: copy
    integer                   :: kmp1,lmp1,k,l,k2,l2,kk,ll

    ! initialize arrays:
    sp = 0.0 ; copy = xy

    kmp1 = km + 1; lmp1 = lm + 1; k2 = mx - km; l2 = ny - lm

    call ft_2d(copy,mx,ny,-1)

    !$omp parallel do schedule(dynamic) private(k,l,kk,ll)
    !do k=1,kmp1; do l=1,lmp1
    do k=1,km; do l=1,lm

        kk = km + k; ll = lm + l

        ! waves: 0 <= k <= +/-kmax; 0 <= l <= +/-lmax:
        sp(k,l) = copy(k,l)/real(mx*ny)

        ! +/-kmax <= k <= -1; 0 <= l <= +/-lmax:
        if (k .gt. 1) then
            sp(kk,l) = copy(k2+k,l)/real(mx*ny)
        endif

        ! 0 <= k <= +/-kmax; +/-lmax <= l <= -1:
        if (l .gt. 1) then
            sp(k,ll) = copy(k,l2+l)/real(mx*ny)
        endif

        ! +/-kmax <= k <= -1; +/-lmax <= l <= -1:
        if ((k .gt. 1) .and. (l .gt. 1)) then
            sp(kk,ll) = copy(k2+k,l2+l)/real(mx*ny)
        endif

    enddo; enddo
    !$omp end parallel do

    return
END SUBROUTINE xy_to_sp
!========================================================================

!========================================================================
SUBROUTINE sp_to_xy(sp,xy,km,lm,mx,ny)
! Map an (km,lm) spectral array onto a _bigger_ grid point array.
! Input: sp(2*km,2*lm) --- a spectral array.
! Output: xy(mx,ny) --- a grid point array.

    implicit none

    integer,                       intent(in)  :: km,lm,mx,ny
    complex, dimension(:,:),       intent(in)  :: sp
    real,    dimension(mx,ny),     intent(out) :: xy

    complex :: copy(mx,ny)
    integer :: kmp1,lmp1,k,l,k2,l2,kk,ll

    copy = 0.0

    kmp1 = km + 1; lmp1 = lm + 1; k2 = mx - km; l2 = ny - lm

    !$omp parallel do schedule(dynamic) private(k,l,kk,ll)
    !do 10 k = 1,kmp1; do 10 l = 1,lmp1
    do k = 1,km ; do l = 1,lm

        ! waves: 0 <= k <= +/-kmax; 0 <= l <= +/-lmax:
        copy(k,l) = sp(k,l)

        kk = km + k; ll = lm + l

        ! +/-kmax <= k <= -1; 0 <= l <= +/-lmax:
        if (k .ge. 1) then
           copy(k2+k,l) = sp(kk,l)
        endif

        ! 0 <= k <= +/-kmax; +/-lmax <= l <= -1:
        if (l .ge. 1) then
           copy(k,l2+l) = sp(k,ll)
        endif

        ! +/-kmax <= k <= -1; +/-lmax <= l <= -1:
        if ((k .ge. 1) .and. (l .ge. 1)) then
           copy(k2+k,l2+l) = sp(kk,ll)
        endif

     enddo ; enddo
    !$omp end parallel do

    call ft_2d(copy,mx,ny,1)
    xy = real(copy)

    return
END SUBROUTINE sp_to_xy
!========================================================================

!========================================================================
SUBROUTINE init(infile,itime,tbxy,ttxy,pvxy)
! Initialize pv and theta fields

    implicit none

    character(len=*),                    intent(in)  :: infile
    integer,                             intent(in)  :: itime
    real, dimension(2*kmax,2*lmax),      intent(out) :: tbxy,ttxy
    real, dimension(2*kmax,2*lmax,pmax), intent(out) :: pvxy

    real, dimension(2*kmax,2*lmax,2) :: theta

    real :: x,y,z,dx,dy,dz
    real :: xcen,ycen,amdx,amdy,rr,rrz
    real :: xn
    integer :: i,j,k
    integer :: iseed
    integer :: ncid, varid, start(4), count(4)

    real, parameter :: pva = 0.0
    real, parameter :: amp = -2.0, asig = 0.5
    real, parameter :: asx = 1.0,  asy  = 1.0, asz = 0.1
    real, parameter :: zlev = 0.5

    if (itime == 0) then ! make IC

        dx = XL / real(2*kmax)
        dy = YL / real(2*lmax)
        dz = ZH / real(pmax)

        do i=1,2*kmax; do j=1,2*lmax

            x = (i-1)*dx; y = (j-1)*dy

            xcen = kmax*dx; ycen = lmax*dy
            amdx=min( abs(x-xcen),abs(XL+x-xcen),abs(XL-x+xcen) )
            amdy=min( abs(y-ycen),abs(YL+y-ycen),abs(YL-y+ycen) )

            rr = (((amdx/asx)**2) + ((amdy/asy)**2))**0.5

            ! gaussian:
            tbxy(i,j) = 0.
            ttxy(i,j) = amp*exp(-1.0 * ((rr/asig)**2))

            ! second deriv of gaussian:
            !tbxy(i,j) = 0.
            !ttxy(i,j) = -amp*(1.-(rr*rr))*exp(-1.*((rr/asig)**2))

            ! random:
            !tbxy(i,j) = amp*(ran1(iseed)-0.5)
            !ttxy(i,j) = amp*(ran1(iseed)-0.5)

            ! plane waves
            !tbxy(i,j) = 0.
            !ttxy(i,j) = amp*cos((2.*pi*x/XL) ) ! + (2.*pi*y/YL))

            do k = 1, pmax

                Z = (real(k) - 0.5)*dz

                xn = xcen !- (z - 1.) ! this term adds vertical tilt

                amdx=min( abs(x-xn),abs(XL+x-xn),abs(XL-x+xn) )
                rr = (((amdx/asx)**2) + ((amdy/asy)**2))**0.5
                rrz = abs(zlev-z)
                pvxy(i,j,k) = pva * (1.0 - (rr*rr))        * &
                exp(-1.0 * ((rr/asig)**2)) * &
                exp(-1.0 * ((rrz/asz)**2))

                if ( ( i.eq.kmax) .and. (j.eq.lmax) ) then
                if (verbose .gt. 1) print*,'pv: ',k,pvxy(i,j,k)
                endif

            enddo

        enddo; enddo

    else ! read from a file

        tbxy = 0.0 ; ttxy = 0.0 ; pvxy = 0.0

        start(1) = 1     ; count(1) = 2*kmax
        start(2) = 1     ; count(2) = 2*lmax
        start(3) = 1     ; count(3) = 2
        start(4) = itime ; count(4) = 1

        call nc_check( nf90_open(trim(infile), NF90_NOWRITE, ncid), 'init', 'open, ' // trim(infile) )
        call nc_check( nf90_inq_varid(ncid, 'theta', varid), 'init', 'inq_varid theta, ' // trim(infile) )
        call nc_check( nf90_get_var(ncid, varid, theta, start, count), 'init', 'get_var theta, ' // trim(infile) )
        call nc_check( nf90_close(ncid), 'init', 'close, ' // trim(infile) )

        tbxy = theta(:,:,1) ; ttxy = theta(:,:,2)

    endif

    return
END SUBROUTINE init
!========================================================================

!========================================================================
FUNCTION ran1(idum)
! Numerical Recipes random number generator:

    implicit none

    REAL    ran1
    INTEGER idum

    INTEGER IA,IM,IQ,IR,NTAB,NDIV
    REAL    AM,EPS,RNMX
    PARAMETER (IA=16807,IM=2147483647,AM=1./IM,IQ=127773,IR=2836, &
               NTAB=32,NDIV=1+(IM-1)/NTAB,EPS=1.2e-7,RNMX=1.-EPS)
    INTEGER j,k,iv(NTAB),iy
    SAVE iv,iy
    DATA iv /NTAB*0/, iy /0/
    if (idum.le.0.or.iy.eq.0) then
        idum=max(-idum,1)
        do j=NTAB+8,1,-1
        k=idum/IQ
        idum=IA*(idum-k*IQ)-IR*k
        if (idum.lt.0) idum=idum+IM
        if (j.le.NTAB) iv(j)=idum
        enddo
        iy=iv(1)
    end if
    k=idum/IQ
    idum=IA*(idum-k*IQ)-IR*k
    if (idum.lt.0) idum=idum+IM
    j=1+iy/NDIV
    iy=iv(j)
    iv(j)=idum
    ran1=min(AM*iy,RNMX)

    return
END FUNCTION ran1
!========================================================================

!========================================================================
SUBROUTINE dump(u,v,theta,phi,w,pv,ub,vb,ut,vt,tb,tt,phib,phit,it,outfile)
! Write to disk.

    implicit none

    complex, dimension(2*kmax,2*lmax,pmax), intent(in) :: u,v,theta,phi,w,pv
    complex, dimension(2*kmax,2*lmax),      intent(in) :: ub,vb,ut,vt,tb,tt,phib,phit
    integer,                                intent(in) :: it
    character(len=*),                       intent(in) :: outfile

    real, dimension(pmax)               :: pheight
    real, dimension(2*kmax,2*lmax,pmax) :: otheta,ophi,ou,ov,ow,opv
    real, dimension(2*kmax,2*lmax,2)    :: otheta_s,ophi_s,ou_s,ov_s
    integer :: k
    real :: dz
    real :: grav,tnot,Rhos,Ns,km,Cor,Ls,Hs,Us,Ws,Ps,Ts,PVs

    call scaling(grav,tnot,Rhos,Ns,km,Cor,Ls,Hs,Us,Ws,Ps,Ts,PVs)

    dz = ZH/real(pmax)
    if (verbose .gt. 1) then
        print*,'Scaling parameters...'
        print*,'grav = ',grav,' m s^-2'
        print*,'Cor  = ',Cor,' s^-1'
        print*,'Ls   = ',Ls, ' m'
        print*,'Hs   = ',Hs, ' m'
        print*,'Us   = ',Us, ' m s^-1'
        print*,'Ws   = ',Ws*100, ' cm s^-1'
        print*,'Ps   = ',Ps, ' m^2 s^-2'
        print*,'Ts   = ',Ts, ' K'
        print*,'Ns   = ',Ns, ' s^-1'
        print*,'PVs  = ',PVs, ' PVU'
    endif

    ! lower boundary theta
    call sp_to_xy(tb,otheta_s(:,:,1),kmax,lmax,2*kmax,2*lmax); otheta_s(:,:,1) = Ts*otheta_s(:,:,1)
    if (verbose .gt. 1) print*,'max z=0 dim |t| = ',maxval(abs(otheta_s(:,:,1)))

    ! upper boundary theta
    call sp_to_xy(tt,otheta_s(:,:,2),kmax,lmax,2*kmax,2*lmax); otheta_s(:,:,2) = Ts*otheta_s(:,:,2)
    if (verbose .gt. 1) print*,'max z=H dim |t| = ',maxval(abs(otheta_s(:,:,2)))

    ! lower boundary geopotential
    call sp_to_xy(phib,ophi_s(:,:,1),kmax,lmax,2*kmax,2*lmax); ophi_s(:,:,1) = Ps*ophi_s(:,:,1)
    if (verbose .gt. 1) print*,'max z=0 dim |p| = ',maxval(abs(ophi_s(:,:,1)))

    ! upper boundary geopotential
    call sp_to_xy(phit,ophi_s(:,:,2),kmax,lmax,2*kmax,2*lmax); ophi_s(:,:,2) = Ps*ophi_s(:,:,2)
    if (verbose .gt. 1) print*,'max z=0 dim |p| = ',maxval(abs(ophi_s(:,:,2)))

    ! lower boundary u & v
    call sp_to_xy(ub,ou_s(:,:,1),kmax,lmax,2*kmax,2*lmax); ou_s(:,:,1) = Us*ou_s(:,:,1)
    if (verbose .gt. 1) print*,'max z=0 dim |u| = ',maxval(abs(ou_s(:,:,1)))
    call sp_to_xy(vb,ov_s(:,:,1),kmax,lmax,2*kmax,2*lmax); ov_s(:,:,1) = Us*ov_s(:,:,1)
    if (verbose .gt. 1) print*,'max z=0 dim |v| = ',maxval(abs(ov_s(:,:,1)))

    ! upper boundary u & v
    call sp_to_xy(ut,ou_s(:,:,2),kmax,lmax,2*kmax,2*lmax); ou_s(:,:,2) = Us*ou_s(:,:,2)
    if (verbose .gt. 1) print*,'max z=H dim |u| = ',maxval(abs(ou_s(:,:,2)))
    call sp_to_xy(vt,ov_s(:,:,2),kmax,lmax,2*kmax,2*lmax); ov_s(:,:,2) = Us*ov_s(:,:,2)
    if (verbose .gt. 1) print*,'max z=H dim |v| = ',maxval(abs(ov_s(:,:,2)))

    ! interior theta, geopotential, u, v, w and pv
    do k = 1, pmax

        if (verbose .gt. 1) print*,'level ',k

        call sp_to_xy(theta(:,:,k),otheta(:,:,k),kmax,lmax,2*kmax,2*lmax)
        otheta(:,:,k) = Ts*otheta(:,:,k)
        if (verbose .gt. 1) print*,'max dim |t| = ',maxval(abs(otheta(:,:,k)))

        call sp_to_xy(phi(:,:,k),ophi(:,:,k),kmax,lmax,2*kmax,2*lmax)
        ophi(:,:,k) = Ps*ophi(:,:,k)
        if (verbose .gt. 1) print*,'max dim |p| = ',maxval(abs(ophi(:,:,k)))

        call sp_to_xy(u(:,:,k),ou(:,:,k),kmax,lmax,2*kmax,2*lmax)
        ou(:,:,k) = Us*ou(:,:,k)
        if (verbose .gt. 1) print*,'max dim |u| = ',maxval(abs(ou(:,:,k)))

        call sp_to_xy(v(:,:,k),ov(:,:,k),kmax,lmax,2*kmax,2*lmax)
        ov(:,:,k) = Us*ov(:,:,k)
        if (verbose .gt. 1) print*,'max dim |v| = ',maxval(abs(ov(:,:,k)))

        call sp_to_xy(w(:,:,k),ow(:,:,k),kmax,lmax,2*kmax,2*lmax)
        ow(:,:,k) = Ws*ow(:,:,k)
        if (verbose .gt. 1) print*,'max dim |w| = ',maxval(abs(ow(:,:,k)))

        call sp_to_xy(pv(:,:,k),opv(:,:,k),kmax,lmax,2*kmax,2*lmax);
        opv(:,:,k) = PVs*opv(:,:,k)
        if (verbose .gt. 1) print*,'max dim |pv| = ',maxval(abs(opv(:,:,k)))

        pheight(k) = Hs*(real(k) - 0.5)*dz

    enddo

    if (verbose .gt. 0) print*,'Writing to disk...'
    call write_diag(outfile,it,ou_s,ov_s,otheta_s,ophi_s,ou,ov,otheta,ophi,ow,opv)

    return
END SUBROUTINE dump
!========================================================================

!========================================================================
SUBROUTINE ft_2d(f,ni,nj,isign)
! FFT-calling SUBROUTINE.

    implicit none

    integer,                   intent(in)    :: ni, nj, isign
    complex, dimension(ni,nj), intent(inout) :: f

    real, dimension(ni,nj) :: re, im

    re = real(f); im = aimag(f)
    call fft(re,im,ni*nj,ni,ni,   isign)
    call fft(re,im,ni*nj,nj,nj*ni,isign)
    f = cmplx(re,im)

    return
END SUBROUTINE ft_2d
!========================================================================

!========================================================================
SUBROUTINE invert_old(tbsp,ttsp,pvsp,phi0b,phi0t,phi0)
! Invert PV for streamfunction; compute spectral derivatives on
! the transform (advection) grid.
! this is the old inversion routine. the new one has the quadrant work
! hidden in operators, so it scales better for +1 code.

    implicit none

    complex, dimension(2*kmax,2*lmax),      intent(in)  :: tbsp,ttsp
    complex, dimension(2*kmax,2*lmax,pmax), intent(in)  :: pvsp
    complex, dimension(2*kmax,2*lmax),      intent(out) :: phi0b,phi0t
    complex, dimension(2*kmax,2*lmax,pmax), intent(out) :: phi0

    complex, dimension(2*kmax,2*lmax,pmax) :: bpvsp
    complex, dimension(pmax)               :: psi
    real :: ak,bl,kap,sqg,amod,dz
    integer :: k,l,kk,ll

    dz = ZH/real(pmax)

    ! initialize arrays to zero:
    phi0b=0.;phi0t=0.;phi0=0.

    ! add boundary theta to a copy of spectral pv:
    bpvsp = pvsp
    !$omp parallel do schedule(dynamic) private(k,l)
    do k=1,2*kmax; do l=1,2*lmax
        bpvsp(k,l,1) = bpvsp(k,l,1) + (tbsp(k,l)/dz)
        bpvsp(k,l,pmax) = bpvsp(k,l,pmax) - (ttsp(k,l)/dz)
    enddo; enddo
    !$omp end parallel do

    ! inversion of (kmax,lmax) waves into (mmax,nmax) arrays:
    !$omp parallel do schedule(dynamic) private(k,l,ak,bl,kk,ll,psi)
    do k = 1,kmaxp1 ; do l = 1,lmaxp1

        ! waves: 0 <= k <= +/-kmax; 0 <= l <= +/-lmax:
        ak = facx*real(k - 1); bl = facy*real(l - 1)

        if (k .eq. 1 .and. l .eq. 1) then
            phi0(k,l,:) = 0.; phi0(k,l,:) = 0.
        else
            call matinv(bpvsp(k,l,:),ak,bl,psi,1)
            phi0b(k,l) = psi(1) - (0.5*dz*tbsp(k,l))
            phi0t(k,l) = psi(pmax) + (0.5*dz*ttsp(k,l))
            phi0(k,l,:) = psi
        endif

        kk = kmax + k; ll = lmax + l

        ! +/-kmax <= k <= -1; 0 <= l <= +/-lmax:
        if (k .le. kmax) then
            ak = -1.*facx*real(kmaxp1 - k); bl = facy*real(l - 1)
            call matinv(bpvsp(kk,l,:),ak,bl,psi,1)
            phi0b(kk,l) = psi(1) - (0.5*dz*tbsp(kk,l))
            phi0t(kk,l) = psi(pmax) + (0.5*dz*ttsp(kk,l))
            phi0(kk,l,:) = psi
        endif

        ! 0 <= k <= +/-kmax; +/-lmax <= l <= -1:
        if (l .le. lmax) then
            ak = facx*real(k-1); bl = -1.*facy*real(lmaxp1 - l)
            call matinv(bpvsp(k,ll,:),ak,bl,psi,1)
            phi0b(k,ll) = psi(1) - (0.5*dz*tbsp(k,ll))
            phi0t(k,ll) = psi(pmax) + (0.5*dz*ttsp(k,ll))
            phi0(k,ll,:) = psi
        endif

        ! +/-kmax <= k <= -1; +/-lmax <= l <= -1:
        if (k .le. kmax .and. l .le. lmax) then
            ak=-1.*facx*real(kmaxp1 - k); bl=-1.*facy*real(lmaxp1 - l)
            call matinv(bpvsp(kk,ll,:),ak,bl,psi,1)
            phi0b(kk,ll) = psi(1) - (0.5*dz*tbsp(kk,ll))
            phi0t(kk,ll) = psi(pmax) + (0.5*dz*ttsp(kk,ll))
            phi0(kk,ll,:) = psi
        endif

    enddo ; enddo
    !$omp end parallel do

    return
END SUBROUTINE invert_old
!========================================================================

!========================================================================
SUBROUTINE uvwtp(tb,tt,phi0,phi0b,phi0t,F1,G1,P1,phi1, &
                 u,ub,ut,v,vb,vt,w,thF)
! recover winds, theta, and geopotential from potential data.

    implicit none

    complex, dimension(2*kmax,2*lmax,pmax), intent(in)  :: phi0,phi1
    complex, dimension(2*kmax,2*lmax,pmax), intent(in)  :: F1,G1,P1
    complex, dimension(2*kmax,2*lmax),      intent(in)  :: phi0b,phi0t
    complex, dimension(2*kmax,2*lmax),      intent(in)  :: tb,tt
    complex, dimension(2*kmax,2*lmax,pmax), intent(out) :: u,v,w
    complex, dimension(2*kmax,2*lmax,pmax), intent(out) :: thF
    complex, dimension(2*kmax,2*lmax),      intent(out) :: ub,ut,vb,vt

    complex, dimension(2*kmax,2*lmax)      :: dx,dy,Id,uF,vF,temp,ftemp,oftemp
    complex, dimension(2*kmax,2*lmax,pmax) :: u1,v1
    complex, dimension(2*kmax,2*lmax,pmax) :: phi
    complex, dimension(2*kmax,2*lmax)      :: ug,vg
    real,    dimension(2*kmax,2*lmax)      :: rtemp,rutemp,rvtemp,ws,wsg
    real    :: dz
    integer :: k,l
    complex :: gtemp,ogtemp

    dz = ZH/real(pmax)

    ! derivative operators
    call d_setup(dx,dy,Id)

    ! 3D leading-order winds
    if (order .eq. 0 .or. order .eq. 2) then

        !$omp parallel do schedule(dynamic) private(k,uF,vF)
        do k = 1,pmax
            call d_s2s(phi0(:,:,k),uF,1,-dy); u(:,:,k) = uF
            call d_s2s(phi0(:,:,k),vF,1, dx); v(:,:,k) = vF
        enddo
        !$omp end parallel do

        ! leading-order boundary winds
        call d_s2s(phi0b,uF,1,-dy); ub = uF
        call d_s2s(phi0t,uF,1,-dy); ut = uF

        call d_s2s(phi0b,vF,1,dx); vb = vF
        call d_s2s(phi0t,vF,1,dx); vt = vF

    else

        u = 0.; v = 0.

    endif ! order 0

    if (order .eq. 1 .or. order .eq. 2 .and. Ross .ne. 0.) then

        if (verbose .gt. 0) print*,'computing next-order corrections to u and v...'

         ! next-order corrections
         ! v1 =  phi1_x - G1_z
         ! u1 = -phi1_y - F1_z
         ! w  =  F1_x + G1_y

         ! compute vertical derivatives by finite differences
        !$omp parallel do schedule(dynamic) private(k)
        do k=1,pmax
            if (k .eq. 1) then ! linear interp; zero BC
                u1(:,:,1) = -(F1(:,:,1) + F1(:,:,2))/(2*dz)
                v1(:,:,1) = -(G1(:,:,1) + G1(:,:,2))/(2*dz)
                ub(:,:) = ub(:,:) - Ross*(2*F1(:,:,1)/dz)
                vb(:,:) = vb(:,:) - Ross*(2*G1(:,:,1)/dz)
            elseif (k .eq. pmax) then ! linear interp; zero BC
                u1(:,:,pmax) = (F1(:,:,pmax) + F1(:,:,pmax-1))/(2*dz)
                v1(:,:,pmax) = (G1(:,:,pmax) + G1(:,:,pmax-1))/(2*dz)
                ut(:,:) = ut(:,:) + Ross*(2*F1(:,:,pmax)/dz)
                vt(:,:) = vt(:,:) + Ross*(2*G1(:,:,pmax)/dz)
            else
                u1(:,:,k) = -(F1(:,:,k+1) - F1(:,:,k-1))/(2*dz)
                v1(:,:,k) = -(G1(:,:,k+1) - G1(:,:,k-1))/(2*dz)
            endif
        enddo ! k
        !$omp end parallel do

        ! u1 = 0; v1 = 0 ! test of phi contrib
        ! compute horizonal derivatives spectrally
        !$omp parallel do schedule(dynamic) private(k,temp)
        do k = 1, pmax
            call d_s2s(P1(:,:,k),temp,1, dx); v1(:,:,k) = v1(:,:,k) + temp
            call d_s2s(P1(:,:,k),temp,1,-dy); u1(:,:,k) = u1(:,:,k) + temp
            if (k .eq. 1) then ! recall phi1_z = 0
                call d_s2s(P1(:,:,k),temp,1, dx); vb = vb + Ross*temp
                call d_s2s(P1(:,:,k),temp,1,-dy); ub = ub + Ross*temp
            endif
            if (k .eq. pmax) then ! recall phi1_z = 0
                call d_s2s(P1(:,:,k),temp,1, dx); vt = vt + Ross*temp
                call d_s2s(P1(:,:,k),temp,1,-dy); ut = ut + Ross*temp
            endif
            call d_s2s(F1(:,:,k),w(:,:,k),1,dx)
            call d_s2s(G1(:,:,k),temp,    1,dy)
            w(:,:,k) = Ross*(w(:,:,k) + temp)
        enddo
        !$omp end parallel do

    else

        u1 = 0.; v1 = 0.; w = 0.

    endif ! order 1

    u = u + Ross*u1
    v = v + Ross*v1

    ! OLD
    ! compute theta and get it on the regular, staggered vertical grid
    !if (order .eq. 0) then
    !   phi = phi0
    !else
    !   phi = phi0 + (Ross*phi1) ! total geopotential
    !endif
    !thF = 0.
    ! compute vertical derivatives by finite differences
    ! on staggered grid starting dz off boundary
    !do p=1,pmax-1
    !   thF(:,:,p) = (phi(:,:,p+1) - phi(:,:,p))/dz
    !enddo ! p

    ! NEW---try computing theta from potential functions
    ! theta = Phi_z + G_x - F_y
    if (1 .eq. 1) then
        thF = 0.
        if (order .eq. 0) then
            phi = phi0             ! total phi potential -- different from above!
        else
            phi = phi0 + (Ross*P1) ! total phi potential -- different from above!
        endif
        !$omp parallel do schedule(dynamic) private(k,temp)
        do k=1,pmax
            !P1_z:
            if (k .eq. 1) then
                thF(:,:,k) = (phi(:,:,k+1) - (phi(:,:,k)-tb*dz))/(2.*dz)
            elseif (k .eq. pmax) then
                thF(:,:,k) = ((phi(:,:,k)+tt*dz) - phi(:,:,k-1))/(2.*dz)
            else
                thF(:,:,k) = (phi(:,:,k+1) - phi(:,:,k-1))/(2.*dz)
            endif
            if (order .ne. 0) then
                !G_x:
                call d_s2s(G1(:,:,k),temp,1,dx); thF(:,:,k) = thF(:,:,k) + Ross*temp
                !F_y
                call d_s2s(F1(:,:,k),temp,1,dy); thF(:,:,k) = thF(:,:,k) - Ross*temp
            endif
        enddo ! k
        !$omp end parallel do
    endif

    ! NEW NEW NEW--theta at half-levels
    if (1 .eq. 0) then
        thF = 0.
        phi = phi0
        !$omp parallel do schedule(dynamic) private(k,temp,ftemp)
        do k=1,pmax-1
            thF(:,:,k) = (phi(:,:,k+1) - phi(:,:,k))/(dz)
            thF(:,:,k) = thF(:,:,k) + Ross*((P1(:,:,k+1) - P1(:,:,k))/(dz))
            if (order .ne. 0) then
                !G_x:
                call d_s2s(G1(:,:,k),temp,1,dx); call d_s2s(G1(:,:,k+1),ftemp,1,dx)
                thF(:,:,k) = thF(:,:,k) + Ross*(temp + ftemp)/2.
                !F_y
                call d_s2s(F1(:,:,k),temp,1,dy); call d_s2s(F1(:,:,k+1),ftemp,1,dy);
                thF(:,:,k) = thF(:,:,k) - Ross*(temp + ftemp)/2.
            endif
        enddo ! k
        !$omp end parallel do
    endif

    ! compute the geostrophic wind from the _total_ geopotential
    do k = 1, pmax
        call d_s2s(phi(:,:,k),vg,1, dx)
        call d_s2s(phi(:,:,k),ug,1,-dy)
        ! move geostrophic wind onto physical grid
        rutemp = 0.; rvtemp = 0.
        call sp_to_xy(vg,rvtemp,kmax,lmax,2*kmax,2*lmax)
        call sp_to_xy(ug,rutemp,kmax,lmax,2*kmax,2*lmax)
        !call d_s2s(vg,rvtemp,1,Id); call d_s2s(ug,rutemp,1,Id)
        wsg = sqrt(((rutemp**2) + (rvtemp**2)))
        ! move full wind onto physical grid
        rutemp = 0.; rvtemp = 0.
        call sp_to_xy(v(:,:,k),rvtemp,kmax,lmax,2*kmax,2*lmax)
        call sp_to_xy(u(:,:,k),rutemp,kmax,lmax,2*kmax,2*lmax)
        !call d_s2s(v(:,:,k),rvtemp,1,Id); call d_s2s(u(:,:,k),rutemp,1,Id)
        ws = sqrt(((rutemp**2) + (rvtemp**2)))
        if (verbose .gt. 1) print*,'max wind speed for full and geostrophic winds (total phi) at level: ', k, maxval(ws),maxval(wsg)
    enddo

    return
END SUBROUTINE uvwtp
!========================================================================

!========================================================================
SUBROUTINE invert_R(tbsp,ttsp,phi0b,phi0t,phi0,F1,G1,P1,phi1,phi1b,phi1t,pvsp)
! Invert for O(R) potentials.
! phi0 : leading order geopotential.
! phi1 : O(Ross) geopotential correction.
! phi1b, phi1t : O(Ross) boundary geopotential corrections.
! F1, G1, P1 : 3D O(Ross) potentials.
! pz, pzz, pxy, etc. : derivatives of phi0.
! bpvsp : boundary condition trick (make homogeneous BCs; correct int).

    implicit none

    complex, dimension(2*kmax,2*lmax),      intent(in)  :: tbsp,ttsp,phi0b,phi0t
    complex, dimension(2*kmax,2*lmax,pmax), intent(in)  :: phi0,pvsp
    complex, dimension(2*kmax,2*lmax),      intent(out) :: phi1b,phi1t
    complex, dimension(2*kmax,2*lmax,pmax), intent(out) :: F1,G1,P1,phi1

    complex, dimension(2*kmax,2*lmax)      :: bcb,bct,errsp,tb,tt
    complex, dimension(2*kmax,2*lmax,pmax) :: F1s,G1s,P1s,phi1s,lap
    complex, dimension(2*kmax,2*lmax,pmax) :: pzS,pzzS

    real,    dimension(mmax,nmax) :: pzx,pzy,pzz,pxx,pyy,pxy,Rtemp,pvxy,tbx,tby,pzxold,pzyold
    complex, dimension(mmax,nmax) :: Ctemp
    complex, dimension(pmax)      :: psi

    real,    dimension(2*kmax,2*lmax) :: Ptemp,errxy
    complex, dimension(2*kmax,2*lmax) :: solvB,solvT

    complex,dimension(2*kmax,2*lmax), save :: dx,dy,Id

    ! filler for setup_tanh
    real, dimension(pmax) :: facz,faczo,zl,zhl

    real    :: ak,bl,dz,solv
    integer :: j,k,l,kk,ll,p

    dz = ZH/real(pmax)

    ! derivative operators
    call d_setup(dx,dy,Id)

    ! factors for tanh vertical grid
    call setup_tanh(facz,faczo,zl,zhl)

    tb = tbsp; tt = ttsp;
    tb(1,1) = 0.; tt(1,1) = 0.

    ! solvability condition [\int q dV = \int \theta^1 + \phi0_z(\phi0_xx + \phi0_yy)dxdy]_0^1
    ! \theta^1 = 0 ==> q = C/H where C = avg[\phi0_z(\phi0_xx + \phi0_yy)].
    ! bottom
    call d_s2b(phi0b,Ctemp,1,dx*dx); call ft_2d(Ctemp,mmax,nmax,1); pxx = real(Ctemp)
    call d_s2b(phi0b,Ctemp,1,dy*dy); call ft_2d(Ctemp,mmax,nmax,1); pyy = real(Ctemp)
    ! this is an estimate of pzz 1/8 dz off the surface; pre-use solvB
    solvB = (((phi0(:,:,1)-phi0b)*2./(facz(1)*dz)) - tb)*4./(facz(1)*dz)
    call d_s2b(solvB,Ctemp,1,Id); call ft_2d(Ctemp,mmax,nmax,1); pzz = real(Ctemp)
    Rtemp = pzz*(pxx + pyy); call xy_to_sp(cmplx(Rtemp,0.),solvB,mmax,nmax,kmax,lmax)
    ! top
    call d_s2b(phi0t,Ctemp,1,dx*dx); call ft_2d(Ctemp,mmax,nmax,1); pxx = real(Ctemp)
    call d_s2b(phi0t,Ctemp,1,dy*dy); call ft_2d(Ctemp,mmax,nmax,1); pyy = real(Ctemp)
    ! this is an estimate of pzz 1/8 dz off the surface; pre-use solvT
    solvT = (tt - ((phi0t - phi0(:,:,pmax))*2./(facz(pmax)*dz)))*4./(facz(pmax)*dz)
    call d_s2b(solvT,Ctemp,1,Id); call ft_2d(Ctemp,mmax,nmax,1); pzz = real(Ctemp)
    Rtemp = pzz*(pxx + pyy); call xy_to_sp(cmplx(Rtemp,0.),solvT,mmax,nmax,kmax,lmax)
    solv = (solvT(1,1) - solvB(1,1))/H
    if (verbose .gt. 1)  print*,'*-*-*-* O(R) solvability constant : ',solv

    ! phi_z at half levels; phi_zz on grid levels
    do k = 1,pmax
        if (k .le. pmax-1) then
        pzS(:,:,k) = (phi0(:,:,k+1) - phi0(:,:,k)) / (faczo(k)*dz)
        endif
        if (k .eq. 1) then
        pzzS(:,:,k) = (pzS(:,:,k) - tb) / (facz(k)*dz)
        elseif (k .eq. pmax) then
        pzzS(:,:,k) = (tt - pzS(:,:,k-1)) / (facz(k)*dz)
        else
        pzzS(:,:,k) = (pzS(:,:,k) - pzS(:,:,k-1)) / (facz(k)*dz)
        endif
    enddo
    pzzS(1,1,:) = 0.

    ! terms with mixed derivatives involving d/dz at half levels on big grid
    !$omp parallel do schedule(dynamic) private(k)
    do k=1,pmax

        pzz=0.;pzx=0.;pzy=0.;pxx=0.;pyy=0.;pxy=0.

        if (k .gt. 2 .and. k .lt. pmax) then
            pzxold = pzx; pzyold = pzy
        endif
        if (k .le. pmax-1) then
            call d_s2b(pzS(:,:,k),Ctemp,1,dx)
            call ft_2d(Ctemp,mmax,nmax,1); pzx = real(Ctemp)
            call d_s2b(pzS(:,:,k),Ctemp,1,dy)
            call ft_2d(Ctemp,mmax,nmax,1); pzy = real(Ctemp)
        endif
        ! start averaging back to grid levels
        if (k .eq. 1) then
            pzxold = pzx; pzyold = pzy
            call d_s2b(tb,Ctemp,1,dx)
            call ft_2d(Ctemp,mmax,nmax,1); tbx = real(Ctemp)
            call d_s2b(tb,Ctemp,1,dy)
            call ft_2d(Ctemp,mmax,nmax,1); tby = real(Ctemp)
            tbx = -tbx; tby = -tby
            pzx = (pzx + tbx)/2.; pzy = (pzy + tby)/2.
        elseif (k .eq. pmax) then ! recycle tbx & tby names here
            call d_s2b(tt,Ctemp,1,dx)
            call ft_2d(Ctemp,mmax,nmax,1); tbx = real(Ctemp)
            call d_s2b(tt,Ctemp,1,dy)
            call ft_2d(Ctemp,mmax,nmax,1); tby = real(Ctemp)
            pzx = (pzxold + tbx)/2.; pzy = (pzyold + tby)/2.
        else
            if (k .eq. pmax-1) then
            pzxold = pzx; pzyold = pzy
            endif
            pzx = (pzx + pzxold)/2.; pzy = (pzy + pzyold)/2.
        endif
        ! end averaging back to grid levels

        call d_s2b(pzzS(:,:,k),Ctemp,1,Id)
        call ft_2d(Ctemp,mmax,nmax,1); pzz = real(Ctemp)
        call d_s2b(phi0(:,:,k),Ctemp,1,dx*dx)
        call ft_2d(Ctemp,mmax,nmax,1); pxx = real(Ctemp)
        call d_s2b(phi0(:,:,k),Ctemp,1,dy*dy)
        call ft_2d(Ctemp,mmax,nmax,1); pyy = real(Ctemp)
        call d_s2b(phi0(:,:,k),Ctemp,1,dx*dy)
        call ft_2d(Ctemp,mmax,nmax,1); pxy = real(Ctemp)

        ! compute source terms
        Rtemp = 2*(pzx*pxy - pzy*pxx)
        call xy_to_sp(cmplx(Rtemp,0.),F1s(:,:,k),mmax,nmax,kmax,lmax)
        Rtemp = 2*(pzx*pyy - pzy*pxy)
        call xy_to_sp(cmplx(Rtemp,0.),G1s(:,:,k),mmax,nmax,kmax,lmax)
        Rtemp = (pzx*pzx) + (pzy*pzy) - ((pxx + pyy)*pzz)
        !     Rtemp = (pzx*pzx) + (pzy*pzy) + (pzz*pzz) ! test for zero pv
        call xy_to_sp(cmplx(Rtemp,0.),P1s(:,:,k),mmax,nmax,kmax,lmax)
        P1s(1,1,k) = P1s(1,1,k) + solv
        Rtemp = 2*(pxx*pyy - pxy*pxy)
        call xy_to_sp(cmplx(Rtemp,0.),phi1s(:,:,k),mmax,nmax,kmax,lmax)

    enddo
    !$omp end parallel do

    ! inversion (mod 19 July to correct F,G calls to matinv)
    !$omp parallel do schedule(dynamic) private(k,l,ak,bl,psi)
    do k=1,2*kmax; do l=1,2*lmax

        ! get wavenumbers
        call get_waves(k,l,ak,bl)

        ! spectral inversion (zero mean)
        if (k .eq. 1 .and. l .eq. 1) then
            F1s(k,l,:)=0.; G1s(k,l,:)=0.; P1s(k,l,:)=0.; phi1s(k,l,:)=0.
            !        F1s(k,l,:)=0.; G1s(k,l,:)=0.; phi1s(k,l,:)=0.
            F1(k,l,:) = 0.; G1(k,l,:)=0.; P1(k,l,:) = 0.
            !     endif
        else
            ! phi (geopotential) 7/20/04 mods here
            ! phi1 BC turns out to be homog. Neumann
            call matinv(phi1s(k,l,:),ak,bl,psi,1)
            phi1b(k,l) = psi(1)
            phi1t(k,l) = psi(pmax)
            phi1(k,l,:) = psi(:)
            ! F
            call matinv(F1s(k,l,:),ak,bl,psi,-1); F1(k,l,:) = psi(:)
            ! G
            call matinv(G1s(k,l,:),ak,bl,psi,-1); G1(k,l,:) = psi(:)
            ! Phi
            call matinv(P1s(k,l,:),ak,bl,psi,1); P1(k,l,:) = psi(:)
        endif
    enddo; enddo
    !$omp end parallel do

    ! finally, correct phi from Phi
    phi1 = phi1 + P1
    phi1b = phi1b + P1(:,:,1) ! P1_z = 0--> P1(z=0) = P1(z=dz/2)
    phi1t = phi1t + P1(:,:,pmax)

    !---------- Phi1 solution check
    bcb = 0.; bct = 0.
    call laplacian(P1,bcb,bct,lap)

    do k = 1,pmax
        errsp = P1s(:,:,k) - lap(:,:,k); errsp(1,1) = 0.
        call sp_to_xy(errsp,errxy,kmax,lmax,2*kmax,2*lmax)
        if (verbose .gt. 1) print*,'max P1 solution error at level ',k,' : ',maxval(abs(errxy))

    enddo
        !---------- F1 solution check
    bcb = 2*F1(:,:,1)/dz; bct = -2*F1(:,:,pmax)/dz
    call laplacian(F1,bcb,bct,lap)
    do k = 1,pmax
        errsp = F1s(:,:,k) - lap(:,:,k)
        call sp_to_xy(errsp,errxy,kmax,lmax,2*kmax,2*lmax)
        if (verbose .gt. 1) print*,'max F1 solution error at level ',k,' : ',maxval(abs(errxy))
    enddo

    !---------- G1 solution check
    bcb = 2*G1(:,:,1)/dz; bct = -2*G1(:,:,pmax)/dz
    call laplacian(G1,bcb,bct,lap)
    do k = 1,pmax
        errsp = G1s(:,:,k) - lap(:,:,k)
        call sp_to_xy(errsp,errxy,kmax,lmax,2*kmax,2*lmax)
        if (verbose .gt. 1) print*,'max G1 solution error at level ',k,' : ',maxval(abs(errxy))
    enddo

    return
END SUBROUTINE invert_R
!========================================================================

!========================================================================
SUBROUTINE w_check(u,v,win,wout)

    implicit none

    complex, dimension(2*kmax,2*lmax,pmax), intent(in)  :: u,v,win
    complex, dimension(2*kmax,2*lmax,pmax), intent(out) :: wout

    complex, dimension(2*kmax,2*lmax)      :: dx,dy,Id
    complex, dimension(2*kmax,2*lmax,pmax) :: w
    complex, dimension(2*kmax,2*lmax)      :: temp,wnew,wold,ux,vy,wF
    real,    dimension(2*kmax,2*lmax)      :: wxo,wxn

    real    :: dz
    integer :: p

    ! vertical grid distance
    dz = ZH/real(pmax)

    ! derivative operators
    call d_setup(dx,dy,Id)

    ! integrate u_x + v_y = -Ross w_z
    ! note that u and v are next-order, so Ross on RHS cancels w/ those that
    ! are on the LHS.
    wold = 0.
    do p = 1, pmax
        call d_s2s(u(:,:,p),ux,1,dx)
        call d_s2s(v(:,:,p),vy,1,dy)
        wnew = wold - (ux + vy)*dz
        !print*,'max ux = ',maxval(abs(ux)),maxval(abs(vy)),dz, maxval(abs(ux+vy))
        !wF = (wnew + wold)/2. !why???
        wF = wnew
        wout(:,:,p) = wnew
        wold = wnew; wxn = 0.
        call sp_to_xy(win(:,:,p),wxo,kmax,lmax,2*kmax,2*lmax)
        call sp_to_xy(wF,wxn,kmax,lmax,2*kmax,2*lmax)
        if (verbose .gt. 1) print*,'max w comparison (old,new): ',maxval(abs(wxo)), maxval(abs(wxn))
    enddo

    return
END SUBROUTINE w_check
!========================================================================

!========================================================================
SUBROUTINE laplacian(func,fb,ft,lap)
! compute an isotropic 3D laplacian given a staggered field & BCs.

    implicit none

    complex, dimension(2*kmax,2*lmax,pmax), intent(in)  :: func
    complex, dimension(2*kmax,2*lmax),      intent(in)  :: fb,ft
    complex, dimension(2*kmax,2*lmax,pmax), intent(out) :: lap

    complex, dimension(2*kmax,2*lmax,pmax-1) :: fz
    complex, dimension(2*kmax,2*lmax)        :: fxx,fyy,fzz
    complex, dimension(2*kmax,2*lmax)        :: dx,dy,Id

    integer :: k
    real    :: dz

    dz = ZH/real(pmax)

    ! derivative operators
    call d_setup(dx,dy,Id)

    dz = ZH/real(pmax)

    ! first z-loop for f_z at at intermediate levels
    !$omp parallel do schedule(dynamic) private(k)
    do k = 1,pmax-1
        fz(:,:,k) = (func(:,:,k+1) - func(:,:,k)) / dz
    enddo
    !$omp end parallel do

    ! second z-loop for laplacian at grid levels and f_zz; lap calc
    !$omp parallel do schedule(dynamic) private(k,fxx,fyy,fzz)
    do k = 1,pmax

        call d_s2s(func(:,:,k),fxx,1,dx*dx)
        call d_s2s(func(:,:,k),fyy,1,dy*dy)

        if (k .eq. 1) then
            fzz = (fz(:,:,k) - fb) / dz
        elseif (k .eq. pmax) then
            fzz = (ft - fz(:,:,k-1)) / dz
        else
            fzz = (fz(:,:,k) - fz(:,:,k-1)) / dz
        endif

        ! this is critical to get the inversion to check; the mean is zero in the inversion
        fzz(1,1) = 0.

        lap(:,:,k) = fxx + fyy + fzz

    enddo
    !$omp end parallel do
    ! END leading-order PV check

    return
END SUBROUTINE laplacian
!========================================================================

!========================================================================
SUBROUTINE vortdiv(u,v,vort,div)
! given horizontal wind (u,v) compute vertical component of
! vorticity and horizontal divergence. NONDIMENSIONAL I/O.
! make these values dimensional by multiplying by U/L

    implicit none

    complex, dimension(2*kmax,2*lmax,pmax), intent(in)  :: u,v
    real,    dimension(2*kmax,2*lmax,pmax), intent(out) :: vort,div

    complex, dimension(2*kmax,2*lmax) :: ux,uy,vx,vy
    complex, dimension(2*kmax,2*lmax) :: dx,dy,Id,vs,ds

    integer :: k

    ! derivative operators
    call d_setup(dx,dy,Id)

    !$omp parallel do schedule(dynamic) private(k,ux,vx,uy,vy,vs,ds)
    do k = 1,pmax
        call d_s2s(v(:,:,k),vx,1,dx); call d_s2s(v(:,:,k),vy,1,dy)
        call d_s2s(u(:,:,k),ux,1,dx); call d_s2s(u(:,:,k),uy,1,dy)
        ! spectral vorticity and divergence
        vs = vx - uy
        ds = ux + vy
        ! grid point
        call sp_to_xy(vs,vort(:,:,k),kmax,lmax,2*kmax,2*lmax)
        call sp_to_xy(ds,div(:,:,k),kmax,lmax,2*kmax,2*lmax)
    enddo
    !$omp end parallel do

    return
END SUBROUTINE vortdiv
!========================================================================

!========================================================================
SUBROUTINE setup_tanh(facz,faczo,zl,zhl)
! setup Jacobian factors for tanh vertical grid

    implicit none

    real, dimension(pmax), intent(out) :: facz, faczo,zl,zhl

    integer :: k
    real    :: fac,alpha,acosh,znot,z,arg,dz

    ! if we're using the regular grid, set jacobian values to one:
    if (.not. tanhgrid) then
        facz = 1.; faczo = 1.
        return
    endif

    dz = ZH/real(pmax)

    ! resolution ratio factor (near-boundary res is fac * mid level)
    fac = 10;

    ! nondimensional midlevel in *both* coordinates
    znot = .5;

    ! grid stretch factor
    alpha = 2.*acosh(sqrt(fac));

    ! first do Jacobian factors for interior grid levels that have the
    ! first level located one-half grid level off the boundary.
    do k = 1,pmax
        z = (real(k) - .5)*dz
        arg = alpha*(z-.5)
        facz(k) = 2.*tanh(alpha/2.)*(1./(1. - (tanh(arg))**2))/alpha
        zl(k) = znot + (tanh(alpha*(z-.5))) / (2.*tanh(alpha/2.))
        if ( verbose .gt. 1 ) print*,'jacobian:',k,z,zl(k),facz(k)
    enddo

    !now do the "half levels", which are actually spaced evenly with no staggering.
    do k = 1,pmax-1
        z = real(k)*dz
        arg = alpha*(z-.5)
        faczo(k) = 2.*tanh(alpha/2.)*(1./(1. - (tanh(arg))**2))/alpha
        zhl(k) = znot + (tanh(alpha*(z-.5))) / (2.*tanh(alpha/2.))
        if ( verbose .gt. 1 ) print*,'jacobian:',k,z,zhl(k),faczo(k)
    enddo

    return
END SUBROUTINE setup_tanh
!========================================================================

!========================================================================
FUNCTION acosh(x)

    implicit none

    real, intent(in) :: x
    real :: acosh

    if (x .lt. 1) then
        print*,'error in acosh routine...bad argument'
        return
    else
        acosh = log(x + sqrt(x*x - 1))
    endif

    return
END FUNCTION acosh
!========================================================================

!========================================================================
SUBROUTINE epv_pseudoheight(u,v,t,ub,vb,tb,ut,vt,tt,epv_out)
! inputs: horizontal winds on grid levels: u,v.
!         horizontal winds on grid boundaries: ub,vb.
!         potential temperature on grid levels: t
!         potential temperature on grid boundaries: tb
! output: Ertel potential vorticity: out
!
! epv = v_x - u_y + t_z + Ross * [(v_x - u_y) * t_z - v_z * t_x + u_z * t_y]
!
! all variables are SPECTRAL and NON-DIMENSIONAL, _except_ for out.!

    implicit none

    complex, dimension(2*kmax,2*lmax,pmax), intent(inout) :: u,v,t
    complex, dimension(2*kmax,2*lmax),      intent(in)    :: ub,vb,tb,ut,vt,tt
    real,    dimension(2*kmax,2*lmax,pmax), intent(out)   :: epv_out

    complex, dimension(2*kmax,2*lmax,pmax) :: uz,vz,tz
    complex, dimension(2*kmax,2*lmax)      :: dx,dy,Id,vx,uy,tx,ty
    complex, dimension(2*kmax,2*lmax)      :: tmp

    real    :: dz
    integer :: i,j,k

    ! vertical grid distance
    dz = ZH/real(pmax); if (verbose .gt. 1) print*,'dz = ',dz

    ! derivative operators
    call d_setup(dx,dy,Id)

    ! zero mean theta
    !tb(1,1) = 0.; tt(1,1) = 0.; t(1,1,:) = 0.

    ! vertical derivatives by centered finite differences, except near boundaries.
    ! the near-boundary points use centered differencing and averaging to get the
    ! mid-grid points (e.g. grid level 1.5).
    !$omp parallel do schedule(dynamic) private(k)
    do k=1,pmax
        if (k .eq. 1) then
            vz(:,:,k) = (v(:,:,k+1) + v(:,:,k) - 2.*vb(:,:)) / (2.*dz)
            uz(:,:,k) = (u(:,:,k+1) + u(:,:,k) - 2.*ub(:,:)) / (2.*dz)
            tz(:,:,k) = (t(:,:,k) - tb(:,:)) / dz
        elseif (k .eq. pmax) then
            vz(:,:,k) = (2.*vt(:,:) - v(:,:,k-1) - v(:,:,k)) / (2.*dz)
            uz(:,:,k) = (2.*ut(:,:) - u(:,:,k-1) - u(:,:,k)) / (2.*dz)
            tz(:,:,k) = (tt(:,:) - t(:,:,k-1)) / dz
        else
            vz(:,:,k) = (v(:,:,k+1)-v(:,:,k-1)) / (2.*dz)
            uz(:,:,k) = (u(:,:,k+1)-u(:,:,k-1)) / (2.*dz)
            tz(:,:,k) = (t(:,:,k) - t(:,:,k-1)) / dz
        endif
        if (verbose .gt. 1) print*,'max sp vz = ',maxval(abs(vz(:,:,k)))
    enddo
    !$omp end parallel do

    ! spectral horizontal derivatives and epv at each level
    epv_out = 0.
    do k=1,pmax

        call d_s2s(v(:,:,k),vx,1,dx)
        call d_s2s(u(:,:,k),uy,1,dy)
        call d_s2s(t(:,:,k),tx,1,dx)
        call d_s2s(t(:,:,k),ty,1,dy)

        tmp = vx - uy + tz(:,:,k) + Ross*( (vx - uy)*tz(:,:,k) - vz(:,:,k)*tx + uz(:,:,k)*ty )
        ! map back to physical space
        call sp_to_xy(tmp,epv_out(:,:,k),kmax,lmax,2*kmax,2*lmax)
        !call sp_to_xy(ty(:,:),out(:,:,k),kmax,lmax,2*kmax,2*lmax)
        if (verbose .gt. 1) print*,'max recovered pv at level ',k,' : ',maxval(abs(epv_out(:,:,k)))

    enddo

    return
END SUBROUTINE epv_pseudoheight
!========================================================================

!========================================================================
SUBROUTINE epv_pseudoheight_new(u,v,t,tb,tt,phi0,phi0b,phi0t,epv_out)
! inputs: horizontal winds on grid levels: u,v.
!         horizontal winds on grid boundaries: ub,vb.
!         potential temperature on grid levels: t
!         potential temperature on grid boundaries: tb
! output: Ertel potential vorticity: out
!
! epv = v_x - u_y + t_z + Ross * [(v_x - u_y) * t_z - v_z * t_x + u_z * t_y]
!
! all variables are SPECTRAL and NON-DIMENSIONAL, _except_ for out.!

    implicit none

    complex, dimension(2*kmax,2*lmax,pmax), intent(inout) :: u,v,t,phi0
    complex, dimension(2*kmax,2*lmax),      intent(inout) :: tb,tt,phi0b,phi0t
    real, dimension(2*kmax,2*lmax,pmax),    intent(out)   :: epv_out
    !
    complex, dimension(2*kmax,2*lmax,pmax) :: uz,vz,tz,u0,v0
    complex, dimension(2*kmax,2*lmax)      :: dx,dy,Id,vx,uy,tx,ty,v0x,u0y,u0b,u0t,v0b,v0t
    complex, dimension(2*kmax,2*lmax)      :: tmp

    real    :: dz
    integer :: i,j,k

    ! vertical grid distance
    dz = ZH/real(pmax); if (verbose .gt. 1) print*,'dz = ',dz

    ! derivative operators
    call d_setup(dx,dy,Id)

    ! zero mean theta
    tb(1,1) = 0.; tt(1,1) = 0.; t(1,1,:) = 0.

    ! leading order winds
    !$omp parallel do schedule(dynamic) private(k)
    do k = 1,pmax
        call d_s2s(phi0(:,:,k),u0(:,:,k),1,-dy)
        call d_s2s(phi0(:,:,k),v0(:,:,k),1, dx)
    enddo
    !$omp end parallel do

    ! leading-order boundary winds
    call d_s2s(phi0b,u0b,1,-dy); call d_s2s(phi0t,u0t,1,-dy)
    call d_s2s(phi0b,v0b,1, dx); call d_s2s(phi0t,v0t,1, dx)

    ! vertical derivatives
    !$omp parallel do schedule(dynamic) private(k)
    do k=1,pmax
        if (k .eq. 1) then
            vz(:,:,k) = (v0(:,:,k+1) + v0(:,:,k) - 2.*v0b(:,:)) / (2.*dz)
            uz(:,:,k) = (u0(:,:,k+1) + u0(:,:,k) - 2.*u0b(:,:)) / (2.*dz)
            tz(:,:,k) = (t(:,:,k) - tb(:,:)) / dz
        elseif (k .eq. pmax) then
            vz(:,:,k) = (2.*v0t(:,:) - v0(:,:,k-1) - v0(:,:,k)) / (2.*dz)
            uz(:,:,k) = (2.*u0t(:,:) - u0(:,:,k-1) - u0(:,:,k)) / (2.*dz)
            tz(:,:,k) = (tt(:,:) - t(:,:,k-1)) / dz
        else
            vz(:,:,k) = (v0(:,:,k+1)-v0(:,:,k-1)) / (2.*dz)
            uz(:,:,k) = (u0(:,:,k+1)-u0(:,:,k-1)) / (2.*dz)
            tz(:,:,k) = (t(:,:,k) - t(:,:,k-1)) / dz
        endif
    enddo
    !$omp end parallel do

    ! spectral horizontal derivatives and epv at each level
    epv_out = 0.
    do k=1,pmax
        call d_s2s(v(:,:,k),vx,1,dx)
        call d_s2s(u(:,:,k),uy,1,dy)
        call d_s2s(v0(:,:,k),v0x,1,dx)
        call d_s2s(u0(:,:,k),u0y,1,dy)
        call d_s2s(t(:,:,k),tx,1,dx)
        call d_s2s(t(:,:,k),ty,1,dy)
        ! linear piece
        tmp = vx - uy + tz(:,:,k)
        ! nonlinear piece
        !tmp = tmp + Ross*( (v0x - u0y)*tz(:,:,k) - vz(:,:,k)*tx + uz(:,:,k)*ty )
        ! map back to physical space
        call sp_to_xy(tmp,epv_out(:,:,k),kmax,lmax,2*kmax,2*lmax)
        !call sp_to_xy(vx-uy,out(:,:,k),kmax,lmax,2*kmax,2*lmax)
        !call sp_to_xy(tz(:,:,k),out(:,:,k),kmax,lmax,2*kmax,2*lmax)
        print*,'--max recovered pv at level ',k,' : ',maxval(abs(epv_out(:,:,k)))
        !print*,'max recovered pv at level ',k,' : ',minval(out(:,:,k))
    enddo

    return
END SUBROUTINE epv_pseudoheight_new
!========================================================================

!========================================================================
SUBROUTINE write_diag(output_file,it,u_s,v_s,theta_s,phi_s,u,v,theta,phi,w,pv)
! write to disk

    implicit none

    character(len=*),                    intent(in) :: output_file
    integer,                             intent(in) :: it
    real, dimension(2*kmax,2*lmax,2),    intent(in) :: u_s,v_s,theta_s,phi_s
    real, dimension(2*kmax,2*lmax,pmax), intent(in) :: u,v,theta,phi,w,pv

    integer :: ncid, vardim(5), svardim(4), pvardim(4), varid
    integer :: count(4), start(4), ierr
    real    :: time

    if ( it .eq. 0 ) then

        ! Create a new NetCDF file
        call nc_check( nf90_create(output_file, ior(NF90_CLOBBER,NF90_64BIT_OFFSET), ncid), 'write_diag', 'create ' // trim(output_file) )

        ! Define dimensions
        call nc_check( nf90_def_dim(ncid, "nx",   2*kmax,         vardim(1)), 'write_diag', 'def_dim, nx '   // trim(output_file) )
        call nc_check( nf90_def_dim(ncid, "ny",   2*lmax,         vardim(2)), 'write_diag', 'def_dim, ny '   // trim(output_file) )
        call nc_check( nf90_def_dim(ncid, "nz",   2,              vardim(3)), 'write_diag', 'def_dim, nz '   // trim(output_file) )
        call nc_check( nf90_def_dim(ncid, "np",   pmax,           vardim(4)), 'write_diag', 'def_dim, np '   // trim(output_file) )
        call nc_check( nf90_def_dim(ncid, "time", NF90_UNLIMITED, vardim(5)), 'write_diag', 'def_dim, time ' // trim(output_file) )

        ! Define variable dimensions
        svardim(1) = vardim(1) ; pvardim(1) = vardim(1)
        svardim(2) = vardim(2) ; pvardim(2) = vardim(2)
        svardim(3) = vardim(3) ; pvardim(3) = vardim(4)
        svardim(4) = vardim(5) ; pvardim(4) = vardim(5)

        ! Define variables
        call nc_check( nf90_def_var(ncid, "time",    NF90_FLOAT, vardim(5), varid),  'write_diag', 'def_var, time '    // trim(output_file) )
        call nc_check( nf90_def_var(ncid, "u_s",     NF90_FLOAT, svardim,   varid),  'write_diag', 'def_var, u_s '     // trim(output_file) )
        call nc_check( nf90_put_att(ncid, varid, "description", "surface zonal wind"), 'write_diag', 'put_att, description ' // trim(output_file) )
        call nc_check( nf90_def_var(ncid, "v_s",     NF90_FLOAT, svardim,   varid),  'write_diag', 'def_var, v_s '     // trim(output_file) )
        call nc_check( nf90_put_att(ncid, varid, "description", "surface meridional wind"), 'write_diag', 'put_att, description ' // trim(output_file) )
        call nc_check( nf90_def_var(ncid, "theta_s", NF90_FLOAT, svardim,   varid),  'write_diag', 'def_var, theta_s ' // trim(output_file) )
        call nc_check( nf90_put_att(ncid, varid, "description", "surface potential temperature"), 'write_diag', 'put_att, description ' // trim(output_file) )
        call nc_check( nf90_def_var(ncid, "phi_s",   NF90_FLOAT, svardim,   varid),  'write_diag', 'def_var, phi_s '   // trim(output_file) )
        call nc_check( nf90_put_att(ncid, varid, "description", "surface geopotential"), 'write_diag', 'put_att, description ' // trim(output_file) )
        call nc_check( nf90_def_var(ncid, "u",       NF90_FLOAT, pvardim,   varid),  'write_diag', 'def_var, u '       // trim(output_file) )
        call nc_check( nf90_put_att(ncid, varid, "description", "interior zonal wind"), 'write_diag', 'put_att, description ' // trim(output_file) )
        call nc_check( nf90_def_var(ncid, "v",       NF90_FLOAT, pvardim,   varid),  'write_diag', 'def_var, v '       // trim(output_file) )
        call nc_check( nf90_put_att(ncid, varid, "description", "interior meridional wind"), 'write_diag', 'put_att, description ' // trim(output_file) )
        call nc_check( nf90_def_var(ncid, "theta",   NF90_FLOAT, pvardim,   varid),  'write_diag', 'def_var, theta '   // trim(output_file) )
        call nc_check( nf90_put_att(ncid, varid, "description", "interior potential temperature"), 'write_diag', 'put_att, description ' // trim(output_file) )
        call nc_check( nf90_def_var(ncid, "phi",     NF90_FLOAT, pvardim,   varid),  'write_diag', 'def_var, phi '     // trim(output_file) )
        call nc_check( nf90_put_att(ncid, varid, "description", "interior geopotential"), 'write_diag', 'put_att, description ' // trim(output_file) )
        call nc_check( nf90_def_var(ncid, "w",       NF90_FLOAT, pvardim,   varid),  'write_diag', 'def_var, w '       // trim(output_file) )
        call nc_check( nf90_put_att(ncid, varid, "description", "interior vertical velocity"), 'write_diag', 'put_att, description ' // trim(output_file) )
        call nc_check( nf90_def_var(ncid, "pv",      NF90_FLOAT, pvardim,   varid),  'write_diag', 'def_var, pv '      // trim(output_file) )
        call nc_check( nf90_put_att(ncid, varid, "description", "interior potential vorticity"), 'write_diag', 'put_att, description ' // trim(output_file) )

        ! Put global attributes
        call nc_check( nf90_put_att(ncid, NF90_GLOBAL, "model", model), 'write_diag', 'put_att, model ' // trim(output_file) )
        call nc_check( nf90_put_att(ncid, NF90_GLOBAL, "ntims", ntims), 'write_diag', 'put_att, ntims ' // trim(output_file) )
        call nc_check( nf90_put_att(ncid, NF90_GLOBAL, "dt",    dt),    'write_diag', 'put_att, dt '    // trim(output_file) )
        call nc_check( nf90_put_att(ncid, NF90_GLOBAL, "iplot", iplot), 'write_diag', 'put_att, iplot ' // trim(output_file) )
        call nc_check( nf90_put_att(ncid, NF90_GLOBAL, "XL",    XL),    'write_diag', 'put_att, XL '    // trim(output_file) )
        call nc_check( nf90_put_att(ncid, NF90_GLOBAL, "YL",    YL),    'write_diag', 'put_att, YL '    // trim(output_file) )
        call nc_check( nf90_put_att(ncid, NF90_GLOBAL, "H",     H),     'write_diag', 'put_att, H '     // trim(output_file) )
        call nc_check( nf90_put_att(ncid, NF90_GLOBAL, "Ross",  Ross),  'write_diag', 'put_att, Ross '  // trim(output_file) )
        call nc_check( nf90_put_att(ncid, NF90_GLOBAL, "gamma", gamma), 'write_diag', 'put_att, gamma ' // trim(output_file) )
        call nc_check( nf90_put_att(ncid, NF90_GLOBAL, "n",     n),     'write_diag', 'put_att, n '     // trim(output_file) )
        call nc_check( nf90_put_att(ncid, NF90_GLOBAL, "tau",   tau),   'write_diag', 'put_att, tau '   // trim(output_file) )
        call nc_check( nf90_put_att(ncid, NF90_GLOBAL, "trl",   trl),   'write_diag', 'put_att, trl '   // trim(output_file) )
        call nc_check( nf90_put_att(ncid, NF90_GLOBAL, "amu",   amu),   'write_diag', 'put_att, amu '   // trim(output_file) )
        call nc_check( nf90_put_att(ncid, NF90_GLOBAL, "shear", shear), 'write_diag', 'put_att, shear ' // trim(output_file) )

        call nc_check( nf90_enddef(ncid), 'write_diag', 'enddef, ' // trim(output_file) )
        call nc_check( nf90_close(ncid),  'write_diag', 'close, '  // trim(output_file) )

    else

        time = (it-1)*dt

        ! Open the netCDF file
        call nc_check( nf90_open(output_file, NF90_WRITE, ncid), 'write_diag', 'open, '           // trim(output_file) )

        ! Write time variable
        call nc_check( nf90_inq_varid(ncid, "time",    varid) ,  'write_diag', 'inq_varid, time ' // trim(output_file) )

        count(1) = 2*kmax;  start(1) = 1
        count(2) = 2*lmax;  start(2) = 1
        count(3) = 2;       start(3) = 1
        count(4) = 1;       start(4) = it

        ! Write surface variables
        call nc_check( nf90_put_var(ncid, varid, time, (/it/)),          'write_diag', 'put_var, time '      // trim(output_file) )
        call nc_check( nf90_inq_varid(ncid, "u_s",     varid) ,          'write_diag', 'inq_varid, u_s '     // trim(output_file) )
        call nc_check( nf90_put_var(ncid, varid, u_s,     start, count), 'write_diag', 'put_var, u_s '       // trim(output_file) )
        call nc_check( nf90_inq_varid(ncid, "v_s",     varid) ,          'write_diag', 'inq_varid, v_s '     // trim(output_file) )
        call nc_check( nf90_put_var(ncid, varid, v_s,     start, count), 'write_diag', 'put_var, v_s '       // trim(output_file) )
        call nc_check( nf90_inq_varid(ncid, "theta_s", varid),           'write_diag', 'inq_varid, theta_s ' // trim(output_file) )
        call nc_check( nf90_put_var(ncid, varid, theta_s, start, count), 'write_diag', 'put_var, theta_s '   // trim(output_file) )
        call nc_check( nf90_inq_varid(ncid, "phi_s",   varid),           'write_diag', 'inq_varid, phi_s '   // trim(output_file) )
        call nc_check( nf90_put_var(ncid, varid, phi_s,   start, count), 'write_diag', 'put_var, phi_s '     // trim(output_file) )

        count(1) = 2*kmax;  start(1) = 1
        count(2) = 2*lmax;  start(2) = 1
        count(3) = pmax;    start(3) = 1
        count(4) = 1;       start(4) = it

        ! Write interior variables
        call nc_check( nf90_inq_varid(ncid, "u",     varid) ,          'write_diag', 'inq_varid, u '     // trim(output_file) )
        call nc_check( nf90_put_var(ncid, varid, u,     start, count), 'write_diag', 'put_var, u '       // trim(output_file) )
        call nc_check( nf90_inq_varid(ncid, "v",     varid) ,          'write_diag', 'inq_varid, v '     // trim(output_file) )
        call nc_check( nf90_put_var(ncid, varid, v,     start, count), 'write_diag', 'put_var, v '       // trim(output_file) )
        call nc_check( nf90_inq_varid(ncid, "theta", varid),           'write_diag', 'inq_varid, theta ' // trim(output_file) )
        call nc_check( nf90_put_var(ncid, varid, theta, start, count), 'write_diag', 'put_var, theta '   // trim(output_file) )
        call nc_check( nf90_inq_varid(ncid, "phi",   varid),           'write_diag', 'inq_varid, phi '   // trim(output_file) )
        call nc_check( nf90_put_var(ncid, varid, phi,   start, count), 'write_diag', 'put_var, phi '     // trim(output_file) )
        call nc_check( nf90_inq_varid(ncid, "w",     varid) ,          'write_diag', 'inq_varid, w '     // trim(output_file) )
        call nc_check( nf90_put_var(ncid, varid, w,     start, count), 'write_diag', 'put_var, w '       // trim(output_file) )
        call nc_check( nf90_inq_varid(ncid, "pv",    varid) ,          'write_diag', 'inq_varid, pv '    // trim(output_file) )
        call nc_check( nf90_put_var(ncid, varid, pv,    start, count), 'write_diag', 'put_var, pv '      // trim(output_file) )

        ! Close the netCDF file
        call nc_check( nf90_close(ncid),                               'write_diag', 'close, '           // trim(output_file) )

    endif

    return
END SUBROUTINE write_diag
!========================================================================

!========================================================================
SUBROUTINE scaling(grav,tnot,Rhos,Ns,km,Cor,Ls,Hs,Us,Ws,Ps,Ts,PVs)
! scaling parameters:
! The _independent_ scaling parameters are:
! - Rossby number.
! - horizontal length scale.
! - vertical length scale.
! - surface temperature (tnot).
! - Coriolis parameter.
! - gravitational constant.
! - density.

! The _dependent_ scaling parameters are:
! - N ~ f * L / H
! - U ~ R * f * L
! - W ~ R * U * H / L
! - P ~ U * f * L
! - T ~ U * f * L * tnot / (g * H)
! - PV ~ U * f^2 * L * tnot / (rho * g * H^2)

    implicit none

    real, intent(out) :: grav,tnot,Rhos,Ns,km,Cor,Ls,Hs,Us,Ws,Ps,Ts,PVs
    grav = 9.81        ! gravitational constant (m s^-2)
    tnot = 300.        ! theta constant at Z = 0 (K)
    km   = 1000.       ! 1000 m in a kilometer
    Cor  = 1.e-4       ! Coriolis parameter (s^-1)
    Hs   = 10.*km      ! vertical length scale (m)
    Ls   = 1000.*km    ! horizontal length scale (m)
    Rhos = 1.0         ! density (kg m^-3)
    Ns   = Cor*Ls/Hs   ! bouyancy frequency (s-^-1)

    if (Ross .ne. 0) then
        Us = Ross*Cor*Ls   ! Horizontal wind scale (m s^-1)
        if (verbose .gt. 1)  print*,'Us = ',Us
    else
        Us = 1.0
    endif

    Ws  = Ross*Hs*Us/Ls                 ! Vertical wind scale (m s^-1)
    Ps  = Us*Cor*Ls                     ! geopotential scale (m^2 s^-2)
    Ts  = Ps*tnot/(grav*Hs)             ! potential temperature (K)
    PVs = 1.e6 * Cor * Ts / (Rhos * Hs) ! PV scaling (PVU)

    return
END SUBROUTINE scaling
!========================================================================

!========================================================================
SUBROUTINE dx_echo(dxs,dys,dzs,dts)
! send back dimensional values for dx, dy, dz, and dt.

    implicit none

    real, intent(out) :: dxs,dys,dzs,dts
    real :: grav,tnot,Rhos,Ns,km,Cor,Ls,Hs,Us,Ws,Ps,Ts,PVs

    call scaling(grav,tnot,Rhos,Ns,km,Cor,Ls,Hs,Us,Ws,Ps,Ts,PVs)

    dxs = Ls*XL/real(2*kmax) ! meters
    dys = Ls*YL/real(2*kmax) ! meters
    dzs = Hs*ZH/real(pmax)   ! meters
    dts = Ls*dt/Us           ! seconds

    return
END SUBROUTINE dx_echo
!========================================================================

!========================================================================
SUBROUTINE nc_check(ierr,subr_name,context)

  ! check for netcdf errors

  implicit none

  integer,          intent(in) :: ierr
  character(len=*), intent(in) :: subr_name, context

  character(len=129) :: error_msg

  if (ierr /= nf90_noerr) then
    error_msg = trim(subr_name) // ': ' // trim(context) // ': ' // trim(nf90_strerror(ierr))
    print*,trim(adjustl(error_msg))
    stop
  end if

  return
END SUBROUTINE nc_check
!========================================================================

!========================================================================
FUNCTION get_ntimes(infile)

    implicit none

    character(len=*), intent(in) :: infile
    integer :: get_ntimes
    integer :: ncid, dimid

    call nc_check( nf90_open(trim(infile), NF90_NOWRITE, ncid), 'get_ntimes', 'open, ' // trim(infile) )
    call nc_check( nf90_inq_dimid(ncid, 'time', dimid), 'get_ntimes', 'inq_dimid time, ' // trim(infile) )
    call nc_check( nf90_inquire_dimension(ncid, dimid, len = get_ntimes), 'get_ntimes', 'inquire_dimension time,' // trim(infile) )
    call nc_check( nf90_close(ncid), 'get_ntimes', 'close, ' // trim(infile) )

    return

END FUNCTION get_ntimes
!========================================================================

END MODULE pvinv_mod
!========================================================================


      subroutine var_adj(nt,a,var,work)
      real var(nt),work(nt,*),rmu,dmu
      complex a(nt,nt)

      parameter (dmu = .9,pmin=.2,pmax=1.0001)

      do k = 1,nt
         work(k,nt+1) = var(k)
      enddo
      rmu = 1.0
10    continue
      do k = 1,nt
         work(k,nt+2) = 1.0
      enddo

      do l = 1,nt
         do k = 1,nt
            if(k.eq.l) then
               work(k,l) = 1.-rmu*abs(a(k,k))**2.
            else
               work(k,l) = rmu*abs(a(k,l))**2.*var(l)/var(k)
            endif
         enddo
      enddo

      iw1 = nt+2
      iw2 = nt+3
      call sgesv(nt,1,work,nt,work(1,iw2),work(1,iw1),nt,info)
      do i = 1,nt
         if((work(i,iw1).le.pmin).or.(work(i,iw1).gt.pmax)) then
            rmu = rmu*dmu
            go to 10
         endif
      enddo
      do i = 1,nt
         var(i) = var(i)*work(i,iw1)
      enddo
      return
      end


      function deriv5 (dx,y,n,i)
C...function       deriv5    01 08 74
      dimension y(n)
      if(n-4) 15,9,10
   15 if(n-2) 17,16,8
   16 a=(y(2)-y(1))/dx
      go to 17
    8 if(i-2) 11,12,13
   11 a=-18*y(i)+24*y(i+1)-6*y(i+2)
      go to 7
   12 a=-6*y(i-1)+6*y(i)
      go to 7
   13 a=6*y(i-2)-24*y(i-1)+18*y(i)
      go to 7
    9 if(i-2) 2,3,14
   14 if(i-4) 13,4,4
   10 if(i-2) 2,3,1
    1 if(n-1-i) 4,5,6
    2 a=4*y(4)-18*y(3)+36*y(2)-22*y(1)
      go to 7
    3 a=-2*y(4)+12*y(3)-6*y(2)-4*y(1)
      go to 7
    4 a=22*y(i)-36*y(i-1)+18*y(i-2)-4*y(i-3)
      go to 7
    5 a=3*y(i+1)+11*y(i)-21*y(i-1)+9*y(i-2)-2*y(i-3)
      go to 7
    6 a=-y(i+2)+8*y(i+1)-8*y(i-1)+y(i-2)
    7 deriv5=a*0.083333/dx
   17 return
      end


!
! This is included in the "linear_with_nonlinear_friction" friction computation.
! We make _FRICTION_DEPTH_POWER_ a constant for efficiency
!

    !subroutine precompute_friction_work(domain)
    !    type(domain_type), intent(inout) :: domain
    !
    !    integer(ip) :: i, j, jp1, ip1
    !    real(dp) :: depth_iph, depth_jph, nsq_iph, nsq_jph
    !    real(dp), parameter :: manning_depth_power = NEG_SEVEN_ON_THREE_dp, chezy_depth_power = -2.0_dp

    !    if(domain%timestepping_method /= 'leapfrog_linear_plus_nonlinear_friction') &
    !        stop 'precompute_friction_work can only be called with timestepping_method=leapfrog_linear_plus_nonlinear_friction'

    !    if(.not. allocated(domain%friction_work)) &
    !        stop 'friction_work is not allocated: ensure elevation is set before running this routine'

            if(truely_linear) then
                !
                ! Here we evaluate the depth assuming stage=domain%msl_linear
                ! This is the standard logic for the "truely linear" linear shallow water equations.
                !

                !$OMP PARALLEL DEFAULT(PRIVATE) SHARED(domain)
                !$OMP DO SCHEDULE(STATIC)
                do j = 1, domain%nx(2)
                    do i = 1, domain%nx(1)

                        ! UH component
                        ip1 = min(i+1, domain%nx(1))
                        depth_iph = 0.5_dp * (domain%msl_linear - domain%U(i,j,ELV) + domain%msl_linear - domain%U(ip1,j, ELV))
                        depth_iph = max(depth_iph, minimum_allowed_depth)
                        !nsq_iph = (0.5_dp * (sqrt(domain%manning_squared(i,j)) + sqrt(domain%manning_squared(ip1,j))))**2
                        nsq_iph = 0.5_dp * (domain%manning_squared(i,j) + domain%manning_squared(ip1,j))
                        domain%friction_work(i,j,UH) = gravity * nsq_iph * depth_iph**_FRICTION_DEPTH_POWER_

                        ! VH component
                        jp1 = min(j+1, domain%nx(2))
                        depth_jph = 0.5_dp * (domain%msl_linear - domain%U(i,j,ELV) + domain%msl_linear - domain%U(i,jp1, ELV))
                        depth_jph = max(depth_jph, minimum_allowed_depth)
                        !nsq_jph = (0.5_dp * (sqrt(domain%manning_squared(i,j)) + sqrt(domain%manning_squared(i,jp1))) )**2
                        nsq_jph = 0.5_dp * (domain%manning_squared(i,j) + domain%manning_squared(i,jp1))
                        domain%friction_work(i,j,VH) = gravity * nsq_jph * depth_jph**_FRICTION_DEPTH_POWER_

                    end do
                end do
                !$OMP END DO
                !$OMP END PARALLEL

                ! For a truely-linear solver, this is only required once
                domain%friction_work_is_setup = .true.
            else
                !
                ! This differs from the above, in that we use domain%U(:,:,STG) to compute the depth.
                ! This means we generally need to recompute the friction-work terms at every time-step.
                !
                !$OMP PARALLEL DEFAULT(PRIVATE) SHARED(domain)
                !$OMP DO SCHEDULE(STATIC)
                do j = 1, domain%nx(2)
                    do i = 1, domain%nx(1)

                        ! UH component
                        ip1 = min(i+1, domain%nx(1))
                        depth_iph = 0.5_dp * (domain%U(i,j,STG) - domain%U(i,j,ELV) + domain%U(ip1,j,STG) - domain%U(ip1,j, ELV))
                        depth_iph = max(depth_iph, minimum_allowed_depth)
                        !nsq_iph = (0.5_dp * (sqrt(domain%manning_squared(i,j)) + sqrt(domain%manning_squared(ip1,j))))**2
                        nsq_iph = 0.5_dp * (domain%manning_squared(i,j) + domain%manning_squared(ip1,j))
                        domain%friction_work(i,j,UH) = gravity * nsq_iph * depth_iph**_FRICTION_DEPTH_POWER_

                        ! VH component
                        jp1 = min(j+1, domain%nx(2))
                        depth_jph = 0.5_dp * (domain%U(i,j,STG) - domain%U(i,j,ELV) + domain%U(i,jp1,STG) - domain%U(i,jp1, ELV))
                        depth_jph = max(depth_jph, minimum_allowed_depth)
                        !nsq_jph = (0.5_dp * (sqrt(domain%manning_squared(i,j)) + sqrt(domain%manning_squared(i,jp1))) )**2
                        nsq_jph = 0.5_dp * (domain%manning_squared(i,j) + domain%manning_squared(i,jp1))
                        domain%friction_work(i,j,VH) = gravity * nsq_jph * depth_jph**_FRICTION_DEPTH_POWER_

                    end do
                end do
                !$OMP END DO
                !$OMP END PARALLEL
            end if

        !end subroutine


      SUBROUTINE mrqmin(x,y,sig,ndata,a,ia,ma,covar,alpha,nca,chisq,
     *funcs,alamda)
      INTEGER ma,nca,ndata,ia(ma),MMAX
      REAL alamda,chisq,funcs,a(ma),alpha(nca,nca),covar(nca,nca),
     *sig(ndata),x(ndata),y(ndata)
      PARAMETER (MMAX=20)
CU    USES covsrt,gaussj,mrqcof
      INTEGER j,k,l,mfit
      REAL ochisq,atry(MMAX),beta(MMAX),da(MMAX)
      SAVE ochisq,atry,beta,da,mfit
      if(alamda.lt.0.)then
        mfit=0
        do 11 j=1,ma
          if (ia(j).ne.0) mfit=mfit+1
11      continue
        alamda=0.001
        call mrqcof(x,y,sig,ndata,a,ia,ma,alpha,beta,nca,chisq,funcs)
        ochisq=chisq
        do 12 j=1,ma
          atry(j)=a(j)
12      continue
      endif
      do 14 j=1,mfit
        do 13 k=1,mfit
          covar(j,k)=alpha(j,k)
13      continue
        covar(j,j)=alpha(j,j)*(1.+alamda)
        da(j)=beta(j)
14    continue
      call gaussj(covar,mfit,nca,da,1,1)
      if(alamda.eq.0.)then
        call covsrt(covar,nca,ma,ia,mfit)
        return
      endif
      j=0
      do 15 l=1,ma
        if(ia(l).ne.0) then
          j=j+1
          atry(l)=a(l)+da(j)
        endif
15    continue
      call mrqcof(x,y,sig,ndata,atry,ia,ma,covar,da,nca,chisq,funcs)
      if(chisq.lt.ochisq)then
        alamda=0.1*alamda
        ochisq=chisq
        do 17 j=1,mfit
          do 16 k=1,mfit
            alpha(j,k)=covar(j,k)
16        continue
          beta(j)=da(j)
17      continue
        do 18 l=1,ma
          a(l)=atry(l)
18      continue
      else
        alamda=10.*alamda
        chisq=ochisq
      endif
      return
      END

NR-Functions/Numerical Recipes- Example & Functions/Functions/mrqmin.for


module mod_sha256

  use, intrinsic :: iso_fortran_env, only : int8, int32, int64, real64

  implicit none

  integer(kind=int32), parameter :: kk(64) = (/ &
     1116352408,  1899447441, -1245643825,  -373957723,   961987163,  1508970993, -1841331548, -1424204075, &
     -670586216,   310598401,   607225278,  1426881987,  1925078388, -2132889090, -1680079193, -1046744716, &
     -459576895,  -272742522,   264347078,   604807628,   770255983,  1249150122,  1555081692,  1996064986, &
    -1740746414, -1473132947, -1341970488, -1084653625,  -958395405,  -710438585,   113926993,   338241895, &
      666307205,   773529912,  1294757372,  1396182291,  1695183700,  1986661051, -2117940946, -1838011259, &
    -1564481375, -1474664885, -1035236496,  -949202525,  -778901479,  -694614492,  -200395387,   275423344, &
      430227734,   506948616,   659060556,   883997877,   958139571,  1322822218,  1537002063,  1747873779, &
     1955562222,  2024104815, -2067236844, -1933114872, -1866530822, -1538233109, -1090935817,  -965641998 /)

  integer(kind=int32),              public,  save :: hh(8)
  integer(kind=int32), allocatable, private, save :: wBuf(:)
  integer(kind=int64),              private, save :: nBlock, nBuf

contains

subroutine sha256_init(inWord)

  character(len=*), intent(in) :: inWord

  character(len=:), allocatable :: inWordP
  integer(kind=int64) :: wLen, wMod, wPad, i, c
  integer(kind=int8)  :: fWord(4)

  if(allocated(wBuf)) deallocate(wBuf)

  ! Initialise hh() with the fractional part
  !  of the square root of the first 8 prime numbers
  hh = [ &
    1779033703,-1150833019,1013904242,-1521486534, &
    1359893119,-1694144372, 528734635, 1541459225  ]

  ! Which is equivalent to doing this:
  ! hh(1) = transfer(shiftr(transfer(sqrt( 2.0_real64),   1_int64),20), 1_int32)
  ! hh(2) = transfer(shiftr(transfer(sqrt( 3.0_real64),   1_int64),20), 1_int32)
  ! hh(3) = transfer(shiftr(transfer(sqrt( 5.0_real64)-1, 1_int64),20), 1_int32)
  ! hh(4) = transfer(shiftr(transfer(sqrt( 7.0_real64)-1, 1_int64),20), 1_int32)
  ! hh(5) = transfer(shiftr(transfer(sqrt(11.0_real64)-2, 1_int64),20), 1_int32)
  ! hh(6) = transfer(shiftr(transfer(sqrt(13.0_real64)-2, 1_int64),20), 1_int32)
  ! hh(7) = transfer(shiftr(transfer(sqrt(17.0_real64)-3, 1_int64),20), 1_int32)
  ! hh(8) = transfer(shiftr(transfer(sqrt(19.0_real64)-3, 1_int64),20), 1_int32)

  ! Calculate padding needed until the message is
  !  (length+1)%512 = 448 bit
  wLen = len(inWord)
  wMod = mod(wLen+1,64_int64)
  if(wMod > 56) then
    wPad = 120-wMod
  else
    wPad = 56-wMod
  end if

  ! Count the number of 512 bit blocks and
  !  allocate the buffer of 32 bit words
  nBlock = (wLen+wPad+9)/64
  nBuf   = nBlock * 16
  allocate(wBuf(nBuf))

  ! Add the padding of a single 1 and rest 0
  inWordP = inWord//char(128)//repeat(char(0),wPad)

  ! Loop over all characters in the string in order 4,3,2,1,8,7,6,5,...
  !  and store them as int32 in little endian order
  do i=1,nBuf-2
    c = 4*(i-1)
    fWord(1) = transfer(inWordP(c+4:c+4),1_int8)
    fWord(2) = transfer(inWordP(c+3:c+3),1_int8)
    fWord(3) = transfer(inWordP(c+2:c+2),1_int8)
    fWord(4) = transfer(inWordP(c+1:c+1),1_int8)
    wBuf(i)  = transfer(fWord(1:4),int8)
  end do
  ! Append the length as an int64 spanning 2 words at the end
  wBuf(nBuf:nBuf-1:-1) = transfer(wLen*8,int32,2)

end subroutine sha256_init

subroutine sha256_hash

  integer(kind=int64) :: i,j
  integer(kind=int32) :: h(8),t1,t2
  integer(kind=int32) :: w(64)

  do i=1,nBlock ! Loop over blocks

    ! Initialise the block digest variables
    h = hh

    do j=1,16 ! Read the 16 words in the current block
      w(j) = wBuf(16*(i - 1) + j)
    end do
    do j=17,64 ! Compute the remaining 48 words from the first 16
      w(j) = sigmaRRS1(w(j-2)) + w(j-7) + sigmaRRS0(w(j-15)) + w(j-16)
    end do

    do j=1,64 ! Digest the 64 words
      t1     = h(8) + sigmaRRR1(h(5)) + ch(h(5),h(6),h(7)) + kk(j) + w(j)
      t2     = sigmaRRR0(h(1)) + maj(h(1),h(2),h(3))
      h(2:8) = h(1:7)    ! Shift all the values
      h(5)   = h(5) + t1 ! Add the new word
      h(1)   = t1 + t2   ! Add the new word
    end do

    ! Add the results for this block to the message digest
    hh = h + hh

  end do

end subroutine sha256_hash

subroutine sha256_digest(outWord)
  character(len=64), intent(out) :: outWord
  write(outWord,"(8(z8.8))") hh
end subroutine sha256_digest

! (X AND Y) XOR ((NOT X) AND Z)
pure elemental integer(kind=int32) function ch(x,y,z)
  integer(kind=int32), intent(in) :: x,y,z
  ch = ieor(iand(x,y),iand(not(x),z))
end function ch

! (X AND Y) XOR (X AND Z) XOR (Y AND Z)
pure elemental integer(kind=int32) function maj(x,y,z)
  integer(kind=int32), intent(in) :: x,y,z
  maj = ieor(ieor(iand(x,y),iand(x,z)),iand(y,z))
end function maj

! RotR(X,2) XOR RotR(X,13) XOR RotR(X,22)
pure elemental integer(kind=int32) function sigmaRRR0(x)
  integer(kind=int32), intent(in) :: x
  sigmaRRR0 = ieor(ieor(dshiftr(x,x,2),dshiftr(x,x,13)),dshiftr(x,x,22))
end function sigmaRRR0

! RotR(X,6) XOR RotR(X,11) XOR RotR(X,25)
pure elemental integer(kind=int32) function sigmaRRR1(x)
  integer(kind=int32), intent(in) :: x
  sigmaRRR1 = ieor(ieor(dshiftr(x,x,6),dshiftr(x,x,11)),dshiftr(x,x,25))
end function sigmaRRR1

! RotR(X,7) XOR RotR(X,18) XOR ShR(X,3)
pure elemental integer(kind=int32) function sigmaRRS0(x)
  integer(kind=int32), intent(in) :: x
  sigmaRRS0 = ieor(ieor(dshiftr(x,x,7),dshiftr(x,x,18)),shiftr(x,3))
end function sigmaRRS0

! RotR(X,17) XOR RotR(X,19) XOR ShR(X,10)
pure elemental integer(kind=int32) function sigmaRRS1(x)
  integer(kind=int32), intent(in) :: x
  sigmaRRS1 = ieor(ieor(dshiftr(x,x,17),dshiftr(x,x,19)),shiftr(x,10))
end function sigmaRRS1

end module mod_sha256