require 'getopts'
require 'numru/gphys'
require 'numru/gphys/ep_flux'
include NumRu
include Misc::EMath

module NumRu
  class NetCDFVar

    def get_with_miss_and_scaling2(*args) # make mask before scaling
      __interpret_missing_params if !defined?(@missval)
      packed_data = simple_get(*args)
      scaled_data = scaled_get(*args)
      sf = att('scale_factor')
      ao = att('add_offset')
      if @vmin || @vmax
	if sf && ao
	  csf = sf.get
	  cao = ao.get
	  vmin = (@vmin-cao)/csf if @vmin
	  vmax = (@vmax-cao)/csf if @vmax
	elsif
	  vmin = @vmin; vmax = @vmax
	end
	if vmin
	  mask = (packed_data >= vmin) 
	mask = mask.and(packed_data <= vmax) if vmax
	else
	  mask = (packed_data <= vmax)
	end
	data = NArrayMiss.to_nam(scaled_data, mask)
      elsif @missval	# only missing_value is present.
	eps = 1e-6
	missval = @missval[0].to_f
	vmin = missval - eps
	vmax = missval + eps
	mask = (packed_data <= vmin) 
	mask = mask.or(packed_data >= vmax)
	data = NArrayMiss.to_nam(scaled_data, mask)
      else
	data = scaled_data
      end
      data
    end

    def put_with_miss_after_scaling(data, *args)
      if data.is_a?( NArrayMiss )
	__interpret_missing_params if !defined?(@missval)
	if @missval
	  scaled_put_without_missval(data, *args)
	else
	  scaled_put(data.to_na, *args)
	end
      else
	scaled_put(data, *args)
      end
    end

    def scaled_put_without_missval(var,hash=nil)
      sf = att('scale_factor')
      ao = att('add_offset')
      if ( sf == nil && ao == nil ) then
	# no scaling --> just call put
	simple_put(var,hash)
      else
	if (sf != nil) 
	  csf = sf.get
	  if csf.is_a?(NArray) then  # --> should be a numeric
	    csf = csf[0]
	  elsif csf.is_a?(String)
	    raise TypeError, "scale_factor is not a numeric"
	  end
	  if(csf == 0) then; raise NetcdfError, "zero scale_factor"; end
	else
	  csf = 1.0      # assume 1 if not defined
	end
	if (ao != nil) 
	  cao = ao.get
	if cao.is_a?(NArray) then  # --> should be a numeric
	  cao = cao[0]
	elsif csf.is_a?(String)
	  raise NetcdfError, "add_offset is not a numeric"
	end
	else
	  cao = 0.0      # assume 0 if not defined
	end
	var = var.to_na( @missval[0]*csf + cao)
	simple_put( (var-cao)/csf, hash )
      end
    end  

  end
end



unless getopts("u:", "v:", "omega:", "temp:", 
                                     "temp_is_temperature:true", "output:epflx.nc")
  print "#{$0}:illegal options.\n"
  exit 1
end

z_axs = 'level'
z_range = [100, 1000]

nc_uwnd, var_uwnd =   ($OPT_u).split(/\s*@\s*/)
nc_vwnd, var_vwnd =   ($OPT_v).split(/\s*@\s*/)
nc_omega, var_omega = ($OPT_omega).split(/\s*@\s*/)
nc_temp, var_temp =   ($OPT_temp).split(/\s*@\s*/)

gp_u =     GPhys::IO.open( nc_uwnd,  var_uwnd ).cut(z_axs=>z_range[0]..z_range[-1])
gp_v =     GPhys::IO.open( nc_vwnd,  var_vwnd ).cut(z_axs=>z_range[0]..z_range[-1])
gp_omega = GPhys::IO.open( nc_omega, var_omega).cut(z_axs=>z_range[0]..z_range[-1])
gp_t =     GPhys::IO.open( nc_temp,  var_temp ).cut(z_axs=>z_range[0]..z_range[-1])

ofile = NetCDF.create($OPT_output)


if gp_u.rank == 3
  
  epflx_y, epflx_z, v_rmean, w_rmean, gp_lat, gp_z, = ary =
    GPhys::EP_Flux::ep_full_sphere(gp_u, gp_v, gp_omega, gp_t, true)
  gp_lat.rename('phi')

  ary.each{|gp|                                  #  This part will not 
    gp.data.att_names.each{|nm|                  #  be needed in future.
      gp.data.del_att(nm) if /^valid_/ =~ nm     #  (Even now, it is not
    }                                            #  needed if the valid
  }                                              #  range is wide enough)
  
  epflx_div = GPhys::EP_Flux::div_sphere(epflx_y, epflx_z)

  a       = GPhys::EP_Flux::radius               # get planetaly radius
  h       = GPhys::EP_Flux::scale_height         # get scale height
  cos_phi = cos(gp_lat)                          # cosine phi
  sigma = exp(-gp_z/h)                           # sigma (p/p00)

  GPhys::IO.write(ofile, epflx_y*a)
  GPhys::IO.write(ofile, epflx_z*a)
  GPhys::IO.write(ofile, v_rmean)
  GPhys::IO.write(ofile, w_rmean)
  GPhys::IO.write(ofile, epflx_div/sigma/cos_phi)
  GPhys::IO.write(ofile, gp_lat)
  GPhys::IO.write(ofile, gp_z)

  ofile.close

elsif gp_u.rank > 3
  
  nt = gp_u.shape[-1]
  i = 0
  GPhys::IO.each_along_dims_write([gp_u, gp_v, gp_omega, gp_t], ofile, -1){
    |u, v, omega, t|
    i += 1
    print "processing #{i} / #{nt} ..\n" if (i % (nt/20+1))==1
    
    epflx_y, epflx_z, v_rmean, w_rmean, gp_lat, gp_z, u_mean, theta_mean,
      uv_dash, vt_dash, uw_dash, dtheta_dz = ary =
      GPhys::EP_Flux::ep_full_sphere(u, v, omega, t, true)
    epflx_div = GPhys::EP_Flux::div_sphere(epflx_y, epflx_z)
    
    ary.each{|gp|                                  #  This part will not
      gp.data.att_names.each{|nm|                  #  be needed in future.
	gp.data.del_att(nm) if /^valid_/ =~ nm     #  (Even now, it is not
      }                                            #  needed if the valid
    }                                              #  range is wide enough)
    
    a       = GPhys::EP_Flux::radius               # get planetaly radius
    h       = GPhys::EP_Flux::scale_height         # get scale height
    cos_phi = cos(gp_lat)                          # cosine phi
    sigma = exp(-gp_z/h)                           # sigma (p/p00)
    
    if i==1    # time independent => write only once
      gp_lat.rename('phi')
      GPhys::IO.write(ofile, gp_lat)
      GPhys::IO.write(ofile, gp_z)
    end
    [ epflx_y*a, epflx_z*a, v_rmean, w_rmean, epflx_div/sigma/cos_phi, u_mean, 
      theta_mean, uv_dash, vt_dash, uw_dash, dtheta_dz ]
  }
  
  ofile.close

end