#!/usr/bin/env ruby

require 'micro-optparse'
require 'spcore'
require 'wavefile'
require 'music-transcription'
require 'yaml'

include Music::Transcription

options = Parser.new do |p|
  p.banner = <<-END
Transcribe notes from an audio file.

Usage:
      transcribe [options] <audio_file>

Notes:
  Output directory must already exist.
  FFT actual_size must be a power of two.
  Delta (sec) must be > 0.
  Threshold must be >= 0.
  
Options:
END
  p.version = "0.1"
  p.option :outdir, "set output directory", :default => "./", :value_satisfies => lambda { |path| Dir.exist? path }
  p.option :freq_res, "set min frequency (Hz) resolution in Hz", :default => 1.0, :value_satisfies => lambda { |mfr| mfr > 0 }
  p.option :time_res, "set time (sec) between FFT samples", :default => 0.01, :value_satisfies => lambda { |dt| dt > 0 }
  p.option :threshold, "set theshold (minimum signal energy to run FFT)", :default => 0.5, :value_satisfies => lambda { |t| t >= 0 }
  p.option :min_freq, "set minimum frequency to consider during harmonic series analysis.", :default => PITCHES.first.freq, :value_satisfies => lambda { |f| f > 0 }
  p.option :verbose, "is verbose?", :default => false
end.process!

outdir = File.expand_path(options[:outdir])
freq_res = options[:freq_res]
time_res = options[:time_res]
threshold = options[:threshold]
verbose = options[:verbose]
min_freq = options[:min_freq]

TIME_RESOLUTION = time_res
THRESHOLD = threshold
FREQ_RESOLUTION = freq_res

def force_power_of_two size
  power_of_two = Math::log2(size)
  if power_of_two.floor != power_of_two # input size is not an even power of two
    size = 2**(power_of_two.to_i() + 1)
  end
  return size
end

ARGV.each do |filename|

  if !File.exist?(filename)
    puts "Could not find file #{filename}, skipping."
    next
  end

  if verbose
    puts
  end

  outfile = File.basename(filename, ".*") + ".yml"
  outpath = File.join(outdir, outfile)

  notes = []
  
  WaveFile::Reader.new(filename) do |reader|
    print "Transcribing #{File.basename(filename)} -> #{outfile}   0.0%"
    
    if reader.format.channels == 1
      samples = reader.read(reader.total_sample_frames).samples
    else
      puts "Multi-channel audo files not supported yet. Skipping."
      # TODO handl multi-channel audio
      break
    end
    signal = SPCore::Signal.new(:data => samples, :sample_rate => reader.format.sample_rate)  
    
    ideal_size = (TIME_RESOLUTION * signal.sample_rate).to_i
    fft_size = reader.format.sample_rate / FREQ_RESOLUTION
    fft_size = force_power_of_two fft_size
    
    # more samples are in the chunk than needed by the FFT. Make the FFT size
    # greater or equal to chunk size.
    if ideal_size > fft_size
      fft_size = force_power_of_two ideal_size
    end
    
    i = 0      
    while(i < signal.size)
      print "\b" * 6
      print "%5.1f%%" % (100 * i / signal.size)
      
      actual_size = ideal_size
      if (actual_size + i) > signal.size
        actual_size = signal.size - i
      end
      t = i / signal.sample_rate.to_f
      puts t if verbose
      
      chunk = signal.subset i...(i + actual_size)
      duration_sec = actual_size / signal.sample_rate.to_f
      
      if chunk.energy > THRESHOLD
        
        
        window = SPCore::BlackmanWindow.new(chunk.size)
        chunk.multiply! window.data

        remaining_before = (fft_size - chunk.size) / 2
        remaining_after = fft_size - chunk.size - remaining_before
        
        if remaining_before > 0
          chunk.prepend! Array.new(remaining_before, 0)
        end
        
        if remaining_after > 0
          chunk.append! Array.new(remaining_after, 0)
        end

        series = chunk.harmonic_series(:min_freq => min_freq)
        intervals = []
        if series.nil?
        if series.any?
          fundamental = series.min
          pitch = Pitch.make_from_freq fundamental
          intervals.push Interval.new(:pitch => pitch)
        end
        notes.push Note.new(:duration => duration_sec, :intervals => intervals)
      else
        notes.push Note.new(:duration => duration_sec, :intervals => [])
      end
      
      i += ideal_size
    end
  end
  
  new_notes = []
  
  # simplify the part
  notes.each_index do |i|
    note = notes[i]
    
    if new_notes.any?
      prev_note = new_notes.last
      
      if prev_note.intervals.any? and note.intervals.any?
        if prev_note.intervals == note.intervals
          prev_note.duration += note.duration
        else
          prev_note.intervals.first.link = slur(note.intervals.first.pitch)
          new_notes.push note
        end
      elsif prev_note.intervals.empty? and note.intervals.empty?
        prev_note.duration += note.duration
      else
        new_notes.push note
      end
    else
      new_notes.push note
    end
  end
  
  part = Part.new(:notes => new_notes)
  yaml = part.make_hash.to_yaml
  
  File.open(outpath, 'w') do |file|
    file.write yaml
  end
  
  puts
end