require "numru/ggraph"
require "yaml"

module NumRu
  class VizShot

    @@ms_windows = false
    def self.windows
      @@ms_windows = true
    end

    @@dumpdir = './'
    def self.dumpdir=(dirname)
      @@dumpdir = dirname.sub(/([^\/])$/,'\1/')   # to end with '/'
    end

    @@basename_default = 'vizshot_dump'   ## DCL::swcget('FNAME').strip 
    def self.basename_default=(basename)
      @@basename_default = basename
    end

    ONED_METHODS = [:plot_1d, :line, :mark]
    TWOD_METHODS = [:tone_cont, :vector, :tone, :contour]

    def initialize(opt = nil)
      opt = opt ? opt.dup : Hash.new
      @set = Hash.new
      @set[:admin] = Hash.new
      @set[:admin][:iwidth]  = ( opt.delete(:iwidth)  || 700 )
      @set[:admin][:iheight] = ( opt.delete(:iheight) || 700 )
      @set[:admin][:xdiv]    = ( opt.delete(:xdiv) || 1 )
      @set[:admin][:ydiv]    = ( opt.delete(:ydiv) || 1 )
      @set[:admin][:basename] = ( opt.delete(:basename) || @@basename_default)

      if opt.length > 0
	raise ArgumentError,"Unsupported option(s) #{opt.keys.inspect}" 
      end

      @plots = Array.new
    end

    def plot(arg)
      arg[:file] = File.expand_path(arg[:file]) if arg[:file] && arg[:file].is_a?(String)
      arg[:file2] = File.expand_path(arg[:file2]) if arg[:file2] && arg[:file2].is_a?(String)
      @plots.push(arg) 
      @plots.length - 1      # index of the current plot
    end

    def add(vizshot)
      viz_new = self.semi_deep_clone
      vizshot.plots.each{|plot| viz_new.plot(plot)}
      viz_new
    end

    def replace_plot!(index, arg)
      @plots[index].update(arg)
    end

    def set_fig(opt)
      @set[:fig] = opt
    end

    def set_map(opt)
      @set[:map] = opt
    end

    def set_axes(opt)
      @set[:axes] = opt
    end

    def set_tone(opt)
      @set[:tone] = opt
    end

    def set_contour(opt)
      @set[:contour] = opt
    end

    def execute(opt=nil)
      if opt  # must be a Hash
	opt = opt.dup
	postscript = opt.delete(:postscript)
	image_dump = opt.delete(:image_dump)
	if opt.length > 0
	  raise ArgumentError,"Unsupported option(s) #{opt.keys.inspect}" 
	end
      end

      setup(postscript,image_dump)
      first = true
      @plots.each do |p| 
        exec_plot(p, first)
        first = false
      end
      finish
    end

    def dump_code(basename=nil, all_in_one=false, exec_opt=nil)
      @set[:admin][:basename] = basename if basename
      path = @@dumpdir + @set[:admin][:basename] + '.rb'
      File.open(path,'w'){|f| f.print(gen_code(all_in_one, exec_opt))}
      path
    end

    def dump_code_and_data(basename=nil, all_in_one=false, exec_opt=nil)
      @set[:admin][:basename] = basename if basename
      viz = self.dup
      data_paths = viz.cut_out_data(@set[:admin][:basename])
      code_path = viz.dump_code(basename, all_in_one, exec_opt)
      [code_path] + data_paths
    end

    def gen_code(all_in_one = false, exec_opt=nil)
      if all_in_one
        libsrc = File.open(__FILE__)
        code = ""
        while (line=libsrc.gets)
          break if /^#.*ENDOFLIB/ =~ line
          code << line
        end
        code << "####################\n"
      else
        code = 'require "numru/vizshot"'
      end
      code += "\nset = <<YML\n\n#{@set.to_yaml}\nYML\n\n"
      code += "\nplots = <<YML\n\n#{@plots.to_yaml}\nYML\n\n"
      ecode = <<-EOS
        plots = YAML.load(plots)
        set = YAML.load(set)
        viz = NumRu::VizShot.new(set[:admin])
        viz.set_fig(set[:fig]) if set[:fig]
        viz.set_map(set[:map]) if set[:map]
        viz.set_axes(set[:axes]) if set[:axes]
        viz.set_contour(set[:contour]) if set[:contour]
        viz.set_tone(set[:tone]) if set[:tone]
        plots.each{|p| viz.plot(p)}
        viz.execute(#{exec_opt.inspect})
      EOS
      ecode.gsub!(/^        /,'')
      code << ecode
    end

    ###################################################
    protected

    def semi_deep_clone
      viz_new = NumRu::VizShot.new(@set[:admin].dup)
      viz_new.set_fig(@set[:fig]) if @set[:fig]
      viz_new.set_map(@set[:map]) if @set[:map]
      viz_new.set_axes(@set[:axes]) if @set[:axes]
      viz_new.set_contour(@set[:contour]) if @set[:contour]
      viz_new.set_tone(@set[:tone]) if @set[:tone]
      @plots.each{|plot| viz_new.plot(plot.dup) }
      return viz_new
    end

    def plots
      @plots
    end

    # for dump_code_and_data
    def cut_out_data(basename)
      newplots = Array.new
      data_paths = Array.new
      @plots.each_with_index do |pl, i|
	gphys, gphys2 = get_gphyses(pl.dup)
	cut = pl[:cut]
	slice = pl[:slice]
	cut2 = pl[:cut2] || cut
	slice2 = pl[:slice2] || slice
	method = pl[:method]
	if !cut && !slice && \
	    ( ( ONED_METHODS.include?(method) && gphys.rank==1 ) || \
	      ( TWOD_METHODS.include?(method) && gphys.rank==2 ) )
	  # no slicing needed
	  newplots[i] = pl
	else
	  filename = basename + sprintf("_%03d",i) + '.nc'
	  path = @@dumpdir + filename
	  pl = pl.dup
	  pl[:file] = filename
	  if ONED_METHODS.include?(method)
	    gphys = gphys.first1D
	  else
	    gphys = gphys.first2D
	  end
	  data_paths.push(path)
	  file = NetCDF.create(path)
	  GPhys::IO.write(file,gphys)
	  file.close
	  newplots[i] = pl
	end
	if gphys2
	  if !cut2 && !slice2 && \
	    ( ( ONED_METHODS.include?(method) && gphys2.rank==1 ) || \
	      ( TWOD_METHODS.include?(method) && gphys2.rank==2 ) )
	    # no slicing needed
	    newplots[i] = pl
	  else
	    filename = basename + sprintf("_%03d_2",i) + '.nc'
	    path = @@dumpdir + filename
	    pl = pl.dup
	    pl[:file2] = filename
	    if ONED_METHODS.include?(method)
	      gphys2 = gphys2.first1D
	    else
	      gphys2 = gphys2.first2D
	    end
	    data_paths.push(path)
	    file = NetCDF.create(path)
	    GPhys::IO.write(file,gphys2)
	    file.close
	    newplots[i] = pl
	  end
	end
	[:cut, :slice, :cut2, :slice2].each{|k| pl.delete(k)}
      end
      @plots = newplots

      data_paths
    end

    ###################################################
    private

    def get_gphyses(arg)
      file = arg.delete(:file)
      var = arg.delete(:var)

      if file && var
        gphys = GPhys::IO.open(file,var)
        slice  = arg.delete(:slice)
        cut = arg.delete(:cut)
        gphys = gphys[*slice] if slice
        case cut
        when Hash
          gphys = gphys.cut(cut)
        when Array
          gphys = gphys.cut(*cut)
        end
      else
	raise ArgumentError, "Need to specify a variable by :file and :var"
      end

      file2 = arg.delete(:file2)
      var2 = arg.delete(:var2)
      file2 = file if !file2 && var2
      var2 = var if file2 && !var2 

      if file2 && var2
        gphys2 = GPhys::IO.open(file2,var2)
        slice2  = arg.delete(:slice2) || slice
        cut2  = arg.delete(:cut2) || cut
        gphys2 = gphys2[*slice2] if slice2
        case cut2
        when Hash
          gphys2 = gphys2.cut(cut2)
        when Array
          gphys2 = gphys2.cut(*cut2)
        end
      else
	gphys2 = nil
      end

      [gphys,gphys2]
    end

    def exec_plot(arg, first)
      arg = arg.dup
      arg[:newfrm] = first if arg[:newfrm]==nil

      method = arg.delete(:method)    # :tone_cont, :line, ..

      gphys, gphys2 = get_gphyses(arg)

      opt = arg

      case method
      when :tone_cont
        tone_cont(gphys, opt)
      when :vector
	raise ArgumentError, "2nd variable is not specified" if !gphys2
        vector(gphys, gphys2, opt)
      when :plot_1d
        plot_1d(gphys, opt)
      when :tone
        tone_cont(gphys, opt.update(:tone=>true,:contour=>false))
      when :contour
        tone_cont(gphys, opt.update(:tone=>false,:contour=>true))
      when :line
        plot_1d(gphys, opt.update(:line=>true,:mark=>false))
      when :mark
        plot_1d(gphys, opt.update(:line=>false,:mark=>true))
      else
        raise ArgumentError, "unsuppored draw method: "+method.to_s
      end
    end

    def setup(postscript=nil,image_dump=nil)
      DCL::swiset("iwidth",  @set[:admin][:iwidth])
      DCL::swiset("iheight", @set[:admin][:iheight])
      if postscript
	raise "DCL on MS Windows does not support PostScript file O" if @@ms_windows
        DCL::gropn(2)
      else
        if image_dump
	  if !@@ms_windows
	    DCL::swcset("fname", @@dumpdir + @set[:admin][:basename])
	    DCL::swlset("lwnd", false)
	    DCL::gropn(4)
	    print "DUMP IMAGE FILE\n"
	  else
	    DCL::swlset("ldump", true)
	    DCL::gropn(1)
	  end
	else
	  DCL::gropn(1)
        end
      end
      DCL::sglset("lfull", true)
      DCL::sgpset("isub", 96)
      DCL::gllset("lmiss", true)

      if ( @set[:admin][:xdiv] > 1 || @set[:admin][:ydiv] > 1)
	DCL::sldiv('y',@set[:admin][:xdiv],@set[:admin][:ydiv])
      end

      GGraph::set_fig(@set[:fig]) if @set[:fig]
      GGraph::set_map(@set[:map]) if @set[:map]
      GGraph::set_axes(@set[:axes]) if @set[:axes]
      GGraph::set_contour(@set[:contour]) if @set[:contour]
      GGraph::set_tone(@set[:tone]) if @set[:tone]
    end
 
    def finish
      DCL::grcls
    end

    def tone_cont(gphys, opt)
      newfrm = opt.delete(:newfrm)
      contour = opt.delete(:contour) 
      tone = opt.delete(:tone) 
      color_bar = opt.delete(:color_bar)
      color_bar_options = opt.delete(:color_bar_options)
      if color_bar_options && !color_bar_options.is_a?(Hash)
	raise ":color_bar_options must be a Hash"
      end

      if @set[:fig] && (itr=@set[:fig]['itr']) && itr>=10
	lon = gphys.coord(0)
	lat = gphys.coord(1)
	if lon.min >= 120 && lon.max <= 150 && lat.min >= 20 && lat.max <=50
	  GGraph::set_map({'coast_japan'=>true})
	else
	  GGraph::set_map({'coast_world'=>true})
	end
	if itr==31
	  DCL.sgpset('lclip',true) 
	  GGraph.set_map('vpt_boundary'=>3)
	end
      end

      tone = contour = true if !tone && !contour   # should be what is meant

      if tone 
        GGraph::tone(gphys, newfrm, opt)
        newfrm = false
      end
      if contour
        GGraph::contour(gphys, newfrm, opt)
        newfrm = false
      end
      if tone && color_bar
	cbopt = color_bar_options || Hash.new
	cbopt['log'] = true if opt['log']
	GGraph::color_bar(cbopt)
      end
      newfrm
    end

    def vector(gpx, gpy, opt)
      newfrm = opt.delete(:newfrm)
      GGraph::set_unit_vect_options('vyuoff'=>-0.1,'vxuoff'=>0.07)
      GGraph::vector(gpx, gpy, newfrm, opt)
    end

    def plot_1d(gphys, opt)
      newfrm = opt.delete(:newfrm)
      contour = opt.delete(:contour) 
      line = opt.delete(:line) 
      mark = opt.delete(:mark)

      line = mark = true if !line && !mark   # should be what is meant

      if line 
        GGraph::line(gphys, newfrm, opt)
        newfrm = false
      end
      if mark
        GGraph::mark(gphys, newfrm, opt)
        newfrm = false
      end
      newfrm
    end


  end
end

########## ENDOFLIB ##########

if __FILE__ == $0

  if ARGV.length == 0
    raise <<-EOS

      USAGE: % ruby #{__FILE__} menu [dumpdir]
      where menu=0,1,2,..

    EOS
  end

  test_menu = ARGV[0].to_i
  NumRu::VizShot.dumpdir = ARGV[1] if ARGV[1]

  case test_menu

  when 0
    viz = NumRu::VizShot.new(:iwidth=>700,:iheight=>500)
    viz.set_fig('viewport'=>[0.15,0.85,0.2,0.55])
    viz.set_tone('tonc'=>true)
    viz.plot( :method => :tone_cont,
	      :file => "../../testdata/T.jan.nc", 
	      :var=> "T", :cut => {'level'=>400},
	      :color_bar=>true, "title"=>"T and [U,V]" )
    viz.plot( :method => :vector,
	      :file  => "../../testdata/UV.jan.nc",
	      :var=> "U",  :var2=> "V" , :cut => {'level'=>400},
	      "unit_vect" => true )  
	      # You can give GGraph options using string keys.

    paths = viz.dump_code_and_data(nil, true)
    print "\n!!! Code and data stored in #{paths.inspect} !!!\n\n"

    viz.execute

  when 1

    viz = NumRu::VizShot.new(:iwidth=>700,:iheight=>500)
    viz.set_fig('itr'=>10, 'viewport'=>[0.15,0.85,0.2,0.55])
    viz.plot( :method => :tone_cont,
	      :file => "../../testdata/T.jan.nc", 
	      :var=> "T", :cut => {'lon'=>90..270,'lat'=>0..90,'level'=>600},
	      :color_bar=>true)

    path = viz.dump_code(nil, true, :image_dump => true)
    print "\n!!! Code written in #{path} !!!\n\n"

    viz.execute(:image_dump => true)

  when 2

    viz = NumRu::VizShot.new
    viz.set_fig('itr'=>31, 'viewport'=>[0.15,0.85,0.15,0.85])
    viz.plot( :method => :tone_cont,
	      :file => "../../testdata/T.jan.nc", 
	      :var=> "T", :cut => {'lat'=>10..90,'level'=>600},
	      :color_bar=>true)

    path = viz.dump_code
    print "\n!!! Code written in #{path} !!!\n\n"

    viz.execute

  when 3

    viz = NumRu::VizShot.new
    viz.plot( :method => :plot_1d,
	      :file => "../../testdata/T.jan.nc", 
	      :var=> "T", :cut => {'lon'=>135}, "index"=>2,
	      :line => true, :mark => true )

    path = viz.dump_code
    print "\n!!! Code written in #{path} !!!\n\n"

    viz.execute

  when 4

    # you can use :method=>:line instead of :method=>:plot_1d && :line=>true

    viz = NumRu::VizShot.new
    viz.plot( :method => :line,
	      :file => "../../testdata/T.jan.nc", 
	      :var=> "T", :cut => {'lon'=>135}, "index"=>2, "type"=>2 )

    path = viz.dump_code
    print "\n!!! Code written in #{path} !!!\n\n"

    viz.execute

  when 5

    viz = NumRu::VizShot.new(:iwidth=>700,:iheight=>500,:xdiv=>2,:ydiv=>2)
    viz.set_fig('viewport'=>[0.15,0.85,0.2,0.55])
    [1000,600,400,200].each do |lev|
      viz.plot( :method => :tone_cont, :newfrm=>true,
	        :file => "../../testdata/T.jan.nc", 
	        :var=> "T", :cut => {'level'=>lev},
	        :color_bar=>true)
    end

    path = viz.dump_code
    print "\n!!! Code written in #{path} !!!\n\n"

    viz.execute

  when 6

    viz = NumRu::VizShot.new(:iwidth=>700,:iheight=>500,:xdiv=>2,:ydiv=>2)
    viz.set_fig('viewport'=>[0.2,0.8,0.27,0.6])
    viz.plot( :method => :tone_cont, :newfrm=>true,
	      :file => "../../testdata/T.jan.nc", 
	      :var=> "T", :cut => {'level'=>1000},
	     :color_bar=>true, :color_bar_options=>{'vlength'=>0.25})
    viz.plot( :method => :tone_cont, :newfrm=>true,
	      :file => "../../testdata/T.jan.nc", 
	      :var=> "T", :cut => {'level'=>1000},
	      :color_bar=>true, :color_bar_options=>{'labelintv'=>1})
    viz.plot( :method => :tone_cont, :newfrm=>true,
	      :file => "../../testdata/T.jan.nc", 
	      :var=> "T", :cut => {'level'=>1000},
	      :color_bar=>true, 
              :color_bar_options=>{'landscape'=>true, 'labelintv'=>1,
               'vlength'=>0.6})

    path = viz.dump_code
    print "\n!!! Code written in #{path} !!!\n\n"

    viz.execute

  end

end