module NumRu
  class Grib
    class << self
      # cache due to momery leak in the original grib_api library
      @@gribs = Hash.new
      def is_a_Grib?(fname)
        is = false
        File.open(fname) do |file|
          is = (file.read(4) == "GRIB")
        end
        return is
      end
      alias :_new :new
      def new(fname, mode="r")
        key = [fname,mode]
        if (obj = @@gribs[key])
          obj[1] += 1
          return obj[0]
        else
          grib = _new(fname, mode)
          @@gribs[key] = [grib, 1]
          return grib
        end
      end
      alias :open :new
    end # class << self

    alias :_close :close
    def close
      @@gribs.each do |k, val|
        if self == val[0]
          val[1] -= 1
          if val[1] < 1
            _close
            @@gribs.delete(k)
            return true
          end
        end
      end
      return false
    end

    def inspect
      "Grib: #{path}"
    end
  end # class Grib

  class GribMessage
    def name
      n = get_value("name")
      return n if (n && n != "unknown")
      get_value("parameterName") || get_value("indicatorOfParameter")
    end
    def sname
      sn = get_value("shortName")
      return sn if (sn && sn != "unknown")
      "id" + (get_value("indicatorOfParameter", ::NumRu::Grib::TYPE_LONG) || get_value("parameterNumber", ::NumRu::Grib::TYPE_LONG)).to_s
    end
    def units
      get_value("parameterUnits") || get_value("units")
    end
    alias :unit :units
    def z_type
      zt = get_value("typeOfLevel")
      return zt if zt && zt != "unknown"
      z_sname
    end
    def z_sname
      get_value("indicatorOfTypeOfLevel") || get_value("typeOfLevel") || "level"
    end
    def z_value
      get_value("levels") || get_value("level")
    end
    def z_units
      get_value("pressureUnits")
    end
    def date
      return DateTime.parse(date_raw)
    end
    def date_raw
      d = get_value("dataDate").to_s.rjust(8,"0")
      h = get_value("dataTime").to_s.rjust(4,"0")
      d << h
    end
    def step
      get_value("step")
    end
    def time_interval
      get_value("lengthOfTimeRange")
    end
    def step_units
      get_value("stepUnits")
    end
    MEMBER_NAME = {1 => "", 2 => "n", 3 => "p"}
    def ensemble_member
      pn = get_value("perturbationNumber")
      te = get_value("typeOfEnsembleForecast")
      name = ""
      name << pn.to_s if pn
      name << (MEMBER_NAME[te] || te.to_s) if te
      name
    end
    def gtype
      get_value("typeOfGrid")
    end
    def missing_value
      get_value("missingValue")
    end
    def get_xy
      lo, la, = get_data
      case gtype
      when "regular_ll", "regular_gg"
        lon = Hash.new
        lat = Hash.new
        if get_value("jPointsAreConsecutive") == 1
          lon["ij"] = 1
          lat["ij"] = 0
          shape = [get_value("Nj"), get_value("Ni")]
        else
          lon["ij"] = 0
          lat["ij"] = 1
          shape = [get_value("Ni"), get_value("Nj")]
        end
        lo.reshape!(*shape)
        la.reshape!(*shape)
        lon["short_name"] = "lon"
        lon["long_name"] = "longitude"
        lon["units"] = "degrees_east"
        idx = [0,0]; idx[lon["ij"]] = true
        lon["value"] = lo[*idx]
        lat["short_name"] = "lat"
        lat["long_name"] = "latitude"
        lat["units"] = "degrees_north"
        idx = [0,0]; idx[lat["ij"]] = true
        lat["value"] = la[*idx]
        return [lon,lat]
      else
        raise "not defined yet: #{gtype}"
      end
    end
  end # class GribMessage

  class GribVar
    Day0 = DateTime.new(1900)
    MAXNDIM = 7 # lon, lat, lev, t, step, timeInterval, ensemble
    def rank
      @dims.length
    end
    alias :ndims :rank
    def total
      @dims.inject(1){|t,d| t*d.length}
    end
    def dim_names
      @dims.collect{|d| d.name}
    end
    def shape
      @dims.collect{|d| d.length}
    end
    def dim(index)
      index = dim_names.index(index) if String===index
      return nil if index.nil?
      @dims[index]
    end
    def def_dim(name,index)
      d = ::NumRu::GribDim.new(self,name)
      if index == -1
        @dims.push(d)
      else
        @dims[index] = d
      end
      return d
    end
    def put_att(key,val)
      @attr[key] = val
    end
    alias :set_att :put_att
    def att(key)
      @attr[key]
    end
    def att_names
      @attr.keys
    end
    def typecode
      if missing_value
        NArrayMiss::FLOAT
      else
        NArray::FLOAT
      end
    end
    def missing_value
      @msgs[0].missing_value
    end
    def get(*indices)
      sha = shape
      mask = nil
      first = Array.new(rank-@xy_dims,0)
      if indices.length != 0
        if indices.include?(false)
          sha2 = sha.dup
          sha2.delete(false)
          raise ArgumentError, 'multiple "false" in indices' if sha.length - sha2.length > 1
          indices[indices.index(false)] = [true]*(sha.length-indices.length+1)
          indices.flatten!
        end

        rank.times{|n|
          ind = indices[n]
          case ind
          when true
            indices[n] = 0..sha[n]-1
          when Fixnum
            sha[n] = 1
          when Range
            f = ind.first
            e = ind.end
            e = sha[n]-1 if e==-1
            e -= 1 if ind.exclude_end?
            sha[n] = e-f+1
            indices[n] = f..e
          else
            raise "invalid indices"
          end
        }
        if rank > @xy_dims
          mask = NArray.byte(*shape[@xy_dims..-1])
          mask[*indices[@xy_dims..-1]] = 1
          (rank-@xy_dims).times do |i|
            ind = indices[@xy_dims+i]
            case ind
            when Fixnum
              first[i] = ind
            when Range
              first[i] = ind.first
              first[i] += shape[@xy_dims+i] if first[i] < 0
            end
          end
        end
      end

      shape2 = shape.dup
      sha2 = sha.dup
      @del_dims.each do |i|
        shape2.insert(i,1)
        sha2.insert(i,1)
        first.insert(i-@xy_dims,0)
      end
      mask.reshape!(*shape2[@xy_dims..-1]) if mask
      
      value = missing_value ? NArrayMiss.sfloat(*sha2) : NArray.sfloat(*sha2)
      index = Array.new(MAXNDIM-2)
      @msgs.each_with_index do |msg,i|
        idx = @idx[i]
        next if (mask && mask[*idx]==0)
        val = msg.get_value("values")
        if @xy_dims > 0
          val.reshape!(*shape[0...@xy_dims])
          #        val = msg.get_data[2].reshape!(*shape[0...@xy_dims])
          unless indices.length==0 || indices[0...@xy_dims].inject(true){|t,v| t &&= v==true}
            val = val.slice(*indices[0...@xy_dims])
          end
        end
        (MAXNDIM-2).times do |i| index[i] = idx[i]-first[i] end
        value[*(Array.new(@xy_dims,true)+index)] = val
      end
      ns = sha.length
      ns.times do |ii|
        i = ns-ii-1
        sha.delete_at(i) if indices[i].kind_of?(Fixnum)
      end
      value.reshape!(*sha)
      if missing_value
        value.set_mask value.get_array!.ne(missing_value)
      end
      return value
    end
    alias :[] :get
    alias :val :get
    def inspect
      "GribVar: #{name} in #{file.path}, [#{shape.join(",")}]"
    end
    private
    def get_time(date)
      (DateTime.parse(date) - Day0).to_f*24
    end
    def init
      @attr = Hash.new
      @dims = Array.new
      msgs = get_messages
      msg = msgs[0]
      put_att("long_name", msg.name)
      # put_att("standard_name",std_name)
      put_att("units",msg.units) if msg.units
      put_att("grid_type", msg.gtype)
      put_att("missing_value", [msg.missing_value]) if msg.missing_value
      vdim = Array.new
      xy_dims = 0
      msg.get_xy.sort{|x,y| x["ij"]<=>y["ij"] }.each do |xy|
        val = xy.delete("value")
        sname = xy.delete("short_name")
        ij = xy.delete("ij")
        if val.length > 1
          d = def_dim(sname, -1)
          d.put(val)
          xy.each{|k,v| d.put_att(k,v)}
          xy_dims += 1
        else
          #xy.each{|k,v| va.put_att(k,v)}
        end
      end
      z = Array.new
      t = Array.new
      st = Array.new
      ti = Array.new
      en = Array.new
      hash = Hash.new
      msgs.each_with_index do |msg,i|
        zv = msg.z_value
        tv = msg.date_raw
        stv = msg.step
        tiv = msg.time_interval
        env = msg.ensemble_member
        z.push zv
        t.push tv
        st.push stv
        ti.push tiv
        en.push env
        ary = [zv,tv,stv,tiv,env]
        if hash[ary] # error
          m1 = msgs[hash[ary]]
          m2 = msgs[i]
          a1 = m1.get_keys.map{|k| [k,m1.get_value(k)]}
          a2 = m2.get_keys.map{|k| [k,m2.get_value(k)]}
          a = Array.new
          a1.length.times{|j| a.push [a1[j],a2[j]] if a1[j]!=a2[j]}
          warn "BUG: send the following message to the developers"
          p self
          p ary
          p a
          raise("error")
        end
        hash[ary] = i
      end
      [z, t, st, ti, en].each{|a| a.uniq!}
      #        [z, t, st, ti, en].each{|a| a.uniq!; a.sort!}
      @idx = Array.new(msgs.length)
      hash.each do |ary, i|
        zv, tv, stv, tiv, env = ary
        @idx[i] = [z.index(zv), t.index(tv), st.index(stv), ti.index(tiv), en.index(env)]
      end
      del_dims = Array.new
      if z.length == 1
        put_att("level_type", msg.z_type)
        put_att("level_value", z[0].kind_of?(Numeric) ? [z[0]] : z[0])
        put_att("level_units", msg.z_units) if msg.z_units
        del_dims.push xy_dims
      else
        d = def_dim(msg.z_sname, -1)
        d.put_att("long_name", msg.z_type)
        d.put_att("units", msg.z_units) if msg.z_units
        d.put(NArray.to_na(z))
      end
      if t.length == 1
        put_att("time", msg.date.to_s)
        del_dims.push xy_dims + 1
      else
        d = def_dim("time", -1)
        d.put_att("long_name","time")
        d.put_att("units","hours since #{Day0.strftime('%Y-%m-%d %H:%M:%S')}")
        d.put(NArray.to_na(t.map{|tv|get_time(tv)}))
      end
      if st.length == 1
        put_att("step", [msg.step]) if msg.step
        put_att("step_units", msg.step_units) if msg.step_units
        del_dims.push xy_dims + 2
      else
        d = def_dim("step", -1)
        d.put_att("long_name", "step")
        d.put_att("units", msg.step_units) if msg.step_units
        d.put(NArray.to_na(st))
      end
      if ti.length == 1
        del_dims.push xy_dims + 3
      else
        d = def_dim("time_interval", -1)
        d.put_att("long_name", "time_interval")
        d.put_att("units", msg.step_units) if msg.step_units
        d.put(NArray.to_na(ti))
      end
      if en.length == 1
        del_dims.push xy_dims + 4
      else
        d = def_dim("member", -1)
        d.put_att("long_name", "ensemble_member")
        d.put(NArray.to_na(en))
      end
      @xy_dims = xy_dims
      @del_dims = del_dims
      @msgs = msgs
    end
  end # class GribVar

  class GribDim
    attr_reader :var, :length, :name
    def initialize(var,name)
      @var = var
      @name = name
      @attr = Hash.new
    end
    alias :total :length
    def get
      @ary
    end
    def typecode
      if NArray===@ary
        @ary.typecode
      elsif Array===@ary
        @ary[0]["value"].typecode
      end
    end
    def val
      return @ary
    end
    def [](*ind)
      return val[*ind]
    end
    def put(ary)
      @ary = ary
      @length = val.length
      return @ary
    end
    def put_att(key,val)
      @attr[key]=val
    end
    alias :set_att :put_att
    def att(key)
      @attr[key]
    end
    def att_names
      @attr.keys
    end
    def inspect
      "GribDim: #{name} length=#{length}"
    end
  end # class GribDim

end # module NumRu





#####################################################
if $0 == __FILE__

  include NumRu

  if ARGV.length>0
    infname = ARGV.shift
  else
    infname = "../../../testdata/T.jan.grib"
  end


  Grib.is_a_Grib?(infname) || raise("file is not a Grib dataset")
  p grib = Grib.open(infname)



  print "\nVars\n"
  grib.var_names.each{|vn|
    p v = grib.var(vn)
    p v.dim_names
    v.dim_names.each{|dn| p dn; p v.dim(dn).get }
    p v.shape
    v.att_names.each{|an| print an, " => ", v.att(an), "\n" }
    puts "\n"
    p v.val
  }


end
