#!/usr/bin/python3
# direct translation of extract_wb in python using as little external deps as possible

from __future__ import print_function
import sys
import os
import xml.etree.ElementTree as ET
import subprocess
import shlex

def eprint(*args, **kwargs):
    print(*args, file=sys.stderr, **kwargs)

if len(sys.argv) < 2 :
    sys.exit("Usage: extract_wb <file1> [file2] ...")

IGNORED_PRESETS = {"Auto", "Kelvin", "Measured", "AsShot", "As Shot", "Preset",
                 "Natural Auto", "Multi Auto", "Color Temperature Enhancement",
                 "One Touch WB 1", "One Touch WB 2", "One Touch WB 3",
                 "One Touch WB 4", "Custom WB 1", "Auto0", "Auto1", "Auto2",
                 "Custom", "CWB1", "CWB2", "CWB3", "CWB4", "Black", 
                 "Illuminator1", "Illuminator2", "Uncorrected"}

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",
  "Fine Weather" : "DirectSunlight",
  "Tungsten (Incandescent)" : "Tungsten",
  "ISO Studio Tungsten" : "Tungsten",
  "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 = ["DirectSunlight", "Daylight", "D55", "Shade","Cloudy",
              "Tungsten", "Incandescent","Fluorescent", 
              "WarmWhiteFluorescent", "CoolWhiteFluorescent",
              "DayWhiteFluorescent","DaylightFluorescent",
              "DaylightFluorescent", "NeutralFluorescent", "WhiteFluorescent",
              "HighTempMercuryVaporFluorescent", "HTMercury", 
              "SodiumVaporFluorescent", "Underwater", "Flash", "Unknown"]

PRESET_SORT_MAPPING = {}

for index,name in enumerate(PRESET_ORDER):
    PRESET_SORT_MAPPING[name] = index + 1

cams_from_source = os.path.dirname(os.path.abspath(__file__)) + "/../src/external/rawspeed/data/cameras.xml"
cams_from_dist = os.path.dirname(os.path.abspath(__file__)) + "/../rawspeed/cameras.xml"

CAMERAS = os.path.abspath(cams_from_source) if os.path.exists(os.path.abspath(cams_from_source)) else os.path.abspath(cams_from_dist)

if not os.path.exists(CAMERAS):
    sys.exit("Can't find cameras mapping file, should be in {0}".format(CAMERAS))

exif_name_map = {}
xml_doc = ET.parse(CAMERAS)
for camera in xml_doc.getroot().findall('Camera'):
    maker = exif_maker = camera.get('make')
    model = exif_model = camera.get('model')
    exif_id = maker,model
    if camera.find('ID') is not None:
        cid = camera.find('ID')
        maker = cid.get('make')
        model = cid.get('model')
    exif_name_map[exif_id] = maker,model
    for alias in camera.findall('Aliases/Alias'):
        exif_model = alias.text
        exif_id = exif_maker, exif_model
        exif_name_map[exif_id] = maker,model

found_presets = []

for filename in sys.argv[1:]:
    red = green = blue = maker = model = preset = None
    finetune = fl_count = rlevel = blevel = glevel = 0
    listed_presets = []
    preset_names = {}
    gm_skew = False
    command = "exiftool -Make -Model \"-WBType*\" \"-WB_*\" \"-ColorTemp*\"    "\
    "-WhiteBalance -WhiteBalance2 -WhitePoint -ColorCompensationFilter       "\
    "-WBShiftAB -WBShiftAB_GM -WBShiftAB_GM_Precise -WBShiftGM -WBScale      "\
    "-WhiteBalanceFineTune -WhiteBalanceComp -WhiteBalanceSetting            "\
    "-WhiteBalanceBracket -WhiteBalanceBias -WBMode -WhiteBalanceMode        "\
    "-WhiteBalanceTemperature -WhiteBalanceDetected -ColorTemperature        "\
    "-WBShiftIntelligentAuto -WBShiftCreativeControl -WhiteBalanceSetup      "\
    "-WBRedLevel -WBBlueLevel -WBGreenLevel -RedBalance -BlueBalance         "\
    "\"{0}\"".format(filename)
    if filename.endswith(('.txt','.TXT')):
        command = 'cat "{0}"'.format(filename)
    command = shlex.split(command)
    proc = subprocess.check_output(command, universal_newlines=True)
    for io in proc.splitlines():
        lineparts = io.split(':')
        tag = lineparts[0].strip()
        values = lineparts[1].strip().split(' ')
        if 'Make' in tag.split():
            maker = lineparts[1].strip()
        elif 'Model' in tag.split():
            model = lineparts[1].strip()
        elif tag == "WB RGGB Levels":
            green = (float(values[1])+float(values[2]))/2.0
            red = float(values[0])/green
            blue = float(values[3])/green
            green = 1
        elif tag == "WB RB Levels":
            red = float(values[0])
            blue = float(values[1])
            if len(values) == 4 and values[2] == "256" and values[3] == "256":
                red /= 256.0
                blue /= 256.0
            green = 1
        elif tag == "WB GRB Levels":
            green = float(values[0])
            red = float(values[1])/green
            blue = float(values[2])/green
            green = 1
        # elif tag == "WB GRB Levels Auto" and maker == "FUJIFILM" # fuji seems to use "WB GRB Levels Auto to describe manual finetuning
        #  green = float(values[0])
        #  red = float(values[1])/green
        #  blue = float(values[2])/green
        #  green = 1
        elif tag == "White Point" and len(values) > 3:
            green = (float(values[1])+float(values[2]))/2.0
            red = float(values[0])/green
            blue = float(values[3])/green
            green = 1
        elif tag == "White Balance" or tag == "White Balance 2":
            preset = ' '.join(values)
            if preset in FL_PRESET_REPLACE:
                preset = FL_PRESET_REPLACE[preset]
        elif ' '.join(tag.split()[:2]) == "WB Type":
            preset_names[' '.join(tag.split()[2:])] = ' '.join(values)
        elif ' '.join(tag.split()[:3]) in ['WB RGB Levels', 'WB RGGB Levels', 'WB RB Levels']:
            # todo - this codepath is weird
            p = ''.join(tag.split()[3:])
            if( p in preset_names):
                p = preset_names[p]

            r=g=b=0

            if len(values) == 4 and ' '.join(tag.split()[:3]) in ['WB RB Levels']:
                g = (float(values[2])+float(values[3]))/2.0
                r = float(val
ues[0])/g
                b = float(values[1])/g
                g = 1
            elif len(values) == 4:
                g = (float(values[1])+float(values[2]))/2.0
                r = float(values[0])/g
                b = float(values[3])/g
                g = 1
            elif len(values) == 3:
                g = float(values[1])
                r = float(values[0])/g
                b = 
float(values[2])/g
                g = 1
            elif len(values) == 2 and ' '.join(tag.split()[:3]) in ['WB RB Levels']:
                r = float(values[0])
                b = float(values[2])
                g = 1
            else:
                eprint("Found RGB tag '{0}' with {1} values instead of 2, 3 or 4".format(p, len(values)))
            
            if 'Fluorescent' in p:
                fl_count += 1
            
            if not p:
                p= 'Unknown'
            if p not in IGNORED_PRESETS:
                listed_presets.append(tuple([p,r,g,b]))
        elif tag == "WB Red Level":
            rlevel = float(values[0])
        elif tag == "WB Blue Level":
            blevel = float(values[0])
        elif tag == "WB Green Level":
            glevel = float(values[0])