require "numru/gphys"
require 'numru/gphys/derivative'
require "numru/ganalysis/met"

module NumRu
  module GAnalysis

    # module LogP : log-pressure coordinate support
    #
    module LogP
      module_function

      # P00 = 1e5 Pa (1000 hPa)
      P00 = UNumeric[1e5,"Pa"]

      @@H = UNumeric[7e3,"m"]

      # set the constant log-p scale height (default 7e3 m)
      #
      # Argument
      # * h [Unumeric]
      def set_h(h)
        @@H = h
      end

      # returns the current value of the log-p scale height (in UNumeric)
      def h
        @@H.clone
      end

      # convert pressure to log-p height
      def p2z(p)
        z = (p/P00).convert_units("1").log * (-@@H)
        z.name = "z"
        z.long_name = "log-p z"
        z.del_att("positive")
        z
      end

      # del/delz applied to data with the p coordinate
      def pcdata_dz(gphys, pdim=nil)
        pdim = Met.find_prs_d(gphys) if !pdim
        p = gphys.axis(pdim).to_gphys
        z = p2z(p.data)
        bc = Derivative::LINEAR_EXT
        delz = gphys.threepoint_O2nd_deriv(pdim, bc, z)
        delz.name = gphys.name+"_z"
        delz.long_name = "del/delz (#{gphys.long_name})"
        delz
      end

      # del2/del2z applied to data with the p coordinate
      def pcdata_dz2(gphys, pdim=nil)
        pdim = Met.find_prs_d(gphys) if !pdim
        p = gphys.axis(pdim).to_gphys
        z = p2z(p.data)
        bc = Derivative::LINEAR_EXT
        delz = gphys.deriv2nd(pdim, bc, z)
        delz.name = gphys.name+"_zz"
        delz.long_name = "del2/del2z (#{gphys.long_name})"
        delz
      end

      # p^{-1} del_z ( p del_z self )
      # where del_z = del/del z
      def pcdata_pi_dz_p_dz(gphys, pdim=nil)
        pdim = Met.find_prs_d(gphys) if !pdim
        p = gphys.axis(pdim).to_gphys
        z = p2z(p.data)
        bc = Derivative::LINEAR_EXT
        p_dz = gphys.threepoint_O2nd_deriv(pdim, bc, z) * p
        pi_dz_p_p_dz = p_dz.threepoint_O2nd_deriv(pdim, bc, z) / p
        pi_dz_p_p_dz.long_name = "1/p del_z (p del_z #{gphys.name})"
        pi_dz_p_p_dz
      end

    end
  end

  class GPhys
    # Convert the pressure coordinate in self to log-pressure height (after duplicating self)
    # 
    # Return value: a GPhys
    def logp_coord_p2z(pdim=nil)
      pdim = GAnalysis::Met.find_prs_d(self) if !pdim
      p = self.coord(pdim)
      z = GAnalysis::LogP.p2z(p)
      ax = self.axis(pdim).copy
      ax.set_pos(z)
      ax.name = z.name
      grid = self.grid.copy.set_axis(pdim, ax)
      GPhys.new(grid,self.data)
    end

  end
end


########################################
##### test part ######
if $0 == __FILE__
  require "numru/ggraph"
  include NumRu
  u = GPhys::IO.open("../../../testdata/UV.jan.nc","U")
  u2 = u.logp_coord_p2z
  p u2.axnames
  puts u2.coord(2).units
  DCL.swpset('iwidth',1000)
  DCL.swpset('iheight',500)
  DCL.gropn(4)
  DCL.sldiv('y',2,1)
  DCL.sgpset('isub', 96)      # control character of subscription: '_' --> '`'
  DCL.glpset('lmiss',true)
  GGraph.contour u2[0,false]
  u_zz = GAnalysis::LogP.pcdata_pi_dz_p_dz(u[0,false])
  GGraph.tone u_zz
  GGraph.contour u[0,false], false
  GGraph.color_bar
  puts u2.units, u_zz.units

  #sig = (u.axis(2).to_gphys / GAnalysis::LogP::P00).convert_units("1")
  z = GAnalysis::LogP.p2z(u.axis(2).to_gphys)
  p z.val
  y = u[0,false] * 0 - z
  y_zz = GAnalysis::LogP.pcdata_pi_dz_p_dz(y)
  GGraph.line y_zz[0,false], true, "exchange"=>true
  GGraph.color_bar

  DCL.grcls
end