# = NumRu::GAnalysis::Planet : Library for spherical planets (default: Earth)
#
# ASSUMPTIONS
# * longitude is assumed to increase in the eastward direction.
# * latitude is assumed to increase in the northward direction,
#   and it is zero at the equator.

require "numru/gphys"
require 'numru/gphys/derivative'

module NumRu
  module GAnalysis

    module Planet
      module_function

      #< Pre-defined planets >

      Earth = "Earth"
      def init(planet)
        case planet
        when Earth
          @@R = UNumeric[6.37e6, "m"]
          @@Ome = UNumeric[2*Math::PI/8.64e4,"s-1"]
        else
          raise "Unsupported predefined planet: #{planet}."
        end
      end

      init(Earth)    ###### default setting as the Earth ######

      #< Adjustable planetary settings >

      def radius=(r); @@R = r; end
      def omega=(o);  @@Ome = o; end
      def radius; @@R; end
      def omega; @@Ome; end

      #< Class variables regarding lon & lat >

      @@lonbc = GPhys::Derivative::CYCLIC_OR_LINEAR # this should be always fine
      @@latbc = GPhys::Derivative::LINEAR_EXT       # this should be always fine

      @@lam = nil   # lambda (lon in radian) obtaiend lately (see get_lonlat)
      @@phi = nil   # phi (lat in radian) obtaiend lately (see get_lonlat)

      #< Differentian at the planets's near surface. With suffix "_s" >

      def rot_s(u,v)
        lond, latd, lam, phi = get_lonlat(u)
        cos_phi = phi.cos
        dv_dlam = v.cderiv(lond,@@lonbc,lam)
        duc_dphi  = (u*cos_phi).cderiv(latd,@@latbc,phi)
        rot = (dv_dlam - duc_dphi) / (@@R*cos_phi)
        rot.long_name = "rot(#{u.name},#{v.name})"
        rot.name = "rot"
        rot
      end

      def div_s(u,v)
        lond, latd, lam, phi = get_lonlat(u)
        cos_phi = phi.cos
        du_dlam = u.cderiv(lond,@@lonbc,lam)
        dvc_dphi  = (v*cos_phi).cderiv(latd,@@latbc,phi)
        rot = (du_dlam + dvc_dphi) / (@@R*cos_phi)
        rot.long_name = "div(#{u.name},#{v.name})"
        rot.name = "div"
        rot
      end

      def vor_s(u,v)
        vor = rot_s(u,v)
        vor.long_name = "Relative Vorticity"
        vor.name = "vor"
        vor
      end
      
      def absvor_s(u,v)
        avor = vor_s(u,v) + @@phi.sin * (2*@@Ome)
        avor.long_name = "Absolute Vorticity"
        avor.name = "avor"
        avor
      end

      def grad_s(s)
        lond, latd, lam, phi = get_lonlat(s)
        cos_phi = phi.cos
        xs = s.cderiv(lond,@@lonbc,lam) / (@@R*cos_phi)
        ys = s.cderiv(latd,@@latbc,phi) / @@R
        xs.long_name = "xgrad(#{s.name})"
        xs.name = "xgrad"
        ys.long_name = "ygrad(#{s.name})"
        ys.name = "ygrad"
        [xs,ys]
      end

      #< helper methods >

      def get_lonlat(gp)
        lond, latd = find_lon_lat_dims(gp, true)
        lam = gp.axis(lond).to_gphys.convert_units("rad") # lambda (lon in rad)
        lam.units = ""                                    # treat as non-dim
        phi = gp.axis(latd).to_gphys.convert_units("rad") # phi (lat in rad)
        phi.units = ""                                    # treat as non-dim
        @@lam = lam
        @@phi = phi
        [lond, latd, lam, phi]
      end

      def find_lon_lat_dims(gp, err_raise=false)
        lond = nil
        latd = nil
        (0...gp.rank).each do |dim|
          if /^degrees?_east$/i =~ gp.coord(dim).get_att("units")
            lond = dim
            break
          elsif /^longitude$/i =~ gp.coord(dim).long_name
            lond = dim
            break
          end
        end
        (0...gp.rank).each do |dim|
          if /^degrees?_north$/i =~ gp.coord(dim).get_att("units")
            latd = dim
            break
          elsif /^latitude$/i =~ gp.coord(dim).long_name
            latd = dim
            break
          end
        end
        if err_raise
          raise("Longitude dimension was not found") if !lond
          raise("Latitude dimension was not found") if !latd
        end
        [lond,latd]
      end
    end



  end
end

################################################
##### test part ######
if $0 == __FILE__
  require "numru/ggraph"
  include NumRu
  u = GPhys::IO.open("../../../testdata/UV.jan.nc","U")
  v = GPhys::IO.open("../../../testdata/UV.jan.nc","V")
  ##i = 3
  ##u = GPhys::IO.open("/mnt/disk2/horinout/ncep/daily/pressure/uwnd.2008.nc","uwnd")[{0..-1=>i},{1..-2=>i},true,0]
  ##v = GPhys::IO.open("/mnt/disk2/horinout/ncep/daily/pressure/vwnd.2008.nc","vwnd")[{0..-1=>i},{1..-2=>i},true,0]
  rot = GAnalysis::Planet.rot_s(u,v)
  avor = GAnalysis::Planet.absvor_s(u,v)
  div = GAnalysis::Planet.div_s(u,v)
  print "*** test ***\n"
  p rot
  p rot.units
  p rot.max,rot.min

  DCL.swpset('iwidth',700)
  DCL.swpset('iheight',700)
  DCL.sgscmn(10)  # set colomap
  DCL.gropn(1)
  DCL.sgpset('isub', 96)      # control character of subscription: '_' --> '`'
  DCL.glpset('lmiss',true)
  DCL.sldiv("y",1,2)
  DCL.sgpset("lfull",true)
  GGraph::set_fig "itr"=>10,"viewport"=>[0.1, 0.8, 0.05, 0.4]
  GGraph::set_map "coast_world"=>true
  iz = 4
  GGraph::vector u[false,iz], v[false,iz],true  #,"fact"=>2
  #GGraph::tone u[false,iz],true,"color_bar"=>true
  GGraph::tone rot[false,iz],true,"max"=>2e-5,"min"=>-2e-5,"int"=>2e-6
  ##GGraph::tone rot[false,iz],true,"max"=>1e-4,"min"=>-1e-4,"int"=>1e-5
  GGraph::tone avor[false,iz],true,"max"=>2e-4,"min"=>-2e-4,"int"=>2e-5
  GGraph::color_bar
  GGraph::tone div[false,iz],true,"max"=>0.5e-5,"min"=>-0.5e-5,"int"=>0.5e-6
  GGraph::color_bar
  DCL.grcls
end