#!/usr/bin/env ruby
 
=begin
 
Copyright (C) 2007 Nokia Corporation. All rights reserved.
 
Contact: Felipe Contreras <[email protected]>
 
Permission is hereby granted, free of charge, to any person obtaining a copy of
this software and associated documentation files (the "Software"), to deal in
the Software without restriction, including without limitation the rights to
use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies
of the Software, and to permit persons to whom the Software is furnished to do
so, subject to the following conditions:
 
The above copyright notice and this permission notice shall be included in all
copies or substantial portions of the Software.
 
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
SOFTWARE.
 
=end
 
=begin
 
This script tries to intelligently find the best video frame size that would
look good in the desired device (N770 or N800) while trying to keep the same
aspect ratio as the original clip, as well as trying to fit the aspect ratio of
the device.
 
It will try to change the framerate only to fit the capabilities of the device.
 
=end
 
require 'optparse'
 
class AppException < RuntimeError
end
 
module Video
  class Aspect
    attr_reader :numerator, :denominator
 
    def initialize(num, den)
      def greatest_common_divisor(a, b)
        while a % b != 0
          a, b = b.round, (a % b).round
        end 
        return b
      end
 
      gcd = greatest_common_divisor(num, den)
 
      @numerator = num / gcd
      @denominator = den / gcd
    end
 
    def to_s()
      return "%d:%d" % [numerator, denominator]
    end
 
    def to_f()
      return @numerator.to_f / @denominator.to_f
    end
 
    def /(b)
      self.to_f / b.to_f
    end
  end
 
  class FrameSize
    attr_reader :width, :height
    attr_writer :width, :height
 
    def initialize(width, height)
      @width = width
      @height = height
    end
 
    def ==(size)
      return false if size.width != @width
      return false if size.height != @height
      return true
    end
 
    def to_s()
      return "%dx%d" % [@width, @height] 
    end
 
    def aspect()
      return Aspect.new(@width, @height)
    end
 
    def /(b)
      self.width.to_f / b.width.to_f
    end
  end
 
  class Clip
    attr_reader :file_name, :size, :framerate, :bitrate
    attr_writer :file_name, :size, :framerate, :bitrate
 
    def initialize(file)
      @file_name = file
    end
 
    def to_s()
      return "%s [%s], %s fps, %s kbps" % [@size.to_s, @size.aspect.to_s, @framerate, @bitrate]
    end
  end
end
 
module Device
  class Base
    attr_reader :screen_size, :basic_sizes, :macroblocks_per_second, :max_framerate
 
    def initialize()
      @basic_sizes = []
    end
  end
 
  class N770 < Base
    def initialize()
      super()
      @screen_size = Video::FrameSize.new(800, 480)
      @macroblocks_per_second = 22 * 18 * 15 # 352x288x15
      @max_framerate = 30
 
      @basic_sizes << Video::FrameSize.new(240, 144)
      @basic_sizes << Video::FrameSize.new(352, 208)
      @basic_sizes << Video::FrameSize.new(352, 288)
      @basic_sizes << Video::FrameSize.new(176, 144)
 
      @basic_sizes << Video::FrameSize.new(320, 240)
    end
  end
 
  class N800 < N770
    def initialize()
      super()
      @screen_size = Video::FrameSize.new(800, 480)
      @macroblocks_per_second = 40 * 30 * 15 # 640x480x15
      @max_framerate = 30
 
      @basic_sizes << Video::FrameSize.new(400, 240)
      @basic_sizes << Video::FrameSize.new(640, 480)
    end
  end
end
 
module Transcoder
  class Base
    def initialize(input, output, device, bitrate)
      @input = input
      @output = output
      @device = device
 
      output.bitrate = bitrate
    end
 
    def run()
      raise "Run not implemented"
    end
  end
 
  class Smart < Base
    class Evaluator
      class Variable
        def initialize(element, weight)
          @element = element
          @weight = weight
        end
 
        def get(size)
          case @element
          when Video::FrameSize
            r = compare(@element, size)
          when Video::Aspect
            r = compare(@element, size.aspect)
          end
 
          return r * @weight
        end
 
        private
 
        # Returns the amount of similarity from 0 to 1
        def compare(a, b)
          return ((1.0 / 100) ** (Math.log(a / b) ** 2))
        end
      end
 
      attr_writer :framerate, :max_mbps
 
      def initialize()
        @variables = []
      end
 
      def add(element, weight)
        @variables << Variable.new(element, weight)
      end
 
      def execute(size)
        value = @variables.inject(0) {|sum, n| sum + n.get(size)}
 
        # We don't want a barely playable video
        mbps = (size.width / 16) * (size.height / 16) * @framerate
 
        if mbps >= @max_mbps
          # print "Out of range\n"
          value /= 2
        end
 
        # print "%3dx%3d: %f\n" % [size.width, size.height, value]
 
        return value
      end
    end
 
    def calculate()
      def nearest(num, mul)
        return (0.5 + num / mul).to_i * mul;
      end
 
      max_value = nil
      new_size = nil
      new_framerate = @input.framerate
 
      new_framerate /= 2 while new_framerate > @device.max_framerate
 
      evaluator = Evaluator.new()
      evaluator.framerate = new_framerate
      evaluator.max_mbps = @device.macroblocks_per_second
 
      # How similar to the original frame size?
      evaluator.add(@input.size, 50)
 
      # How similar to the original aspect ratio?
      evaluator.add(@input.size.aspect, 100)
 
      # How similar to the screen's aspect ratio?
      evaluator.add(@input.size.aspect, 75)
 
      @device.basic_sizes.each do |size|
        # Evaluate this frame size
        value = evaluator.execute(size)
 
        # Is this frame size the best or not?
        if not max_value or value > max_value
          new_size = size
          max_value = value
        end
      end
 
      if @input.size.aspect / new_size.aspect > 1.2
        # Change height keep the aspect ratio
        new_size.height = nearest(new_size.width * (1 / @input.size.aspect.to_f), 16)
      end
 
      # p @input.size.aspect / new_size.aspect
 
      @output.framerate = new_framerate
      @output.size = new_size.clone()
    end
  end
 
  class MEncoder < Smart
    def analyze()
      info_map = {}
      cmd = "mplayer -identify -quiet -frames 0 -vc null -vo null -ao null \"%s\"" % [@input.file_name]
      info_raw = %x[#{cmd} 2> /dev/null | grep "^ID_"]
      raise AppException, "Bad input file: \"%s\"" % [@input.file_name] if info_raw == ""
 
      info_array = info_raw.map { |i| i.chomp().split("=")}
      info_array.each { |e| info_map[e[0]] = e[1] }
 
      width = info_map["ID_VIDEO_WIDTH"].to_i
      height = info_map["ID_VIDEO_HEIGHT"].to_i
      @input.framerate = info_map["ID_VIDEO_FPS"].to_f
      @input.bitrate = info_map["ID_VIDEO_BITRATE"].to_i
 
      @input.size = Video::FrameSize.new(width, height)
    end
 
    def generate()
      messages = []
 
      audio_options = []
      audio_options << "-srate 44100"
      audio_options << "-oac mp3lame"
      audio_options << "-lameopts vbr=0:br=128"
      audio_options << "-af volnorm"
 
      video_options = []
      case $options[:output_format]
      when "mpeg4"
        video_options << "-ovc lavc"
        video_options << "-lavcopts vcodec=mpeg4:vbitrate=%d" % [@output.bitrate]
      when "h264"
        video_options << "-ovc x264"
        video_options << "-x264encopts bitrate=%d:nocabac" % [@output.bitrate]
      end
      video_options << "-ofps %f" % [@output.framerate]
 
      messages << "Input: %s." % [@input.to_s]
      messages << "Output: %s." % [@output.to_s]
 
      if @input.size != @output.size
        video_options << "-vf-add scale=%d:%d" % [@output.size.width, @output.size.height]
      end
 
      # For the 770?
      # video_options << "-ffourcc DIVX"
      # video_options << "-ffourcc DX50"
      # video_options << "-noidx"
 
      if $options[:verbose]
        messages.each do |m|
          print "* #{m}\n"
        end
      end
 
      return "mencoder %s -o %s %s %s" % [@input.file_name, @output.file_name, audio_options.join(" "), video_options.join(" ")]
    end
 
    def run()
      analyze()
      calculate()
      cmd = generate()
      print "#{cmd}\n"
    end
  end
end
 
class App
  def initialize(args)
    @args = args
    @quality_presets = [80, 96, 200, 300, 400, 800, 1000, 1500, 2000]
 
    @devices = {}
    @devices[:N770] = Device::N770
    @devices[:N800] = Device::N800
 
    $options = {}
    $options[:quality] = 5
    $options[:device] = Device::N800
    $options[:output_format] = "mpeg4"
  end
 
  def parse_options
    op = OptionParser.new do |opts|
      opts.banner = "Usage: transcode [options]"
 
      opts.on("-i", "--input FILE",
        "Input file") do |i|
        $options[:input_file] = i
      end
 
      opts.on("-o", "--output FILE",
        "Output file") do |o|
        $options[:output_file] = o
      end
 
      opts.on("-f", "--format FORMAT",
        "Output format") do |f|
        $options[:output_format] = f
      end
 
      opts.on("-q", "--quality N", Integer,
        "The quality of the output (1..%d) (default: %d)" % [@quality_presets.length, $options[:quality]]) do |q|
        $options[:quality] = q - 1
      end
 
      opts.on("-d", "--device DEVICE", @devices,
        "Output compatible for this device {%s} (default: %s)" % [@devices.keys.join("|"), $options[:device]]) do |d|
        $options[:device] = d
      end
 
      opts.on("-v", "--[no-]verbose",
        "Run verbosely") do |v|
        $options[:verbose] = v
      end
 
      opts.on_tail("-h", "--help",
        "Show this message") do
        puts opts
        exit
      end
    end
 
    op.parse!(@args)
 
    if not $options[:input_file]
      raise ArgumentError, "You need to specify an input file"
    end
 
    if not $options[:output_file]
      $options[:output_file] = File.basename($options[:input_file], ".*") + "_it.avi"
    end
  end
 
  def transcode
    @bitrate = @quality_presets[$options[:quality]]
    @input = Video::Clip.new($options[:input_file])
    @output = Video::Clip.new($options[:output_file])
    @device = $options[:device].new()
 
    trans = Transcoder::MEncoder.new(@input, @output, @device, @bitrate)
    trans.run()
  end
 
  def run
    begin
      parse_options()
    rescue => msg
      print "Error: #{msg}\n"
      exit
    end
 
    begin
      transcode()
    rescue AppException => msg
      print "Error: #{msg}\n"
      exit
    end
  end
end
 
app = App.new(ARGV)
app.run()