#!/usr/bin/env ruby

if ARGV.size < 1
  $stderr.puts "Usage: extract_wb <file1> [file2] ..."
  exit 2
end

# Presets we just skip
IGNORED_PRESETS=["Auto","Kelvin","Measured","AsShot","Kelvin","Preset",
                 "Multi Auto", "Color Temperature Enhancement", "Custom",
                 "One Touch WB 1", "One Touch WB 2", "One Touch WB 3",
                 "One Touch WB 4", "Custom WB 1", "Auto0", "Auto1", "Auto2"]
# Preset names we adjust
FL_PRESET_REPLACE={
  "Fluorescent" => "CoolWhiteFluorescent",
  "FluorescentP1" => "DayWhiteFluorescent",
  "FluorescentP2" => "DaylightFluorescent",
  "FluorescentM1" => "WarmWhiteFluorescent",
  "FluorescentD"  => "DaylightFluorescent",
  "FluorescentN"  => "NeutralFluorescent",
  "FluorescentW"  => "WhiteFluorescent",
  "Daylight Fluorescent" => "DaylightFluorescent",
  "Day White Fluorescent" => "DayWhiteFluorescent",
  "White Fluorescent" => "WhiteFluorescent",
  "Unknown (0x600)" => "Underwater",
  "Sunny" => "DirectSunlight",
  "Cool WHT FL" => "CoolWhiteFluorescent",
  "Daylight FL" => "DaylightFluorescent",
  "Warm WHT FL" => "WarmWhiteFluorescent",
  "Warm White Fluorescent" => "WarmWhiteFluorescent",
  "White FL" => "WhiteFluorescent",
  "Mercury Lamp" => "HighTempMercuryVaporFluorescent",
  "Day White FL" => "DayWhiteFluorescent",
  "Sodium Lamp" => "SodiumVaporFluorescent",
  "3000K (Tungsten light)" => "Tungsten",
  "4000K (Cool white fluorescent)" => "CoolWhiteFluorescent",
  "5300K (Fine Weather)" => "Daylight",
  "5500K (Flash)" => "Flash",
  "6000K (Cloudy)" => "Cloudy",
  "7500K (Fine Weather with Shade)" => "Shade",
}
PRESET_ORDER=["Unknown", "DirectSunlight", "Daylight","Shade","Cloudy","Tungsten","Incandescent","Fluorescent",
              "WarmWhiteFluorescent","CoolWhiteFluorescent","DayWhiteFluorescent","DaylightFluorescent",
              "DaylightFluorescent", "NeutralFluorescent", "WhiteFluorescent", "HighTempMercuryVaporFluorescent",
              "SodiumVaporFluorescent", "Underwater", "Flash"]

PRESET_SORT_MAPPING={}
PRESET_ORDER.each_with_index {|name, index| PRESET_SORT_MAPPING[name]=index+1}

CAMERAS=File.expand_path("../src/external/rawspeed/data/cameras.xml", File.dirname(__FILE__))

# Create a map between EXIF maker/model and our cleaned up maker/model
require 'nokogiri'
exif_name_map = {}
File.open(CAMERAS) do |f|
  xml_doc  = Nokogiri::XML(f)
  xml_doc.css("Camera").each do |c|
    maker = exif_maker = c.attribute("make").value
    model = exif_model = c.attribute("model").value
    exif_id = [maker,model]
    if c.css("ID")[0]
      maker = c.css("ID")[0].attribute("make").value
      model = c.css("ID")[0].attribute("model").value
    end
    exif_name_map[exif_id] = [maker,model]
    c.css("Alias").each do |a|
      exif_model = a.content
      exif_id = [exif_maker, exif_model]
      exif_name_map[exif_id] = [maker, model]
    end
  end
end

found_presets = []
ARGV.each do |filename|
  red = green = blue = maker = model = preset = nil
  listed_presets = []
  fl_count = 0
  IO.popen("exiftool \"-WB_*\" -WhiteBalance -WhiteBalance2 -WhitePoint -Make -Model \"#{filename}\"") do |io|
    while !io.eof?
      lineparts = io.readline.split(":")
      tag = lineparts[0].strip
      values = lineparts[1].strip.split(" ")

      if tag == "WB RGGB Levels"
        green = (values[1].to_f+values[2].to_f)/2.0
        red = values[0].to_f/green
        blue = values[3].to_f/green
        green = 1
      elsif tag == "WB RB Levels"
        red = values[0].to_f
        blue = values[1].to_f
        if values.size == 4 && values[2] == "256" && values[3] == "256"
          red /= 256.0
          blue /= 256.0
        end
        green = 1
      elsif tag == "WB GRB Levels"
        green = values[0].to_f
        red = values[1].to_f/green
        blue = values[2].to_f/green
        green = 1
      elsif tag == "White Point" && values.length > 3
        green = (values[1].to_f+values[2].to_f)/2.0
        red = values[0].to_f/green
        blue = values[3].to_f/green
        green = 1
      elsif tag.split.include? "Make"
        maker = lineparts[1].strip
      elsif tag.split.include? "Model"
        model = lineparts[1].strip
      elsif tag == "White Balance" || tag == "White Balance 2"
        preset = values.join(" ")
        preset = FL_PRESET_REPLACE[preset] if FL_PRESET_REPLACE[preset]
      elsif ["WB RGB Levels", "WB RGGB Levels"].include? tag.split[0..2].join(" ")
        p = tag.split[3..-1].join

        r=g=b=0
        if values.size == 4
          g = (values[1].to_f+values[2].to_f)/2.0
          r = values[0].to_f/g
          b = values[3].to_f/g
          g = 1
        elsif values.size == 3
          g = values[1].to_f
          r = values[0].to_f/g
          b = values[2].to_f/g
          g = 1
        else
          $stderr.puts "Found tag '#{p}' with #{values.size} values instead of 3 or 4"
        end
        fl_count += 1 if p[0..."Fluorescent".size] == "Fluorescent"
        p = "Unknown" if !p || p==""
        listed_presets << [p,r,g,b]
      end
    end
  end

  # Adjust the maker/model we found with the map we generated before
  if exif_name_map[[maker,model]]
    maker, model = exif_name_map[[maker,model]]
  else
    $stderr.puts "WARNING: Couldn't find model in cameras.xml ('#{maker}', '#{model}')"
  end

  found_presets += listed_presets.map do |preset, red, green, blue|
    # ugly hack. Canon's Fluorescent is listed as WhiteFluorescent in usermanual
    if maker && maker == "Canon" && preset == "Fluorescent"
      preset = "WhiteFluorescent"
    end

    preset = FL_PRESET_REPLACE[preset] if FL_PRESET_REPLACE[preset]
    [maker, model, preset, red, green, blue]
  end

  # Print out the WB value that was used in the file
  preset = "\"#{ARGV[0]}\"" if !preset
  if red && green && blue
    found_presets << [maker, model, preset, red, green, blue]
  end
end

# Get rid of duplicate presets
found_presets = Hash[found_presets.sort.zip].keys

def convert_preset_to_sort(preset)
  if IGNORED_PRESETS.include? preset
    return 0
  elsif PRESET_SORT_MAPPING[preset]
    return PRESET_SORT_MAPPING[preset]
  elsif preset[-1] == "K"
    return preset[0..-2].to_i
  else
    $stderr.puts "WARNING: no defined sort order for '#{preset}'"
    return 0
  end
end

found_presets.sort! do |a, b|
  convert_preset_to_sort(a[2]) <=> convert_preset_to_sort(b[2])
end

min_padding = 0
found_presets.each do |maker, model, preset, red, green, blue, acc|
  min_padding = preset.size if preset.size > min_padding
end

# Print out the presets if they exist
found_presets.each do |maker, model, preset, red, green, blue|
  if IGNORED_PRESETS.include? preset
    $stderr.puts "Ignoring preset '#{preset}'"
  else
    preset = '"'+preset+'"' if preset[-1] == "K"
    preset = preset + " "*(min_padding-preset.size)
    puts "  { \"#{maker}\", \"#{model}\", #{preset}, 0, { #{red}, #{green}, #{blue}, 0 } },"
  end
end

# vim: tabstop=2 expandtab shiftwidth=2 softtabstop=2
# kate: tab-width: 2; replace-tabs on; indent-width 2; tab-indents: off;
# kate: indent-mode ruby; remove-trailing-spaces modified;
