#!/usr/bin/ruby

require 'gsl'
require "rbgsl"
require "stringio"

$: << File::dirname(__FILE__) unless $:.include?(File::dirname(__FILE__))
require "gnuplot_support"

unless Math::methods.include?('log2') then
def Math.log2(x)
  return Math::log(x) / Math::log(2)
end
end

class Wavelet
  @@DEFUALT_OPTIONS = { \
      :window_size => 64, \
      :type => 'daubechies', \
      :order => 4, \
      :data_valid => 0..-1}
  
  def Wavelet::check_options(options)
    @@DEFUALT_OPTIONS.each{|key, v|
      options[key] ||= v
    }
  end
  
  #
  # Waveletで分解をする
  #
  # origの構造は[[t0のデータ1, t0のデータ2, ...], [t1のデータ1, t1のデータ2, ...], ...]
  # 出力は[[データ1のWT結果1, データ1のWT結果2, ...], [データ2のWT結果1, データ2のWT結果2, ...], ...]
  #
  def Wavelet::decompose(orig, options = {})
    check_options(options)
    
    w = GSL::Wavelet.alloc(options[:type], options[:order])
    data_range = options[:data_valid]
    window_width = options[:window_size]
    data = orig.clone
    
    decomposed = []
    (data[0][data_range].size).times{|i| decomposed << []}
    
    while true
      target = (data[0...window_width]).transpose
      target[data_range].each_with_index{|item, i|
        vec = GSL::Vector.alloc(*item)
        decomposed[i] << \
            w.transform(vec, GSL::Wavelet::FORWARD)
        
      }
      data = data[window_width..-1]
      
      if data.empty? then
        break
      elsif data.size < window_width then
        # Padding, 最後のデータを繰り返す
        (window_width - data.size).times{|i|
          data << data[-1]
        }
      end
    end
    
    return decomposed
  end
  
  #
  # Waveletの特定の分解レベルについて変更する
  # decomposed <= [[データ1のWT結果1, データ1のWT結果2, ...], [データ2のWT結果1, データ2のWT結果2, ...], ...]
  #
  def Wavelet::modify(decomposed, target_level)
    return unless block_given?
    max_level = Math::log2(decomposed.first.first.size).to_i
    decomposed.each_with_index{|data, n|
      if target_level == -1 then
        data.each{|item| item[0] = yield item[0], n}
        next
      end
      offset = 0
      data.each{|item|
        max_level.times{|level| # level = -1, i.e., [0]はゴミなので調べない
          next unless target_level == level
          begin_index = 2 ** level
          end_index = (2 ** (level + 1)) - 1
          (begin_index..end_index).each{|i|
            item[i] = yield item[i], n 
          }
        }
        offset += item.size
      }
    }
  end
  
  #
  # Waveletの特定の分解レベルについてitearteする
  # decomposed <= [[データ1のWT結果1, データ1のWT結果2, ...], [データ2のWT結果1, データ2のWT結果2, ...], ...]
  #
  def Wavelet::each(decomposed, target_level)
    modify(decomposed, target_level){|v, n| yield v, n; v}
  end
  
  #
  # Waveletで再構成を行う
  # decomposed <= [[データ1のWT結果1, データ1のWT結果2, ...], [データ2のWT結果1, データ2のWT結果2, ...], ...]
  # @return [[t0のデータ1, t0のデータ2, ...], [t1のデータ1, t1のデータ2, ...], ...]
  #
  def Wavelet::recompose(decomposed, options = {})
    check_options(options)
    
    w = GSL::Wavelet.alloc(options[:type], options[:order])
    
    data = decomposed.clone
    recomposed = []
    data.size.times{|i| recomposed << []}
    
    data.each_with_index{|wted_items, i|
      wted_items.each{|item|
        w.transform(item, GSL::Wavelet::BACKWARD).each{|v|
          recomposed[i] << v
        }
      }
    }
    
    return recomposed.transpose
  end
  
  #
  # TODO: どこが閾値を超えたか調べる
  # @decomposed : decomposeで分解されたもの => [[データ1のWT結果1, データ1のWT結果2, ...], [データ2のWT結果1, データ2のWT結果2, ...], ...]
  # @thresholds : 閾値のリスト {レベル　=> 閾値 }
  # @return : [[[データ1でのインデックス範囲, レベル], [データ1でのインデックス範囲, レベル], ...], ...]
  #
  def Wavelet::over_threshold(decomposed, thresholds = {})
    res = []
    max_level = Math::log2(decomposed.first.first.size).to_i
    decomposed.each{|data_n|
      res << []
      offset = 0
      data_n.each{|item|
        max_level.times{|level| # level = -1, i.e., [0]はゴミなので調べない
          threshold = thresholds[level]
          next unless threshold
          begin_index = 2 ** level
          end_index = (2 ** (level + 1)) - 1
          range_step_size = 2 ** (max_level - level)
          (begin_index..end_index).each_with_index{|i, j|
            range_top = (j * range_step_size) + offset
            range_bottom = range_top + (range_step_size - 1)
            res.last << [(range_top..range_bottom), level] \
                if (item[i] ** 2) >= threshold
          }
        }
        offset += item.size
      }
    }
    return res
  end

  #
  # Waveletで分解されたデータのグラフを描画する
  #  decomposed_data <= [wt_block0, wt_block1, ...] # :multiplot => true
  #  decomposed_data = wt_block0 # :multiplot => false
  #
  def Wavelet::plot(decomposed_data, options = {})
    
    options = {
        :b_scaling => 1, :b_shift => 0,
        :a_scaling => 1, :a_shift => 0, :a_logarithmic => true,
        :z_logarithmic => false,
        :plot_power => true,
        :area_regularization => false}.merge!(options)
    
    b_shift = options[:b_shift]
    
    if options[:a_max] then
      options[:a_scaling] = options[:a_max] / (options[:multi_plot] ? decomposed_data[0].size : decomposed_data.size)
    end
    
    max_z = 0
    
    data_extracter = Proc::new{|data|
      unless (data.kind_of?(GSL::Vector) || data.kind_of?(Array)) then
        raise "Unsupported format => #{decomposed_data.class}"
      end
    
      axis_b = []  # 時間軸に相当
      axis_a = []  # 周波数軸に相当
      axis_z = []
      
      Math::log2(data.size).to_i.times{|level|
        begin_index = 2 ** level
        end_index = (2 ** (level + 1)) - 1
        b_step = options[:b_scaling] * data.size / begin_index
        a_lower = options[:a_shift] + options[:a_scaling] * (2 ** level)
        a_upper = options[:a_shift] + options[:a_scaling] * (2 ** (level + 1))
        lower_line = []
        upper_line = []
        data[begin_index..end_index].to_a.each_with_index{|v, k|
          v **= 2 if options[:plot_power]
          v *= begin_index if options[:area_regularization] #面積で正規化
          max_z = v if v > max_z
          
          b_left = b_shift + b_step * k
          b_right = b_shift + b_step * (k + 1)
        
          lower_line << [b_left, a_lower, v] # 左下
          lower_line << [b_right, a_lower, v] # 右下
          upper_line << [b_left, a_upper, v] # 左上
          upper_line << [b_right, a_upper, v] # 右上
        }
        
        (lower_line + [['', '', '']] + upper_line + [['', '', '']]).each{|b, a, z|
          axis_b << b
          axis_a << a
          axis_z << z
        }
      }
      
      [axis_b, axis_a, axis_z]
    }
    
    drawer = Proc::new{|data, plot|
      plot.set('pm3d', 'map')
      plot.set('nokey')
      
      items = []
      
      if options[:multi_plot] then
        data.each{|item|
          items << Gnuplot::DataSet::new(data_extracter.call(item)){|ds|
                ds.notitle
              }
          b_shift += options[:b_scaling] * item.size
        }
      else
        items << Gnuplot::DataSet::new(data_extracter.call(data)){|ds|
              #ds.with = "dots" #"lines lw 2"
              #ds.title = "Analysis Result"
            }
      end
      
      plot.set "logscale", "y 2" if options[:a_logarithmic]
      
      if options[:z_logarithmic] then
        plot.set "logscale", "cb"
        plot.cbrange("[1E#{Math::log10([max_z, 1E-5].max).to_i - 3}:*]")
      else
        # 変動があまりない場合の対策
        plot.cbrange("[0:#{max_z < 1E-5 ? '1E-5' : '*'}]")
      end
      
      yield(plot) if block_given?
      
      plot.data = items
    }
    
    if options[:fname] then
      Plotter::plot_basic_3d(options[:fname]){|plot|
        drawer.call(decomposed_data, plot)
      }
    else
      sio = StringIO::new
      Plotter::plot_3d(sio){|plot|
        drawer.call(decomposed_data, plot)
        plot.reset_after_plot
      }
      return sio.string
    end
  end
end

if __FILE__ == $0 then

data = <<-__TEXT_DATA__.split(/\s+/).collect{|item| item.to_f}
0.0462458471760794
0.0462458471760794
0.0512458471760794
0.0712458471760795
0.0712458471760795
0.0662458471760795
0.0962458471760795
0.101245847176079
0.116245847176079
0.121245847176079
0.116245847176079
0.106245847176079
0.0912458471760794
0.101245847176079
0.0962458471760795
0.0962458471760795
0.0962458471760795
0.0912458471760794
0.0862458471760795
0.0812458471760795
0.0862458471760795
0.101245847176079
0.111245847176079
0.116245847176079
0.0762458471760795
0.0362458471760795
0.0362458471760795
0.0212458471760795
0.0112458471760795
0.00875415282392056
0.00875415282392056
0.00375415282392055
0.00624584717607946
0.00124584717607945
0.00624584717607946
0.00375415282392055
0.0187541528239206
0.0237541528239205
0.0187541528239206
0.0187541528239206
0.0287541528239205
0.0237541528239205
0.0337541528239205
0.00875415282392056
0.0137541528239206
0.00875415282392056
0.00124584717607945
0.0237541528239205
0.0337541528239205
0.0187541528239206
0.00875415282392056
0.00375415282392055
0.00875415282392056
0.0287541528239205
0.0437541528239205
0.0387541528239205
0.0587541528239205
0.103754152823921
0.123754152823921
0.153754152823921
0.188754152823921
0.213754152823921
0.183754152823921
0.0937541528239205
0.0212458471760795
0.161245847176079
0.306245847176079
0.556245847176079
0.81124584717608
 1.04124584717608
 1.19624584717608
 1.26124584717608
 1.22624584717608
 1.07624584717608
0.81124584717608
0.486245847176079
0.211245847176079
0.0512458471760794
0.0687541528239206
0.128754152823921
0.153754152823921
0.133754152823921
0.103754152823921
0.0687541528239206
0.0687541528239206
0.0637541528239206
0.0687541528239206
0.0587541528239205
0.0587541528239205
0.0587541528239205
0.0737541528239206
0.0637541528239206
0.0637541528239206
0.0637541528239206
0.0537541528239205
0.0737541528239206
0.0887541528239205
0.0887541528239205
0.0787541528239206
0.0737541528239206
0.0687541528239206
0.0837541528239206
0.0737541528239206
0.0637541528239206
0.0537541528239205
0.0687541528239206
0.0687541528239206
0.0837541528239206
0.0887541528239205
0.0887541528239205
0.0687541528239206
0.0687541528239206
0.0737541528239206
0.0837541528239206
0.0937541528239205
0.0787541528239206
0.0887541528239205
0.0837541528239206
0.0887541528239205
0.0937541528239205
0.0887541528239205
0.0787541528239206
0.0787541528239206
0.0737541528239206
0.0687541528239206
0.0837541528239206
0.0887541528239205
0.0687541528239206
0.0687541528239206
0.0637541528239206
0.0637541528239206
0.0887541528239205
0.0837541528239206
0.0737541528239206
0.0687541528239206
0.0537541528239205
0.0687541528239206
0.0737541528239206
0.0887541528239205
0.0787541528239206
0.0687541528239206
0.0687541528239206
0.0637541528239206
0.0837541528239206
0.0937541528239205
0.0937541528239205
0.0787541528239206
0.0737541528239206
0.0837541528239206
0.0937541528239205
0.0987541528239205
0.0987541528239205
0.0887541528239205
0.0937541528239205
0.103754152823921
0.0987541528239205
0.113754152823921
0.108754152823921
0.108754152823921
0.0987541528239205
0.108754152823921
0.128754152823921
0.133754152823921
0.128754152823921
0.113754152823921
0.123754152823921
0.128754152823921
0.133754152823921
0.148754152823921
0.138754152823921
0.133754152823921
0.128754152823921
0.133754152823921
0.148754152823921
0.153754152823921
0.138754152823921
0.128754152823921
0.123754152823921
0.118754152823921
0.113754152823921
0.118754152823921
0.0887541528239205
0.0737541528239206
0.0487541528239205
0.0437541528239205
0.0387541528239205
0.0437541528239205
0.0187541528239206
0.00375415282392055
0.00624584717607946
0.00124584717607945
0.00875415282392056
0.00875415282392056
0.00124584717607945
0.0112458471760795
0.0212458471760795
0.0212458471760795
0.00124584717607945
0.00124584717607945
0.00624584717607946
0.0162458471760795
0.0162458471760795
0.0262458471760795
0.00124584717607945
0.00875415282392056
0.0162458471760795
0.0112458471760795
0.0212458471760795
0.0212458471760795
0.00124584717607945
0.00375415282392055
0.0112458471760795
0.0162458471760795
0.00624584717607946
0.0162458471760795
0.00624584717607946
0.00624584717607946
0.0112458471760795
0.0262458471760795
0.0312458471760795
0.0162458471760795
0.0112458471760795
0.00124584717607945
0.00624584717607946
0.0212458471760795
0.00624584717607946
0.00624584717607946
0.00624584717607946
0.00875415282392056
0.00624584717607946
0.00124584717607945
0.00624584717607946
0.00375415282392055
0.0137541528239206
0.0187541528239206
0.0137541528239206
0.0137541528239206
0.00875415282392056
0.00375415282392055
0.0237541528239205
0.0287541528239205
0.0237541528239205
0.0137541528239206
0.00875415282392056
0.00875415282392056
0.0237541528239205
0.0237541528239205
0.0237541528239205
0.00124584717607945
0.00875415282392056
0.0137541528239206
0.0187541528239206
0.0337541528239205
0.0137541528239206
0.00875415282392056
0.00875415282392056
__TEXT_DATA__

nc = 20

orig = GSL::Vector.alloc(*data)

w = GSL::Wavelet.alloc("daubechies", 4)
#work = GSL::Wavelet::Workspace.alloc(n)

##### Choose as you like...
wt_orig = w.transform(orig, GSL::Wavelet::FORWARD) #, work)
#wt_orig = w.transform(orig, work)
#wt_orig = w.transform(orig)
#wt_orig = w.transform(orig, GSL::Wavelet::FORWARD)
#wt_orig = w.transform_forward(orig, work)
#wt_orig = w.transform_forward(orig)
#wt_orig = GSL::Wavelet.transform(w, orig, GSL::Wavelet::FORWARD, work)
#wt_orig = GSL::Wavelet.transform(w, orig, GSL::Wavelet::FORWARD)
#wt_orig = GSL::Wavelet.transform(w, orig, work)
#wt_orig = GSL::Wavelet.transform_forward(w, orig)
#wt_orig = orig.wavelet_transform(w, GSL::Wavelet::FORWARD, work)
#wt_orig = orig.wavelet_transform_forward(w, work)

=begin
# 単体プロット
Wavelet::plot(wt_orig, { \
    :fname => "hoge.eps", \
    :a_logarithmic => true}){|plot|
  plot.xlabel "'Index'"
  plot.ylabel "'Freq [Hz]'"
  #plot.set "palette rgbformulae 33,13,10"
}
=end

# 複数プロット
Wavelet::plot([wt_orig, wt_orig], { \
    :fname => "hoge.eps", \
    :a_logarithmic => true, \
    :multi_plot => true}){|plot|
  plot.xlabel "'Index'"
  plot.ylabel "'Freq [Hz]'"
  #plot.set "palette rgbformulae 33,13,10"
}

p Wavelet::over_threshold([[wt_orig]], {0 => 0, 3 => 0.5, 4 => 0.5})

perm = wt_orig.abs2.sort_index

(0...(data.size - nc)).each{|i|
  wt_orig[perm[i]] = 0.0  
}

#iwt_wt_orig = w.transform(wt_orig, GSL::Wavelet::BACKWARD, work)
#iwt_wt_orig = w.transform(wt_orig, GSL::Wavelet::BACKWARD)
#iwt_wt_orig = w.transform_inverse(wt_orig, work)
#iwt_wt_orig = w.transform_inverse(wt_orig)
#iwt_wt_orig = GSL::Wavelet.transform(w, wt_orig, GSL::Wavelet::BACKWARD, work)
#iwt_wt_orig = GSL::Wavelet.transform(w, wt_orig, GSL::Wavelet::BACKWARD)
#iwt_wt_orig = GSL::Wavelet.transform_inverse(w, wt_orig, work)
iwt_wt_orig = GSL::Wavelet.transform_inverse(w, wt_orig)
#iwt_wt_orig = wt_orig.wavelet_transform(w, GSL::Wavelet::BACKWARD, work)
#iwt_wt_orig = wt_orig.wavelet_transform_inverse(w, work)

#p iwt_wt_orig
#p Wavelet::recompose([[wt_orig]])

end