# = 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)
        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

      # 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

    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