C Sm0.f
C
C VIS-5D function to calculate first guess of balanced streamfunction (Sm0),
C used for inversion of 3-D Ertel PV distribution (from MEAN fields). 

C ==============================================================================

      INTEGER FUNCTION USERFUNC( OUTGRID, OUTNL, OUTLOWLEV,
     *                           INGRID, NR, NC, NL, LOWLEV, MAXNL,
     *                           NVARS, NAMES,
     *                           DATE, TIME,
     *                           PROJECTION, PROJ_ARGS,
     *                           VERTICAL, VERT_ARGS )
C ARGUMENTS:
      implicit none
      INTEGER NVARS
      INTEGER NR, NC, NL(NVARS), LOWLEV(NVARS), MAXNL
      REAL OUTGRID(NR,NC,MAXNL)
      INTEGER OUTNL, OUTLOWLEV
      REAL INGRID(NR,NC,MAXNL,NVARS)
      CHARACTER*8 NAMES(NVARS)
      INTEGER DATE, TIME
      INTEGER PROJECTION
      REAL PROJ_ARGS(*)
      INTEGER VERTICAL
      REAL VERT_ARGS(*)

C LOCAL VARS:
      integer var, ir, ic, il
      integer iu, iv, ih, it, icor, imap, ih0 
      integer max_iters,iter,noconverg
      real grav, po, cp, rd, kap, ds, ds2, dssqr

      parameter (grav=9.81,po=100000.,cp=1004.5,rd=2.*cp/7.,kap=rd/cp)

      real thrs_psi,over_psi,
     $     res,resmod
      real Fndim
      real dudy,dvdx,
     $     relvor,psiref,delta        
      real p(22),ph(0:22),exnf(22),exnfh(0:22),dexnfh(0:22),
     $     u(200,200,22),v(200,200,22),
     $     h(200,200,22),t(200,200,22),theta(200,200,22),
     $     cor(200,200),
     $     map(200,200),map2(200,200),
     $     phi0(200,200,0:22),
     $     psi0(200,200,0:22),for_psi(200,200,22),
     $     thetaB(200,200),thetaT(200,200)  

C ==============================================================================
      
C Specify number of levels in OUTGRID
      OUTNL = MAXNL
      OUTLOWLEV = 0
      
C Find the Um, Vm, Hm, Tm, CORIOLIS and MAPFACCR variables (MEAN fields)
      iu = -1
      iv = -1
      ih = -1
      it = -1
      icor = -1
      imap = -1
      do var=1,nvars
        if (names(var) .eq. 'Um') iu = var
        if (names(var) .eq. 'Vm') iv = var
        if (names(var) .eq. 'Hm') ih = var         
        if (names(var) .eq. 'Tm') it = var         
        if (names(var) .eq. 'CORIOLIS') icor = var
        if (names(var) .eq. 'MAPFACCR') imap = var                  
      end do

C If Um, Vm, Hm, Tm, CORIOLIS or MAPFACCR not found, return error 1
      if ( iu.eq.-1 .or. iv.eq.-1 .or. ih.eq.-1 .or. it.eq.-1 .or. 
     $     icor.eq.-1 .or. imap.eq.-1 ) then
        print *, 
     $  "Couldn't find Um, Vm, Hm, Tm, CORIOLIS or MAPFACCR vars!"
        userfunc = 1
        return
      endif

C Grid-length ds (m)
      if (projection.Eq.2) then
        ds=proj_args(6)*1000.      
      else if (projection.Eq.3) then  
        ds=proj_args(5)*1000.
      else
        print *, "Uncomprensible projection data!"
        userfunc = 1
        return
      endif

C Definition max_iters,thrs_psi,over_psi,ds2, dssqr
      max_iters=500      
      thrs_psi=1.E4
      over_psi=2.-3.1415927*(2./Real(nc-1)**2.+2./Real(nr-1)**2.)**.5
      ds2=2.*ds
      dssqr=ds*ds
      
C Constants for nondimensionalization of the system
      Fndim=1.E-4
      
C ==============================================================================

C NOTE: In VIS5D (ir:N-->S, ic:W-->E, il:B-->T)
      
C Fill p/ph(Pa), exnf/exnfh(J kg-1 K-1), u(m/s), v(m/s), h(m), t(K), theta(K), 
C cor(s-1) and map(adim.)
       
      do il=1,maxnl
        p(il)=1012.5*exp(vert_args(il)/-7.2)*100.
        exnf(il)=cp*((p(il)/po)**kap)  
      enddo
      do il=1,maxnl-1
        exnfh(il)=0.5*(exnf(il)+exnf(il+1))
        ph(il)=(exnfh(il)*(po**kap)/cp)**(1./kap)
      enddo 
      exnfh(0)=2.*exnf(1)-exnfh(1)
      ph(0)=(exnfh(0)*(po**kap)/cp)**(1./kap)      
      exnfh(maxnl)=2.*exnf(maxnl)-exnfh(maxnl-1)      
      ph(maxnl)=(exnfh(maxnl)*(po**kap)/cp)**(1./kap)
               
      do 10 ir=1,nr
      do 10 ic=1,nc
        cor(ir,ic)=ingrid(ir,ic,1,icor)
        map(ir,ic)=ingrid(ir,ic,1,imap)   
        map2(ir,ic)=map(ir,ic)*map(ir,ic)      
      do 10 il=1,maxnl
        u(ir,ic,il)=ingrid(ir,ic,il,iu)
        v(ir,ic,il)=ingrid(ir,ic,il,iv) 
        h(ir,ic,il)=ingrid(ir,ic,il,ih)                   
        t(ir,ic,il)=ingrid(ir,ic,il,it)+273.16
        theta(ir,ic,il)=t(ir,ic,il)*cp/exnf(il)        
10    continue   

C Remove any unstable layers
      do 15 ir=1,nr
      do 15 ic=1,nc
      do 15 il=2,maxnl
        If (theta(ir,ic,il).Lt.theta(ir,ic,il-1)+1.E-2) 
     $     theta(ir,ic,il)=theta(ir,ic,il-1)+1.E-2
15    continue

C Definition of Neumann boundary conditions thetaB and thetaT (K)
      do 100 ir=1,nr
      do 100 ic=1,nc
        thetaB(ir,ic)=theta(ir,ic,1)
        thetaT(ir,ic)=theta(ir,ic,maxnl)               
100   continue      

C Definition dexnf
      dexnfh(0)=exnfh(1)-exnfh(0)
      do 120 il=1,maxnl-1
       dexnfh(il)=exnfh(il+1)-exnfh(il) 
120   continue      

C ==============================================================================

C If Hm0 available (MEAN fields), retrieve phi0 and skip calculation
      ih0 = -1
      do var=1,nvars
        if (names(var) .eq. 'Hm0') ih0 = var                        
      end do

      If (ih0.Ne.-1) Then       
       Do 11 ir=1,nr
       Do 11 ic=1,nc
         phi0(ir,ic,0)=grav*ingrid(ir,ic,1,ih0)+thetaB(ir,ic)*
     $                 (exnf(1)-exnfh(0))
       Do 11 il=1,maxnl
         phi0(ir,ic,il)=phi0(ir,ic,il-1)+dexnfh(il-1)*
     $                  (grav*ingrid(ir,ic,il,ih0)-
     $                  phi0(ir,ic,il-1))/(exnf(il)-exnfh(il-1))
11     Continue
       print*,'Retrieved PHI0'
       Goto 111 
      Endif 

C CALCULATION OF FIRST-GUESS GEOPOTENTIAL phi0 (m2 s-2)    

      do 30 ir=1,nr
      do 30 ic=1,nc       
        phi0(ir,ic,0)=grav*h(ir,ic,1)+0.5*dexnfh(0)*theta(ir,ic,1)

        do 40 il=1,maxnl
          phi0(ir,ic,il)=phi0(ir,ic,il-1)-dexnfh(il-1)*theta(ir,ic,il)  
40      continue    

30    continue

      print*,'Calculated PHI0'
 
111   CONTINUE 
      
C ==============================================================================

C CALCULATION OF FIRST GUESS STREAMFUNCTION psi0 (m2 s-1)

C Relative vorticity (forcing for_psi) and set initial values
      do 60 ir=2,nr-1
      do 60 ic=2,nc-1
      do 60 il=1,maxnl-1
        dudy=0.5*map(ir,ic)*(u(ir-1,ic,il)+u(ir-1,ic,il+1)-
     $       u(ir+1,ic,il)-u(ir+1,ic,il+1))/ds2     
        dvdx=0.5*map(ir,ic)*(v(ir,ic+1,il)+v(ir,ic+1,il+1)-
     $       v(ir,ic-1,il)-v(ir,ic-1,il+1))/ds2
        relvor=dvdx-dudy
        relvor=Max(relvor,-cor(ir,ic)+1.E-8)              
        for_psi(ir,ic,il)=relvor/map2(ir,ic)                     
        psi0(ir,ic,il)=0.
60    continue      

      DO 65 il=1,maxnl-1

C Set boundary values for each level
       delta=0.
       psiref=phi0(1,1,il)/Fndim
70     continue
       psi0(1,1,il)=psiref 

       do ic=2,nc
         psi0(1,ic,il)=psi0(1,ic-1,il)+
     $                 0.5*ds*(v(1,ic,il)+v(1,ic,il+1)+
     $                         v(1,ic-1,il)+v(1,ic-1,il+1))/
     $                 (map(1,ic)+map(1,ic-1))-
     $                 delta        
       enddo       
       do ir=2,nr
         psi0(ir,nc,il)=psi0(ir-1,nc,il)+
     $                 0.5*ds*(u(ir,nc,il)+u(ir,nc,il+1)+
     $                         u(ir-1,nc,il)+u(ir-1,nc,il+1))/
     $                 (map(ir,nc)+map(ir-1,nc))-
     $                 delta        
       enddo       
       do ic=nc-1,1,-1
         psi0(nr,ic,il)=psi0(nr,ic+1,il)-
     $                 0.5*ds*(v(nr,ic,il)+v(nr,ic,il+1)+
     $                         v(nr,ic+1,il)+v(nr,ic+1,il+1))/
     $                 (map(nr,ic)+map(nr,ic+1))-
     $                 delta        
       enddo       
       do ir=nr-1,1,-1
         psi0(ir,1,il)=psi0(ir+1,1,il)-
     $                 0.5*ds*(u(ir,1,il)+u(ir,1,il+1)+
     $                         u(ir+1,1,il)+u(ir+1,1,il+1))/
     $                 (map(ir,1)+map(ir+1,1))-
     $                 delta        
       enddo   

       If (delta.Eq.0.) Then
        delta=(psi0(1,1,il)-psiref)/Real(2*(nr+nc)-4)
        If (delta.Eq.0.) Goto 75   
        Goto 70
       Endif
       
75     Continue     

C SOR solution for interior values for each level
       Do 80 iter=1,max_iters
       noconverg=0

         Do 85 ir=2,nr-1
         Do 85 ic=2,nc-1
           res=psi0(ir,ic+1,il)+psi0(ir,ic-1,il)+
     $         psi0(ir-1,ic,il)+psi0(ir+1,ic,il)-
     $         4.*psi0(ir,ic,il)-
     $         dssqr*for_psi(ir,ic,il)
           resmod=.25*over_psi*res  
           psi0(ir,ic,il)=psi0(ir,ic,il)+resmod
           If (Abs(resmod).Gt.thrs_psi) noconverg=noconverg+1            
85       Continue 

         If (noconverg.Eq.0) Goto 90
      
80     Continue       

       print*,'Convergence for PSI0(',il,') not reached'
       Goto 65
 
90     Continue 
       print*,'PSI0(',il,') converged at iter=',iter
 
65    CONTINUE

C Finally, top and bottom B.C
      do 95 ir=1,nr
      do 95 ic=1,nc
        psi0(ir,ic,0)=psi0(ir,ic,1)+dexnfh(0)*
     $                thetaB(ir,ic)/Fndim
        psi0(ir,ic,maxnl)=psi0(ir,ic,maxnl-1)-dexnfh(maxnl-1)*
     $                    thetaT(ir,ic)/Fndim
95    continue

C ==============================================================================

C OUTPUT OF FIRST GUESS OF BALANCED STREAMFUNCTION Sm0 (=[Fndim/grav]*psi0, 
C in [grav/Fndim] m2 s-1)

      print*,
     $'Output of first guess of balanced streamfunction ', 
     $'Sm0 ([grav/Fndim] m2 s-1)'
      DO 9999 ir=1,nr
      DO 9999 ic=1,nc
      DO 9999 il=1,maxnl
        outgrid(ir,ic,il)=Fndim*
     $                    ( psi0(ir,ic,il)-(exnfh(il)-exnf(il))*
     $                    (psi0(ir,ic,il)-psi0(ir,ic,il-1))/
     $                    dexnfh(il-1) )/grav
9999  CONTINUE     

C ==============================================================================
               
      USERFUNC = 0
      RETURN
      END

C ==============================================================================
