#!/usr/bin/ruby
#-*- encoding: sjis -*-

=begin
周波数時間解析をWaveletで行う
=end

$: << "~/src/eclipse/autopilot/common/param/swig/build_SWIG"
$: << "~/src/eclipse/autopilot/common/ruby"

require 'gsl'
require 'wavelet'
require 'gnuplot'
require 'gnuplot_support'
require 'stringio'

options = {}

# オプションのチェック
ARGV.reject!{|arg|
  if arg =~ /--([^=]+)=?/ then
    key = $1
    value = $'
    case key
      when /^window(?:_size)?$/
        options[:window_width] = value.to_i
      when 'order'
        options[:window_width] = 1 << value.to_i
      when 'prefix'
        options[:eps_fname_base] = value
      when 'labels'
        options[:labels] = value.split(/,/)
      when 'script'
        Gnuplot::verbose(true)
      else
        if value.empty? then
          value = true
        else
          begin
            value = Float(value)
          rescue
            case value.downcase
              when 'true'; value = true
              when 'false'; value = false
            end
          end
        end
        options[key.intern] = value
    end
    true
  else
    false
  end
}

input_fname = ARGV.shift

DEFAULT_OPTIONS = {
    :window_width => 64,
    :labels => [],
    :with_orig_signal => true,
    :with_fft => false,
    :auto_max => nil,
    :eps_fname_base => "#{File.basename(input_fname || '', ".*")}_"}

$stderr.puts <<-__STRING__
Usage: #{__FILE__} input.csv [options: #{DEFAULT_OPTIONS.to_a.collect{|k, v| "--#{k}=#{v}"}.join(' ')}]
__STRING__

raise "No file specified!!" unless input_fname

options = DEFAULT_OPTIONS.merge(options)
Gnuplot::verbose if options[:gnuplot_debug]

data = []

read_data = Proc::new{|io|
  res = []
  decimation_level = 1
  line_no = 0
  
  decimation_proc_str =<<-__STRING__
    if res.size >= options[:auto_max] then
      i = res.size
      while i > 0
        i -= 2
        res.delete_at(i)
      end
      decimation_level <<= 1
    end
  __STRING__
  
  eval(<<-__STRING__)
  io.each{|line|
    line_no += 1
    #{"next if line_no % decimation_level != 0" if options[:auto_max]}
    
    # 0列: 時刻, あとは入出力データ
    res << line.chop.split(/\s*,\s*/)[0..-1].collect{|item| item.to_f}
    
    #{decimation_proc_str if options[:auto_max]}
  }
  __STRING__
  
  if decimation_level > 1
    options[:window_width] >>= 1
    $stderr.puts "Warning: automatic data reduction from #{line_no} to #{res.size}."
  end
  res
}

if input_fname == '-' then
  data = read_data.call($stdin)
else
  open(input_fname){|io|
    data = read_data.call(io)
  }
end

raise Exception::new("Data too short!!") if data.size < options[:window_width]
  
if options[:filter_bank] then
  $: << File::join(File::dirname(__FILE__), '..', '090820_wavelet')
  require 'haar_filter_bank'
  options[:level] = Math::log2(options[:window_width]).to_i - 1
end

first_t = data[0][0]
end_t = data[-1][0]
step_t = (end_t - first_t) / (data.size - 1) 

data_wt = Wavelet::decompose(
    data, 
    options.dup.update({
        :data_valid => 1..-1,
        :window_size => options[:window_width]}))
end_t += (data_wt[0].inject(0){|res, item| res + item.size} - data.size) * step_t

data_t = data.transpose

# 必要に応じてFFTをしてみる
data_fft_power = []
if options[:with_fft] then
data_fft_freq = GSL::Vector.linspace(0, 1.0/step_t/2, (data.size + 1) / 2).to_a
data_t[1..-1].each{|item|
  ffted = GSL::Vector.alloc(*item).fft
  result = [ffted[0] * ffted[0]]
  # realをFFTするとcomplexが圧縮されている(GSLのFFT 7.1 Data storage scheme)
  # (item.size == 5) => 4([1-4])
  # (item.size == 6) => 4([1-4]) 
  ffted.subvector(1, ((item.size - 1) / 2) * 2).to_complex2.abs.each{|v|
    result << v
  }
  data_fft_power << result
}
end

def to_ascii(chr)
  RUBY_VERSION >= "1.9.0" ? chr.bytes.to_a[0] : chr
end

plot_file = Proc::new{|eps_fname, command|
  Gnuplot.open{|gp|
    if options[:png] then
      png_fname = "#{File::join(File::dirname(eps_fname), File::basename(eps_fname, '.*'))}.png"
      Plotter::plot(gp){|arg|
        $stderr.puts "drawing #{png_fname} ..."
        arg.terminal 'png font "Helvetica,18" size 1600,1200'
        arg.output png_fname
      }
    else
      Plotter::plot_eps(eps_fname, gp){|plot| }
    end
    gp << command
  }
}

# Waveletの絵を描く
data_wt.each_with_index{|item, i|
  col_name1 = (to_ascii(?B) + i)
  col_name2 = nil
  while col_name1 > to_ascii(?Z)
    if col_name2 then
      col_name2 += 1
    else
      col_name2 = to_ascii(?A)
    end
    col_name1 -= (to_ascii(?Z) + 1)
    col_name1 += to_ascii(?A)
  end
  
  eps_fname = eps_fname_fft = options[:eps_fname_base]
  if options[:labels][i] then
    eps_fname += "#{options[:labels][i]}.eps"
    eps_fname_fft += "FFT_#{options[:labels][i]}.eps"
  else
    eps_fname += "col_#{col_name2 ? col_name2.chr : '0'}#{col_name1.chr}.eps"
    eps_fname_fft += "FFT_col_#{col_name2 ? col_name2.chr : '0'}#{col_name1.chr}.eps"
  end
  
  plot_options = options.merge({ \
      :a_logarithmic => true, \
      :multi_plot => true,
      :b_shift => data[0][0],
      :b_scaling => (data[options[:window_width]][0] - data[0][0]) / options[:window_width], \
      :a_scaling => (1.0 / 2) / (data[options[:window_width]][0] - data[0][0])}) # Nyquist
  
  command = StringIO::new
        
  unless options[:with_orig_signal] then
    if options[:filter_bank]
      command << HaarFilterBank::init(options[:level]).plot(data_t[i + 1], plot_options)
    else
      command << Wavelet::plot(item, plot_options){|plot|
        plot.xlabel "'Time [s]'"
        plot.ylabel "'Frequency [Hz]'"
        #plot.set "palette rgbformulae 33,13,10"
      }
    end
    plot_file.call(eps_fname, command.string)
    next
  end
  
  x_range = (first_t.truncate - 1)..(end_t.ceil + 1)
  
  # 元の信号も一緒に書くため、multiplotを使う
  Plotter::plot(command){|plot|
    plot.set('xrange', "[#{x_range.first}:#{x_range.end}]")
    plot.set('lmargin 1')
    plot.set('rmargin 1')
    plot.set('multiplot')
  }
  command.puts
  
  # ウェーブレット
  set_plot_options = Proc::new{|plot|
    plot.set('format', 'x ""')
    plot.set('bmargin', 0)
    plot.set('tmargin', 1)
    plot.set('size', '1,0.8')
    plot.set('origin', '0,0.3')
    plot.ylabel "'Frequency [Hz]'"
    #plot.set('colorbox', 'horiz user origin .1,.02 size .8,.04 front noborder')
  }
  if options[:filter_bank]
    command << HaarFilterBank::init(options[:level]).plot(data_t[i + 1], plot_options){|plot|
      set_plot_options.call(plot)
    }
  else
    command << Wavelet::plot(item, plot_options){|plot|
      set_plot_options.call(plot)
    }
  end
  command.puts
  
  # 元信号
  Plotter::plot_3d(command){|plot|
    plot.xlabel "'Time [s]'"
    plot.set('size', '1.0,0.4')
    plot.set('origin', '0,0.05')
    plot.set('bmargin', 1)
    plot.set('tmargin', 0)
    plot.set('nocolorbox')
    plot.data = [ \
        Gnuplot::DataSet::new([data_t[0], data_t[i + 1], []]){|ds|
          ds.using = "($1):($2):(1)"
          ds.with = "lines"
          ds.notitle
        }]
  }
  command.puts
  
  # 閉じる
  Plotter::plot(command){|plot|
    plot.set('nomultiplot')
  }
  command.puts
  
  plot_file.call(eps_fname, command.string)
  
  next unless options[:with_fft]
  
  # FFTの解析結果を出す
  command.string = ""
  Plotter::plot(command){|plot|
    plot.xlabel "'Frequency [Hz]'"
    plot.ylabel "'Power'"
    plot.set "logscale", "y 10"
    
    items = [
        Gnuplot::DataSet::new([data_fft_freq, data_fft_power[i]]){|ds|
          ds.with = "lines lw 3"
          ds.notitle
        }
      ]
    plot.data = items
  }
  
  plot_file.call(eps_fname_fft, command.string)
}