# = NumRu::GAnalysis::Met : Meteorological analysis

require "numru/gphys"
require 'numru/ganalysis/planet'

module NumRu

  module GAnalysis

    # Meteorological analysis
    # 
    # USEFUL METHODS
    # * temp2theta
    # * pv_on_theta
    # * interpolate_onto_theta
    module Met

      #< Themodynamic constants for the Earth's atmosphere >

      R = UNumeric[287.04,"J.kg-1.K-1"] # Gas constant of dry air
      Rv = UNumeric[461.50,"J.kg-1.K-1"] # Gas constant of water vapor
      Cp = UNumeric[1004.6,"J.kg-1.K-1"] # heat capacity at constant pressure for dry air
      Cpv = UNumeric[1870.0,"J.kg-1.K-1"] # heat capacity at constant pressure for water vapor
      Kappa = (R/Cp).to_f  # = 2.0/7.0
      T0 = UNumeric[273.15,"K"]  # K - Celcius
      Lat0 = UNumeric[2.500780e6,"J.kg-1"]  # Latent heat of vaporization at 0 degC

      #< Earth's parameters >
      @@g = UNumeric[9.8,"m.s-2"]    # adjustable: e.g., 9.7 m.s-2 for stratosphere
      P00 = UNumeric[1e5,"Pa"]

      module_function

      #< module parameters >

      # Sets gravity acceleration in the module (default: 9.8 ms-1).
      #
      # ARGUMENT
      # * g [UNumeric]
      # EXAMPLE
      #   Met::set_g( UNumeric[9.7,"m.s-2"] )
      # RETURN VALUE
      # * The argument g
      def set_g(g)
        @@g = g   # e.g., un[9.8,"m.s-2"]
      end


      # Returns gravity acceleration in the module (default: 9.8 ms-1).
      def g
        @@g.dup
      end

      #< potential temperature >

      # Convert temperature into potential temperature
      # 
      # ARGUMENTS
      # * temp [UNumeric or GPhys or VArray, which supports method "units"] :
      #   temperature
      # * prs [UNumeric or GPhys or VArray, which supports method "units"] :
      #   pressure
      # RETURN VALUE
      # * potential temperature [UNumeric or GPhys or VArray,...] in Kelvin
      #
      def temp2theta(temp, prs=nil)
        prs = get_prs(temp) if !prs
        prs = convert_units2Pa(prs)
        if ( !(temp.units === Units["K"]) )
          temp = convert_units(temp,"K")
        end
        theta = temp * (prs/P00)**(-Kappa)
        if theta.respond_to?(:long_name)  # VArray or GPhys
          theta.name = "theta"
          theta.long_name = "potential temperature"
        end
        theta
      end

      #< for isentropic coordinate >

      # Inverse of the "sigma" density in the theta coordinate:
      # 
      # ARGUMENTS
      # * theta [GPhys or VArray] potential temperature
      # * dim [Integer] : the pressure dimension. If theta is a GPhys,
      #   this argument can be omitted, in which case a pressure dimension
      #   is searched internally by find_prs_d.
      # * prs [GPhys or VArray] : the pressure values. Internally searched
      #   if omitted.
      # RETURN VALUE
      # * 1 / sigma = -g d\theta/dp [GPhys]
      #
      def sigma_inv(theta, dim=nil, prs=nil)
        dim = find_prs_d(theta) if !dim
        prs = get_prs(theta) if !prs
        prs = convert_units2Pa(prs)
        #dtheta_dp = df_dx(theta, prs, dim)
        dtheta_dp = df_dx_vialogscale(theta, prs, dim)
        sig_inv = dtheta_dp * (-@@g)
        if sig_inv.respond_to?(:long_name)  # VArray or GPhys
          sig_inv.name = "sig_inv"
          sig_inv.long_name = "1/sigma"
        end
        sig_inv
      end

      # Interpolate onto the potential temperature coordinate
      #
      # ARGUMENTS
      # * gphys [GPhys] a gphys object that have a pressure dimension
      # * theta [GPhys] potential temperature defined on the same grid
      #   as gphys
      # * theta_vals : 1D NArray or Array
      def interpolate_onto_theta(gphys, theta, theta_levs)
        theta_levs = NArray[*theta_levs].to_f if Array === theta_levs
        th_crd = VArray.new( theta_levs, 
               {"units"=>"K", "long_name"=>"potential temperature"}, "theta" )
        gphys.set_assoc_coords([theta])
        pdnm = gphys.coord(find_prs_d(gphys)).name
        gphys.interpolate(pdnm=>th_crd)
      end

      # Derive Ertel's potential vorticity on the theta (isentropic) 
      # coordinate
      #
      # ARGUMENTS
      # * u [GPhys] : zonal wind on pressure coordinate 
      #   (u, v, and theta must share same coordinates)
      # * v [GPhys] : meridional wind on pressure coordinate
      # * theta [GPhys] : potential temperature on pressure coordinate
      # * theta_levs [NArray or Array] : one dimensional array of 
      #   potential temperature values (Kelvin) on which PV is derived.
      #   
      # RETURN VALUE
      # * potential temperature [GPhys] on a theta coordinate, where
      #   levels are set to theta_levs
      def pv_on_theta(u, v, theta, theta_levs)
        sigi = GAnalysis::Met.sigma_inv(theta)
        uth = GAnalysis::Met.interpolate_onto_theta(u, theta, theta_levs)
        vth = GAnalysis::Met.interpolate_onto_theta(v, theta, theta_levs)
        sigith = GAnalysis::Met.interpolate_onto_theta(sigi, theta, theta_levs)
        avorth = GAnalysis::Planet.absvor_s(uth,vth)
        pv = avorth*sigith
        pv.long_name = "potential vorticity"
        pv.name = "PV"
        pv
      end

      # Derive geostrophic wind from geopotential hight (spherical but fixed f)
      # 
      # ARGUMENTS
      # * z [GPhys] : geopotential height on the pressure (or log-pressure)
      #   coordinate
      # * f [nil or Numeric of UNumeric] : the constant f value 
      #   (Coriolis parameter).
      #   If nil, the value at the north pole is assumed.
      #   If Numeric, units are assumed to be "s-1".
      #
      def z2geostrophic_wind(z, f=nil)
        if f.nil?
          f = 2 * GAnalysis::Planet.omega
        elsif f.is_a?(Numeric)
          f = UNumeric[f,"s-1"]
        end
        z = z.convert_units("m")
        gx, gy = GAnalysis::Planet.grad_s( z * (g/f) )
        u = -gy
        v = gx
        u.name = "u"
        u.long_name = "geostrophic U"
        u.put_att("assumed_f",f.to_s)
        v.name = "v"
        v.long_name = "geostrophic V"
        v.put_att("assumed_f",f.to_s)
        [u, v]
      end

      # Adiabatic frontogenesis function over the sphere. --
      # D/Dt(|gradH theta|) or D/Dt(|gradH theta|^2), where gradH express
      # the horizontal component of gradient. 
      # 
      # if full_adv is true (default),
      #  D/Dt(|gradH theta|) = 
      #    [ -(ux-va)*thetax^2 - (vx+uy)*thetax*thetay - vy*thetay^2
      #      - (wx*thetax + wy*thetay)*theta_z ]
      #    / |gradH theta|
      # or else,
      #  (\del/\del t + u gradx + v grady )(|gradH theta|) = 
      #    [ -(ux-va)*thetax^2 - (vx+uy)*thetax*thetay - vy*thetay^2
      #      - (w*theta_z)_x*thetax - (w*theta_z)_y*thetay ]
      #    / |gradH theta|
      # 
      # Here, the 2nd line (vertical advection) is optional;
      # [vx, vy] = gradH v; [thetax, thetay] = gradH theta;
      # [ux, uy] = cos_phi * gradH (u/cos_phi)
      # va = v/tan_phi/a (a=radius).
      # z and w is the vertical coordinate and the lagrangian
      # "velocity" in that coordinate --- Typically they are
      # p and omega, or log-p height and log-p w.
      # 
      # This formulation is adiabatic; the diabatic heating effect
      # can be easily included if needed.
      # 
      # ARGUMENTS
      # * theta [GPhys] : potential temperature
      # * u [GPhys] : zonal wind
      # * v [GPhys] : meridional wind
      # * w [nil (default) or GPhys] : (optional) "vertical wind", which must
      #   be dimensionally consistent with the vertical coordiante 
      #   (e.g., omega for the pressure coordinate). 
      #   If w is given, the vertical cooridnate is assumed to be
      #   the 3rd one (dim=2).
      # * fstodr [true (default) or false] (optional) if true 
      #   D/Dt(|\NablaH theta|) returned; if false D/Dt(|\NablaH theta|^2)
      #   is returned.
      # * full_adv [true (default) or false] : whether to calculate
      #   full lagrangian tendency or lagragian tendency only in 
      #   horizontal direction
      # 
      # RETURN VALUE
      # * frontogenesis function [GPhys]
      # 
      def frontogenesis_func(theta, u, v, w=nil, fstodr=true, full_adv = true)
        thx, thy = GAnalysis::Planet.grad_s( theta )
        ux = GAnalysis::Planet.grad_sx( u )
        uy = GAnalysis::Planet.grad_sy_cosphifact( u, -1 )
        vx, vy =  GAnalysis::Planet.grad_s( v )
        va = GAnalysis::Planet.weight_tanphi( v, 1, -1 )
        frgf = - (ux-va)*thx*thx - (vx+uy)*thx*thy - vy*thy*thy
        frgf.name = "frgen"
        frgf.long_name = "frontogenesis function"
        if w 
          zdim = 2
          if (wun=w.units) !~ (ztun = theta.coord(zdim).units * Units["s-1"])
            raise "w in #{wun} is inconsistent with the vertical coordinate of theta in #{ztun}"
          else
            w = w.convert_units(ztun)    # For example, Pa/s -> hPa/s
          end
          z = theta.axis(zdim).to_gphys
          if z.units =~ Units["Pa"]
            thz = df_dx_vialogscale(theta, z, zdim)
          else
            thz = df_dx(theta, z, zdim)
          end
          if full_adv
            # full lagragian tendency of theta-gradient strength
            wx,  wy =  GAnalysis::Planet.grad_s( w ) 
            frgf -= (wx*thx + wy*thy)*thz
          else
            # lagragian tendency only in horizontal direction
            wthzx,  wthzy =  GAnalysis::Planet.grad_s( w*thz ) 
            frgf -= wthzx*thx + wthzy*thy
          end
        end
        if fstodr
          frgf /= (thx**2 + thy**2).sqrt 
        else
          frgf *= 2
        end
        frgf
      end

=begin
      def frontogenesis_eulerian(theta, u, v, w=nil, fstodr=true )
        thx, thy = GAnalysis::Planet.grad_s( theta )
        uthxx,  uthxy =  GAnalysis::Planet.grad_s( u*thx )
        vthyx,  vthyy =  GAnalysis::Planet.grad_s( v*thy )
        va = GAnalysis::Planet.weight_tanphi( v, 1, -1 )
        frgf = - (uthxx+vthyx)*thx - (uthxy+vthyy)*thy
        frgf.name = "thgrd_tend"
        frgf.long_name = "Eluerian grad-theta tendency"
        if w 
          zdim = 2
          if (wun=w.units) !~ (ztun = theta.coord(zdim).units * Units["s-1"])
            raise "w in #{wun} is inconsistent with the vertical coordinate of theta in #{ztun}"
          else
            w = w.convert_units(ztun)    # For example, Pa/s -> hPa/s
          end
          z = theta.axis(zdim).to_gphys
          if z.units =~ Units["Pa"]
            thz = df_dx_vialogscale(theta, z, zdim)
          else
            thz = df_dx(theta, z, zdim)
          end
          wthzx,  wthzy =  GAnalysis::Planet.grad_s( w*thz ) 
          frgf -= wthzx*thx + wthzy*thy
        end
        if fstodr
          frgf /= (thx**2 + thy**2).sqrt 
        else
          frgf *= 2
        end
        frgf
      end
=end

      # Find a pressure coordinate in a GPhys object
      # 
      # ARGUMENT
      # * gphys [GPhys]
      # * error [nil/false or true] change the behavior if a
      #   pressure coordinate is not found. Default: returns nil;
      #   if error is true, an exception is raised.
      # RETURN VALUE
      # * Integer to indicate the dimension of the pressure coordinate,
      #   or nil if not found by default (see above)
      def find_prs_d(gphys, error=nil)
        pa = Units.new("Pa")
        (gphys.rank-1).downto(0) do |d|
          un = gphys.axis(d).pos.units
          if ( un =~ pa or un.to_s=="millibar" )
            return(d)
          end
        end
        if error
          raise("Could not find a pressure coordinate.")
        else
          nil
        end
      end

      # Find and return a pressure coordinate in a GPhys object
      # 
      # ARGUMENT
      # * gphys [GPhys]
      # 
      # RETURN VALUE
      # * pressure in a 1D GPhys object created from the pressure 
      #   axis in the GPhys object or a UNumeric object if the pressure
      #   coordinated has been deleted by subsetting.
      def get_prs(gphys)
        begin
          prs = gphys.axis( find_prs_d(gphys, true) ).to_gphys
        rescue
          regexp = /([\d\.]+) *(millibar|hPa)/
          cand = gphys.lost_axes.grep(regexp)
          if (str=cand[0])   # substitution, not ==
            regexp =~ str
            hpa = $1.to_f
            prs = UNumeric[hpa*100,"Pa"]
          else
            raise "No pressure axis was found"
          end
        end
        prs
      end

      # Convert units into Pa. To deal with old versions of NumRu::Units 
      # that do not support "millibar".
      #
      # ARGUMENT
      # * prs [GPhys and UNumeric]
      def convert_units2Pa(prs)
        pa = Units["Pa"]
        un = prs.units
        if ((un =~ pa) and !(un === pa))
          prs = convert_units(prs, pa)
        elsif un.to_s=="millibar"
          if UNumeric === prs
            ret = UNumeric[prs.val*100, "Pa"]
          else
            ret = prs*100
            ret.units = "Pa"
          end
          ret
        else
          prs
        end
      end

      #< misc >

      # Unit conversion good to both GPhys and UNumeric
      def convert_units(obj, un)
        begin
          obj.convert_units(un)   # for GPhys etc
        rescue
          obj.convert2(un)        # for UNumeric etc
        end
      end
      private :convert_units   # this method should be defined somewhere else

      # numerical derivative good to both GPhys and UNumeric
      def df_dx(f, x, dim)
        if GPhys === f
          mdl = NumRu::GPhys::Derivative
          mdl.threepoint_O2nd_deriv(f, dim, mdl::LINEAR_EXT, x)
        else
          mdl = NumRu::Derivative
          mdl.threepoint_O2nd_deriv(f, x, dim, mdl::LINEAR_EXT)
        end
      end
      private :df_dx   # this method should be defined somewhere else

      def df_dx_vialogscale(f, x, dim)
        z = Misc::EMath.log(x)
        if GPhys === f
          mdl = NumRu::GPhys::Derivative
          dfdz = mdl.threepoint_O2nd_deriv(f, dim, mdl::LINEAR_EXT, z)
        else
          mdl = NumRu::Derivative
          dfdz = mdl.threepoint_O2nd_deriv(f, z, dim, mdl::LINEAR_EXT)
        end
        dfdz / x
      end

    end
  end
end

################################################
##### test part ######
if $0 == __FILE__
  require "numru/ggraph"
  include NumRu

  #< data >

  temp = GPhys::IO.open("../../../testdata/T.jan.nc","T")
  u = GPhys::IO.open("../../../testdata/UV.jan.nc","U")
  v = GPhys::IO.open("../../../testdata/UV.jan.nc","V")

  #< theta >

  print "***** test theta *****\n\n"

  theta = GAnalysis::Met.temp2theta(temp)
  p theta
  DCL.swpset('iwidth',800)
  DCL.swpset('iheight',800)
  DCL.sgscmn(10)  # set colomap
  DCL.gropn(1)
  DCL.sgpset('isub', 96)      # control character of subscription: '_' --> '`'
  DCL.glpset('lmiss',true)
  DCL.sldiv("y",2,2)
  GGraph::set_fig "itr"=>2
  #GGraph::tone_and_contour temp.mean(0)
  GGraph::tone_and_contour theta.mean(0)

  p temp.coord(2).val[3]
  t = temp[0..1,0,3]
  p t.val
  p GAnalysis::Met.temp2theta(t).val
  p GAnalysis::Met.temp2theta(t,UNumeric[4e4,"Pa"]).val
  p GAnalysis::Met.temp2theta(UNumeric[20,"degC"],UNumeric[850,"hPa"]).val

  #< isentropic coordinate >

  print "\n***** test isentropic coordinate *****\n\n"
  sig_inv = GAnalysis::Met.sigma_inv(theta)
  GGraph::tone_and_contour 1/sig_inv.mean(0),true,"log"=>true,"title"=>"sigma"
  p( (sig_inv.units**(-1)).reduce5.to_s )

  theta_rg = theta.regrid(u)

  theta_levs = [320.0, 350.0, 400.0, 600.0, 800.0]
  u_th = GAnalysis::Met.interpolate_onto_theta(u, theta_rg, theta_levs)
  p u_th

#  DCL.grfrm
  GGraph::tone u.mean(0)
  GGraph::contour theta_rg.mean(0), false, "levels"=>theta_levs,"index"=>[3]
  GGraph::color_bar
  GGraph::set_fig "itr"=>1
  GGraph::tone_and_contour u_th.mean(0),true,"keep"=>true
  GGraph::color_bar

  #< isentropic potential vorticity >
  pv = GAnalysis::Met.pv_on_theta(u, v, theta_rg, theta_levs)
  p pv.units
  GGraph::tone_and_contour pv.mean(0),true
  GGraph::tone pv.cut("theta"=>320),true,"nlev"=>20,"color_bar"=>true
  GGraph::tone pv.cut("theta"=>600),true,"nlev"=>20,"color_bar"=>true

  #< finish >
  DCL.grcls
end
