require 'gnuplot'

class EqString < String
  
end

module Gnuplot
  FITLOG_FNAME = "fit.log"
  
  class Plot
    def unset(var)
      @unsets ||= []
      @unsets << var
    end
    
    def reset_after_plot
      @reset = true
    end
    
    alias :orig_to_gplot :to_gplot
    def to_gplot(io = "")
      (@unsets ||= []).each{|var| io << "unset #{var}\n"}
      orig_to_gplot(io)
      if (@reset rescue false) then
        @sets.each{|var, val|
          next if var =~ /^no.+/ # TODO: should be changed, because set nokey etc. are deprecated.
          io << "unset #{var}\n"
        }
        @unsets.each{|var|
          io << "set #{var}\n"
        }
      end
      io
    end
  end
  
  class SPlot
    alias :orig_to_gplot :to_gplot
    def to_gplot(io = "")
      (@unsets ||= []).each{|var| io << "unset #{var}\n"}
      orig_to_gplot(io)
      if (@reset rescue false) then
        @sets.each{|var, val|
          next if var =~ /^no.+/ 
          next if var == 'palette'
          io << "unset #{var}\n"
        }
        @unsets.each{|var|
          io << "set #{var}\n"
        }
      end
      io
    end
  end
  
  class DataSet
    def plot_args (io = "")
  
      # Order of these is important or gnuplot barfs on 'em

      io << ( (@data.instance_of? String) ? @data : "'-'" )

      io << " using #{@using}" if @using

        io << case @title
        when /notitle/ then " notitle"
        when nil       then ""
        else
          temp = (EqString == @title.class ? @title : "'#{@title}'")
          " title #{temp}"
        end

      io << " matrix" if @matrix
      io << " with #{@with}" if @with
      io << " linewidth #{@linewidth}" if @linewidth
      io
    end
  end
  
  class << self
    def open_default(persist = true)
      cmd = Gnuplot.gnuplot(persist) or raise 'gnuplot not found'
      IO::popen("#{cmd} 2>&1", 'w+'){|io| yield io}
    end
    alias :open :open_default
    
    def open_debug(before_open, persist = true)
      buf = StringIO::new
      yield buf
      buf.rewind
      script = buf.read
      before_open.call(script)
      open_default(persist){|gp| gp << script}
    end
    
    def info(head, *args)
      $stderr.puts "[GP.#{sprintf('%-6s', head)}] #{args.collect{|arg| "#{arg} "}.join}..." 
    end
    
    def verbose(is_verbose = true)
      unless is_verbose
        alias :open :open_default
        return
      end
      
      require 'stringio'
            
      def Gnuplot.open(persist = true, &b)
        open_debug(proc{|script|
          out_fname = nil
          script.lines.each{|line|
            if line.chomp =~ /^ *set +out(?:put)? +/ then
              out_fname = "#{$'.gsub(/^\"(.+)\"$/, '\1').gsub(/^\'(.+)\'$/, '\1')}.gnuplot"
              break
            end
          }
          out_fname ||= proc{
            fnum = [-1]
            Dir::open('.')::each{|fname|
              fnum << $1.to_i if fname =~ /^(\d+)\.gnuplot/
            }
            sprintf("%04d.gnuplot", fnum.max + 1)
          }.call
          info('script', out_fname)
          Kernel::open(out_fname, 'w'){|io| io << script}
        }, persist, &b)
      end
    end
  end
  
  def Gnuplot.key_front(eps_fname)
    Kernel::open(eps_fname, 'r+'){|f|
      lines = f.readlines
      proc{# M[LCR]showの修正
        delete_lines = [:L, :R, :C].collect{|align|
          def_start = lines.index{|line| line =~ /^\/M#{align}show/}
          next nil unless def_start
          def_length = lines[def_start..-1].index{|line| line.chomp =~ /def$/}
          next nil unless def_length
          [def_start, def_length]
        }.sort!{|a, b| b[0] <=> a[0]} # お尻の方から順に、後ろから消せば安全に消せる
        delete_lines.each{|def_start, def_length|
          (def_length + 1).times{|i| lines.delete_at(def_start)}
        }
        lines.insert(delete_lines[-1][0], <<'__TEXT__')
/MFshowW {gsave 
  0 dict begin /Gshow {[] 
    0 setdash 
    userlinewidth 3 mul setlinewidth 
    LCw setrgbcolor 
    false charpath currentpoint stroke M} def
  MFshow end grestore} def
/MLshow { currentpoint stroke M
  0 exch R
  dup MFshowW
  Blacktext {gsave 0 setgray MFshow grestore} {MFshow} ifelse } bind def
/MRshow { currentpoint stroke M
  exch dup MFwidth neg 3 -1 roll R
  dup MFshowW
  Blacktext {gsave 0 setgray MFshow grestore} {MFshow} ifelse} bind def
/MCshow { currentpoint stroke M
  exch dup MFwidth -2 div 3 -1 roll R
  dup MFshowW
  Blacktext {gsave 0 setgray MFshow grestore} {MFshow} ifelse } bind def
__TEXT__
      }.call
      proc{ # keyの裏側のボックスの除去、凡例線の裏書
        line_no = 0
        lines.reverse_each{|line|
          line_no -= 1
          if line =~ /^([\d.]+) (\d+ ){4}BoxColFill/ then
            lines[line_no] = "%#{line}"
            break
          elsif line.chomp =~ /^% End plot #\d+/ and lines[line_no - 1].chomp =~ /^\d+ \d+ V$/ then
            lines[line_no - 1] = <<"__TEXT__" + lines[line_no - 1]
currentpoint gsave [] 0 setdash userlinewidth 3 mul setlinewidth LCw setrgbcolor #{lines[line_no - 1].chomp} stroke grestore M
__TEXT__
          end
        }
      }.call
      f.rewind
      f.print lines.join
      f.truncate(f.tell)
    }
  end
end

class Plotter
  
  def Plotter::plot(gp = nil)
    Gnuplot::Plot::new(gp){|arg|
      yield arg
    }
  end
  
  def Plotter::plot_3d(gp = nil)
    Gnuplot::SPlot::new(gp){|arg|
      yield arg
    }
  end

  def Plotter::plot_eps(file_name, gp = nil)
    Plotter::plot(gp){|arg|
      Gnuplot::info('draw', file_name)
      arg.terminal "postscript eps enhanced color"
      arg.output file_name
      yield arg
    }
  end
  
  def Plotter::plot_eps_3d(file_name, gp = nil)
    plot_3d(gp){|arg|
      Gnuplot::info('draw', file_name)
      arg.terminal "postscript eps enhanced color"
      arg.output file_name
      yield arg
    }
  end

  def Plotter::auto_key_front
    [:plot_basic, :plot_basic_3d].each{|f|
      f_orig = "#{f}_orig".to_sym
      next if self.methods.include?(f_orig) # TODO 一時的な二重呼び出し防止措置
      self.instance_eval <<-__FUNC__
        alias :#{f_orig} #{f}
        def #{f}(eps_fname, &b)
          #{f_orig}(eps_fname){|plot|
            plot.set('key opaque nobox')
            b.call(plot)
          }
          Gnuplot::key_front(eps_fname)
          Gnuplot::info('modify', eps_fname)
        end
      __FUNC__
    }
  end

  def Plotter::plot_basic file_name
    Gnuplot.open{|gp|
      plot_eps(file_name, gp){|arg|
        yield arg
      }
    }
  end
  
  def Plotter::plot_basic_3d file_name
    Gnuplot.open{|gp|
      plot_eps_3d(file_name, gp){|arg|
        yield arg
      }
    }
  end
  
  def Plotter::draw_eps(fname, x, y)
    x_label, x_value = x.to_a[0]
    plot_basic(fname){|arg|
      arg.xlabel "'#{x_label}'"
        
      yield arg if block_given?
      
      items = []
      y.each{|y_label, y_value|
        items << Gnuplot::DataSet::new([x_value, y_value]){|ds|
            ds.with = "lines lw 3"
            ds.title = y_label
          }
      }
      
      arg.data = items
    }
  end
end

class Array
  def to_gplot
    if ( self[0].kind_of? Array ) then
      begin
        self.transpose.collect{|a| a.join(' ')}.join("\n") + "\ne"
      rescue
        tmp = self[0].zip( *self[1..-1] )
        tmp.collect { |a| a.join(" ") }.join("\n") + "\ne"
      end
    elsif ( self[0].kind_of? Numeric ) then
      s = ""
      self.length.times { |i| s << "#{self[i]}\n" }
      s
    else
      self[0].zip( *self[1..-1] ).to_gplot
    end
  end

  def to_gsplot
    f = ""
    
    if ( self[0].kind_of? Array ) then
      x = self[0]
      y = self[1]
      z = self[2]
      
      (0...x.size).each{|i|
        f << "#{x[i]} #{y[i]} #{z[i]}\n"
      }
    elsif ( self[0].kind_of? Numeric ) then
      self.length.times do |i| f << "#{self[i]}\n" end
    else
      self[0].zip( *self[1..-1] ).to_gsplot
    end
    
    f
  end
end

class DataSet3D < Array
  def to_gsplot
      f = ""
    
    x = self[0]
    y = self[1]
    d = self[2]

    x.each_with_index{|xv, i|
      if i > 0 then f << "\n" end
          y.each_with_index{|yv, j|
              f << "#{xv} #{yv} #{d[i][j]}\n"
      }
    }
    #puts f
    f
  end
end