# = NumRu::GAnalysis::QG : Quasi-geostrophic calculations


require "numru/gphys"
require 'numru/ganalysis/planet'
require 'numru/ganalysis/met'  # for g (gravity)
require 'numru/ganalysis/log_p'
require 'numru/ganalysis/beta_plane'

module NumRu
  module GAnalysis

    # QG_common: collection of common methods for QG, QG_sphere, and QG_sphere_div.
    module QG_common
      ## module_function  # disabled: module functions are specified one-by-one

      # geopotential height (multi-D) -> reference geopotential profile (1D)
      def gph2gpref(gph)
        gp2gpref(gph) * Met::g
      end
      module_function :gph2gpref

      # geopotential (multi-D) -> reference geopotential profile (1D)
      def gp2gpref(gp)
        gpref = Planet::ave_s(gp)  # horizontal ave (spherical)
        if gpref.rank >= 2
          # likely a time sequence. need to reduce more.
          pdim = Met.find_prs_d(gpref)
          idxs = (0...gpref.rank).collect{|i| i}
          idxs.delete(pdim)
          gpref = gpref.mean(*idxs)
        end
        gpref.name = "gpref"
        gpref.long_name = "Reference geopotential"
        gpref
      end
      module_function :gp2gpref

      # geopotential height to geopotential deviation from the global&time mean
      def gph2gpd_gpref(gph)
        gp = gph * Met::g
        gpref = gp2gpref(gp)
        gpd = gp - gpref
        gpd.name = "gpd"
        gpd.long_name = "Geopotential deviation"
        [gpd, gpref]
      end
      module_function :gph2gpd_gpref

      # reference geopotential -> buoyancy frequency squared
      def gpref2n2(gpref)
        gp_z = LogP.pcdata_dz( gpref )
        gp_zz = LogP.pcdata_dz2( gpref )
        gp_zz[0] = gp_zz[1]   # At boundary, it's safer to extend lapse rate
        gp_zz[-1] = gp_zz[-2] # At boundary, it's safer to extend lapse rate
        n2 = gp_zz + gp_z * (Met::Kappa / LogP.h)
        n2.name = "N2"
        n2.long_name = "N**2 (log-p)"
        #p "@@@@@ N2 @@@@",n2.coord(0).val.to_a, n2.val.sqrt.to_a
        n2
      end
      module_function :gpref2n2

      # [ (p/b) gp_z ]_z /p
      def gpd2qzz(gp, b)
        bunits = Units["s-2"]   # this is assumed!
        pdim = Met.find_prs_d(gp)
        p = gp.coord(pdim)
        z = LogP.p2z(p)
        zunits = z.units
        g = Derivative::b_expand_linear_ext( gp.val, pdim )
        z = Derivative::b_expand_linear_ext( z.val, 0 )
        p = Derivative::b_expand_linear_ext( p.val, 0 )
        b = b.val
        b = b.to_na if b.respond_to?(:to_na)  # likely a NArrayMiss
        b = Derivative::b_expand_linear_ext( b, 0 )
        pb = p/b

        pbm = (pb[0..-2] + pb[1..-1]) / 2.0  # pb_{i+1/2} (for i=0..-2)
        pbm01 = pbm[0..-2]                   # pb_{i-1/2} (for i=1..-2)
        pbm12 = pbm[1..-1]                   # pb_{i+1/2} (for i=1..-2)
        dz20 = z[2..-1] - z[0..-3]           # z_{i+1} - z_{i-1} (for i=1..-2)
        dz21 = z[2..-1] - z[1..-2]           # z_{i+1} - z_{i}   (for i=1..-2)
        dz10 = z[1..-2] - z[0..-3]           # z_{i} - x_{i-1}   (for i=1..-2)
        pc = p[1..-2]                        # p_{i}   (for i=1..-2)

        a2 = 2*pbm12/(dz21*dz20)/pc
        a0 = 2*pbm01/(dz10*dz20)/pc
        a1 = -a2 - a0

        to_rankD = [1]*pdim + [true] + [1]*(gp.rank-pdim-1)
        a2 = a2.reshape(*to_rankD)
        a1 = a1.reshape(*to_rankD)
        a0 = a0.reshape(*to_rankD)

        vqzz =   g[ *([true]*pdim+[2..-1,false]) ] * a2 \
               + g[ *([true]*pdim+[1..-2,false]) ] * a1 \
               + g[ *([true]*pdim+[0..-3,false]) ] * a0

        qzz = gp.copy
        qzz.data.replace_val(vqzz)
        qzz.name = "qzz"
        qzz.long_name = "z-deriv term in QG PV"
        qzz.units = qzz.units / zunits**2 / bunits
        qzz
      end
      module_function :gpd2qzz


=begin
      def gpd2qzz(gp, b)
        pdim = Met.find_prs_d(gp)
        p = gp.axis(pdim).to_gphys
        gp_z = LogP.pcdata_dz( gp )
        qzz = LogP.pcdata_dz( gp_z * (p/b) ) / p
        qzz.name = "qzz"
        qzz.long_name = "z-deriv term in QG PV"
        qzz
      end
=end

      # Extend the bottom pressure level by the lowest thickness
      # (a hypothetical "Under-ground" level is created)
      # If value of the extended bottom level is set to
      # val_extended (Numeric or NArray etc), if it is specified (non nil).
      # If nil, the value at the original bottom level is simply copied.
      def extend_bottom(z, val_extended=nil)
        pdim = Met.find_prs_d(z)
        plev = z.coord(pdim).val
        raise("Only one pressure level is found; 2 or more needed") if (plev==1)
        bottom_first = ( plev[0] - plev[1] > 0 )
        np = z.shape[pdim]
        idx = (0...np).collect{|i| i}
        if bottom_first       # The first level is the bottom one
          idx.unshift(0)      # idx => [0,0,1,2,...,np-1]
          ihb = 0             # index of the extended bottom level
          dp = plev[0] - plev[1]
          phb = plev[0] + dp  # pressure of the extended bottom level
        else                # The last level is the bottom one
          idx.push(np-1)     # idx => [0,1,2,...,np-1,np-1]
          ihb = np            # index of the extended bottom level
          dp = plev[-1] - plev[-2]
          phb = plev[-1] + dp  # pressure of the extended bottom level
        end
        ze = z[ *([true]*pdim + [idx,false]) ].copy  # add one level below
        ze.coord(pdim)[ihb] = phb
        if val_extended
          ze[ *([true]*pdim + [ihb,false]) ] = val_extended
        end
        ze
      end
      module_function :extend_bottom

      def cut_bottom(z)
        pdim = Met.find_prs_d(z)
        plev = z.coord(pdim).val
        if plev[0] - plev[1] > 0
          z[ *([true]*pdim + [1..-1,false]) ]
        else
          z[ *([true]*pdim + [0..-2,false]) ]
        end
      end
      module_function :cut_bottom


      ######################################################
      ######################################################
      # The following are instance methods
      # for use (or "inherited") in QG or QG_sphere or QG_sphere_div
      # (not module functions)
      ######################################################


      # div of WAF 
      # 
      # (p cos_phi)^-1 div(p waf) =
      # (cos_phi)^-1 ( div_h(fx,fy) + p^-1 d_z (p fz) )
      #
      # * fx, fy, fz (GPhys) : the x, y and z components of waf
      # * bottom_treatment (true (==default) or false) :
      #   If true, the lowest level vertical divergence is 
      #   computed by assuming that fz is zero at the extended
      #   "underground" level. The thickness assumed (=p[1]-p[0]) is
      #   consistent with the ((<extend_bottom>)) method.
      def div_waf(fx, fy, fz, bottom_treatment=true)
        cosphi = cos_phi(fx)
        p = Met.get_prs(fx)

        fz_z = LogP.pcdata_dz( fz*p ) / p

        #>>>>>> the lowest layer treatment consistent with qb, in which
        # geopotential (or stream function) is extended by extend_bottom. 
        # Assumption: the first level is the lowest (bottom) one
        if bottom_treatment
          # using the relation p^{-1} d/dz = -H^{-1} d/dp 
          # and assuming fz=0 below the bottom (the "underground" level),
          # p^{-1} d/dz (p fz) = -H^{-1} d/dp (p fz),
          # which is H^{-1} p fz / delta_p, at the lowest level with a
          # "thickness" of delta_p. 
          w = p[0..1].val
          dp = w[1] - w[0]
          p0 = w[0]
          pdim = Met.find_prs_d(fz)
          sel0 = [true]*pdim + [0,false]  # to specify the first level
          fz_z[*sel0] = fz[*sel0]*p0 / (LogP.h*dp)
        end
        #<<<<<<

        divh = ( div_h(fx, fy) + fz_z ) / cosphi
               # ^ div_h is defined in QG, QG_sphere,..., but not in QG_common
        divh.name = "divwaf"
        divh.long_name = "div of waf (#{fx.name},..)"
        divh
      end


   end

    ######################################################

    # Collection of common methods for QG_sphere and QG_sphere_div
    module QG_sphere_common

      # Coriolis parameter f
      def f(gphys)
        lam, phi, = Planet::get_lambda_phi(gphys)
        f = phi.sin * (2*Planet::omega)
        f.name = "f"
        f.long_name = "Coriolis parameter"
        f
      end

      # mask where f=0
      def f_mask0(gphys)
        f = f(gphys)
        v = f.val
        vm = NArrayMiss.to_nam(v, v.ne(0))
        f.replace_val(vm)
        f
      end
    end

    ######################################################

    # module QG: quasi-geostrophic analysis module for Cartesian coordinates
    module QG

      module_function
      extend QG_common

      # This class is for internal use only
      class Uninitialized
        def method_missing(method_name)
          raise("Reference latitude has not been set. Call QG::set_lat0 to use the module QG.")
        end
      end

      @@bp = Uninitialized.new   # a BetaPlane to be initialized by set_lat0

      # Initialize the QG module by setting a reference latitude.
      def set_lat0(lat0_or_latary)
        @@bp = BetaPlane.new(lat0_or_latary)
      end

      # returns the BetaPlane object created by initialization (((<set_lat0>)))
      def bp
        @@bp
      end

      # Returns the current f0 (the Coriolis parameter at the reference latitude)
      def f0; @@bp.f0; end

      #def get_x_y(gphys); @@bp.get_x_y(gphys); end

      # geopotential height to quasi-geostrophic potential vorticity (QGPV)
      def gph2q(gph)
        psi, gpref = gph2psi_gpref(gph)
        n2 = gpref2n2(gpref)
        psi2q(psi, n2)
      end

      # same as gph2q, but the QGPV is extended to reflect the lowest-level temperature anomalies
      def gph2qb(gph)
        psi, gpref = gph2psi_gpref(gph)
        n2 = gpref2n2(gpref)
        psi2qb(psi, n2)
      end

      # geopotential height -> geostrophic winds
      def gph2ug_vg(gph)
        psi, gpref = gph2psi_gpref(gph)
        psi2ug_vg(psi)
      end

      # geopotential height -> QG stream function and the reference geopotential
      def gph2psi_gpref(gph)
        gpd, gpref = gph2gpd_gpref(gph)
        psi = gpd / @@bp.f0
        psi.name = "psi"
        psi.long_name = "QG stream function"
        [psi, gpref]
      end

      # geopotential height -> QG stream function
      def gph2psi(gph, gpref)
        gpd = gph * Met::g - gpref
        psi = gpd / @@bp.f0
        psi.name = "psi"
        psi.long_name = "QG stream function"
        psi
      end

      # QG stream function -> QGPV
      def psi2q(psi, n2, perturbation=false)
        x, y = @@bp.get_x_y(psi)
        bc = GPhys::Derivative::CYCLIC_OR_LINEAR
        f0 = @@bp.f0
        vor = psi.deriv2nd(0,bc,x) + psi.deriv2nd(1,bc,y)
        if !perturbation
          avor = vor + (f0 + @@bp.beta*y)
          avor.name = "qgavor"
          avor.long_name = "QG abs vor"
        else
          vor.name = "qgvor"
          vor.long_name = "QG vorticity"
          avor = vor
        end
        qzz = gpd2qzz(psi, n2) * (f0*f0)
        q = avor + qzz
        q.name = "q"
        q.long_name = "QG PV"

        [q, avor, qzz]
      end

      # same as psi2q, but the QGPV is extended to reflect the lowest-level temperature anomalies
      def psi2qb(psi, n2, perturbation=false)
        psie = extend_bottom(psi, nil)
        n2e = extend_bottom(n2, nil)
        results = psi2q(psie, n2e, perturbation)
        results.collect{|z| cut_bottom(z)}
      end

      # QG stream function -> geostrophic winds
      def psi2ug_vg(psi)
        bc = GPhys::Derivative::CYCLIC_OR_LINEAR
        x, y = @@bp.get_x_y(psi)
        vg = psi.cderiv(0,bc,x)
        ug = -psi.threepoint_O2nd_deriv(1,bc,y)
        ug.name = "ug"
        vg.name = "vg"
        ug.long_name = "ug"
        vg.long_name = "vg"
        [ug, vg]
      end

      # QG stream function -> the Q-vector
      def psi2Qvector(psi)
        bc = GPhys::Derivative::CYCLIC_OR_LINEAR
        f0 = @@bp.f0
        x, y = @@bp.get_x_y(psi)
        p = Met.get_prs(psi).convert_units("Pa")
        psi_x = psi.threepoint_O2nd_deriv(0,bc,x)
        psi_y = psi.threepoint_O2nd_deriv(1,bc,y)
        psi_xp = psi_x.threepoint_O2nd_deriv(2,bc,p)
        psi_yp = psi_y.threepoint_O2nd_deriv(2,bc,p)
        psi_xy = psi_x.threepoint_O2nd_deriv(1,bc,y)
        psi_xx = psi.deriv2nd(0,bc,x)
        psi_yy = psi.deriv2nd(1,bc,y)
        q1 = (-psi_xy*psi_xp + psi_xx*psi_yp) * f0
        q2 = (-psi_yy*psi_xp + psi_xy*psi_yp) * f0
        q1.name = q1.long_name = "Q1"
        q2.name = q2.long_name = "Q2"
        [q1,q2]
      end

      # same as psi2Qvector, but temperature is given independently
      #
      # p (nil or UNumeric or VArray or..) : specify pressure if the input data 
      # does not have a pressure axis
      def psi_T2Qvector(psi, temp, p=nil)
        bc = GPhys::Derivative::CYCLIC_OR_LINEAR
        f0 = @@bp.f0
        x, y = @@bp.get_x_y(psi)
        if p
          if p.respond_to?(:convert_units)
            p = p.convert_units("Pa")
          else  
            # UNumeric
            p = p.convert2("Pa")
          end
        else
          p = Met.get_prs(psi).convert_units("Pa")
        end
        psi_xy = psi.threepoint_O2nd_deriv(0,bc,x).threepoint_O2nd_deriv(1,bc,y)
        psi_xx = psi.deriv2nd(0,bc,x)
        psi_yy = psi.deriv2nd(1,bc,y)
        t_x = temp.threepoint_O2nd_deriv(0,bc,x)
        t_y = temp.threepoint_O2nd_deriv(1,bc,y)
        q1 = (psi_xy*t_x - psi_xx*t_y) * (Met::R / p)
        q2 = (psi_yy*t_x - psi_xy*t_y) * (Met::R / p)
        #puts "@@@ psi_T2Qvector @@@", psi.units, psi_xx.units, t_x.units, q1.units 
        q1.name = q1.long_name = "Q1"
        q2.name = q2.long_name = "Q2"
        [q1,q2]
      end

      # horizontal gradient (Cartesian)
      def grad_h(gphys)
        @@bp.grad_h(gphys)
      end

      # horizontal divergence (Cartesian)
      def div_h(gphys_u, gphys_v)
        @@bp.div_h(gphys_u, gphys_v)
      end

    end

    ######################################################
    # QG on sphere with non-divergent but inaccurate geostrophic wind
    # 
    ######################################################
    module QG_sphere
      module_function
      extend QG_common
      extend QG_sphere_common

      # geopotential height to quasi-geostrophic potential vorticity (QGPV)
      def gph2q(gph)
        psi, gpref = gph2psi_gpref(gph)
        n2 = gpref2n2(gpref)
        psi2q(psi, n2)
      end

      # same as gph2q, but the QGPV is extended to reflect the lowest-level temperature anomalies
      def gph2qb(gph)
        psi, gpref = gph2psi_gpref(gph)
        n2 = gpref2n2(gpref)
        psi2qb(psi, n2)
      end

      # geopotential height -> geostrophic winds
      def gph2ug_vg(gph)
        psi, gpref = gph2psi_gpref(gph)
        psi2ug_vg(psi)
      end

      # geopotential height -> QG stream function and the reference geopotential
      def gph2psi_gpref(gph)
        gpd, gpref = gph2gpd_gpref(gph)
        f = f_mask0(gph)
        psi = gpd / f
        psi.name = "psi"
        psi.long_name = "QG stream function"
        [psi, gpref]
      end

      # geopotential height -> QG stream function
      def gph2psi(gph, gpref)
        gpd = gph * Met::g - gpref
        f = f_mask0(gph)
        psi = gpd / f
        psi.name = "psi"
        psi.long_name = "QG stream function"
        psi
      end

      # QG stream function -> QGPV
      def psi2q(psi, n2, perturbation=false)
        ug, vg = psi2ug_vg(psi)

        if !perturbation
          avor = Planet::absvor_s(ug,vg)
          avor.name = "qgavor"
          avor.long_name = "QG abs vor"
        else
          vor = Planet::vor_s(ug,vg)
          vor.name = "qgvor"
          vor.long_name = "QG vorticity"
          avor = vor
        end

        f = f_mask0(psi)
        qzz = gpd2qzz(psi, n2) * (f*f)

        q = avor + qzz
        q.name = "q"
        q.long_name = "QG PV"

        [q, avor, qzz, ug, vg]
      end

      # cosine of latitude
      def cos_phi(gphys)
        lam, phi, = Planet::get_lambda_phi(gphys)
        phi.cos
      end

      #########
      # same as psi2q, but the QGPV is extended to reflect the lowest-level temperature anomalies
      def psi2qb(psi, n2, perturbation=false)
        psie = extend_bottom(psi, nil)
        n2e = extend_bottom(n2, nil)
        results = psi2q(psie, n2e, perturbation)
        results.collect{|z| cut_bottom(z)}
      end

      # QG stream function -> geostrophic winds
      def psi2ug_vg(psi)
        f = f_mask0(psi)
        gpx, gpy = Planet::grad_s(psi)
        vg = gpx
        ug = -gpy
        ug.name = "ug"
        vg.name = "vg"
        ug.long_name = "ug"
        vg.long_name = "vg"
        [ug, vg]
      end

      # horizontal gradient (spherical)
      def grad_h(gphys)
        Planet::grad_s(gphys)
      end

      # horizontal divergence (spherical)
      def div_h(fx, fy)
        Planet::div_s(fx, fy)
      end


      #############################################
      # wave activity flux

      # divergence of wave activity flux (redirected to ((<div_waf>)))
      def self.div_waf(*args)
        super(*args)    # defined in QG_common
      end

      # Flux of the pseudo-momentum in x direction by Plumb (1986).
      # Specifically,  B_2j in Eq.(2.9), but without the factor p.
      # This flux is relative to the mean flow.
      # Averaged over time (if the data is 4D).
      #
      def waf_plumb1986_B2(psi, n2)
        psi_x, psi_y = Planet::grad_s(psi)
        psi_z = LogP.pcdata_dz( psi )
        f2 = f_mask0(psi) ** 2
        cosphi = cos_phi(psi)
        fx = (psi_x**2 - psi_y**2 - psi_z**2 * f2 / n2) * cosphi / 2.0
        fy = psi_x * psi_y * cosphi 
        fz = psi_x * psi_z * (cosphi * f2) / n2
        fx.name = "B21"
        fy.name = "B22"
        fz.name = "B23"
        fx.long_name = "WAF B2x"
        fy.long_name = "WAF B2y"
        fz.long_name = "WAF B2z"
        if psi.rank >= 4
          fx = fx.mean(-1)
          fy = fy.mean(-1)
          fz = fz.mean(-1)
        end
        [fx, fy, fz]
      end

      # Flux of the pseudo-momentum in y direction by Plumb (1986).
      # Specifically,  B_1j in Eq.(2.9), but without the factor p.
      # This flux is relative to the mean flow.
      # Averaged over time (if the data is 4D).
      #
      def waf_plumb1986_B1(psi, n2)
        psi_x, psi_y = Planet::grad_s(psi)
        psi_z = LogP.pcdata_dz( psi )
        f2 = f_mask0(psi) ** 2
        cosphi = cos_phi(psi)
        fx = -psi_x * psi_y * cosphi
        fy = (psi_x**2 - psi_y**2 + psi_z**2 * f2 / n2) * cosphi / 2.0
        fz = -psi_y * psi_z * (cosphi * f2) / n2
        fx.name = "B_11"
        fy.name = "B_12"
        fz.name = "B_13"
        fx.long_name = "x-comp of Px flux Plumb86"
        fy.long_name = "y-comp of Px flux Plumb86"
        fz.long_name = "z-comp of Px flux Plumb86"
        if psi.rank >= 4
          fx = fx.mean(-1)
          fy = fy.mean(-1)
          fz = fz.mean(-1)
        end
        [fx, fy, fz]
      end
    end

    ######################################################
    # QG on sphere with divergence in the geostrophic wind
    #
    # The geostrophic wind is defined as
    # 
    #    [ug, vg] = 1/f curl gpd,
    # 
    # where gpd is the geostrophic geopotential, 
    # which is the deviation of geopotential from some globally uniform
    # background. A background can be defined as the global
    # mean geopotential, as is done in the method gph2gpd_gpref.
    # 
    ######################################################
    module QG_sphere_div
      module_function
      extend QG_common
      extend QG_sphere_common

      # geopotential height to quasi-geostrophic potential vorticity (QGPV)
      def gph2q(gph)
        gpd, gpref = gph2gpd_gpref(gph)
        n2 = gpref2n2(gpref)
        gpd2q(gpd, n2)
      end

      # same as gph2q, but the QGPV is extended to reflect the lowest-level temperature anomalies
      def gph2qb(gph)
        gpd, gpref = gph2gpd_gpref(gph)
        n2 = gpref2n2(gpref)
        gpd2qb(gpd, n2)
      end

      # geopotential height -> geostrophic winds
      def gph2ug_vg(gph)
        gpd, gpref = gph2gpd_gpref(gph)
        gpd2ug_vg(gpd)
      end

      #########
      def gpd2q(gpd, n2)
        ug, vg = gpd2ug_vg(gpd)

        avor = Planet::absvor_s(ug,vg)
        avor.name = "qgavor"
        avor.long_name = "QG abs vor"

        f = f_mask0(gpd)
        qzz = gpd2qzz(gpd, n2) * f

        q = avor + qzz
        q.name = "q"
        q.long_name = "QG PV"

        [q, avor, qzz, ug, vg]
      end

      # geopotential height (gpd: deviation from the reference profie) to the extended QGPV 
      #  
      # qb: q including the contribution from the bottom boundary
      def gpd2qb(gpd, n2)
        gpde = extend_bottom(gpd, nil)
        n2e = extend_bottom(n2, nil)
        results = gpd2q(gpde, n2e)
        results.collect{|z| cut_bottom(z)}
      end

      def gpd2ug_vg(gpd)
        f = f_mask0(gpd)
        gpx, gpy = Planet::grad_s(gpd)
        vg = gpx/f
        ug = gpy* (-1/f)
        ug.name = "ug"
        vg.name = "vg"
        ug.long_name = "ug"
        vg.long_name = "vg"
        [ug, vg]
      end

    end

  end
end


########################################
##### test part ######
if $0 == __FILE__
  include NumRu
  GAnalysis::QG.set_lat0(45.0)
  p GAnalysis::QG.f0
end