#!/usr/bin/env python

"""
   Copyright 2016 Alex Cortelyou

   Licensed under the Apache License, Version 2.0 (the "License");
   you may not use this file except in compliance with the License.
   You may obtain a copy of the License at

       http://www.apache.org/licenses/LICENSE-2.0

   Unless required by applicable law or agreed to in writing, software
   distributed under the License is distributed on an "AS IS" BASIS,
   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
   See the License for the specific language governing permissions and
   limitations under the License.
"""

import os
import argparse
import xml.etree.ElementTree as ET
from datetime import datetime
from glob import glob
import exifread
#arguments
parser = argparse.ArgumentParser(description='Extracts DJI image metadata into a Papywizard panohead data file for use with stitching software',
                                 formatter_class=argparse.ArgumentDefaultsHelpFormatter)
parser.add_argument('file', metavar='INPUT', nargs='+',
                    help='image files to process')
parser.add_argument('-t', dest='title', default='Untitled Panorama',
                    help='title for panohead file')
parser.add_argument('-c', dest='comment',
                    default='Generated by ' + os.path.basename(__file__),
                    help='comment for panohead file')
parser.add_argument('-o', dest='output', default='pano.xml',
                    help='filename for output')
parser.add_argument('--no-bracket', dest='no_bracket', action='store_true',
                    help='ignore hdr bracketing information')
args = parser.parse_args()

#constants
djiDateFormat = "%Y:%m:%d %H:%M:%S"
panDateFormat = "%Y-%m-%d_%Hh%Mm%Ss"
tagTime = 'EXIF DateTimeOriginal'
tagBias = 'EXIF ExposureBiasValue'
tagFocal = 'EXIF FocalLengthIn35mmFilm'
tagFile = 'Image ImageDescription'
tagNotes = 'Image ApplicationNotes'
tagGimbalPitch = '{http://www.dji.com/drone-dji/1.0/}GimbalPitchDegree'
tagGimbalRoll = '{http://www.dji.com/drone-dji/1.0/}GimbalRollDegree'
tagGimbalYaw = '{http://www.dji.com/drone-dji/1.0/}GimbalYawDegree'

#variables
shots = []
brackets = {}
startTime = None
endTime = None
focal = None

#process files
for a in args.file:
    for f in glob(a):
        #extract metadata
        handle = open(f, 'rb')
        exif = exifread.process_file(handle, debug=True)
        exif.update(ET.fromstring(exif[tagNotes].printable)[0][0].attrib)
        handle.close()

        #check times
        time = datetime.strptime(exif[tagTime].printable, djiDateFormat)
        if not startTime or startTime > time:
            startTime = time
        if not endTime or endTime < time:
            endTime = time

        #check bracket
        bias = str(exif[tagBias])
        if bias not in brackets.keys():
            brackets[bias] = len(brackets) + 1

        #check focal
        if not focal:
            focal = exif[tagFocal]

        #add shot to list
        shots.append({
            "time": time,
            "bracket": bias,
            "file": str(exif[tagFile]),
            "roll": str(exif[tagGimbalRoll]),
            "yaw": str(exif[tagGimbalYaw]),
            "pitch": str(exif[tagGimbalPitch]),
        })

#quit if nothing to output
if len(shots) == 0:
    print "no shots found"
    quit()

#sort the output by order taken
shots = sorted(shots, key=lambda k: k['file'])

#create panohead document
root = ET.Element("papywizard", version="c")
header = ET.SubElement(root, "header")
general = ET.SubElement(header, "general")
ET.SubElement(general, "title").text = args.title
ET.SubElement(general, "gps").text = ''
ET.SubElement(general, "comment").text = args.comment
shooting = ET.SubElement(header, "shooting", mode="mosaic")
ET.SubElement(shooting, "headOrientation").text = "up"
ET.SubElement(shooting, "cameraOrientation").text = "landscape"
ET.SubElement(shooting, "stabilizationDelay").text = "1"
ET.SubElement(shooting, "counter").text = "1"
ET.SubElement(shooting, "startTime").text = startTime.strftime(panDateFormat)
ET.SubElement(shooting, "endTime").text = endTime.strftime(panDateFormat)
camera = ET.SubElement(header, "camera")
ET.SubElement(camera, "timeValue").text = "1"
ET.SubElement(camera, "bracketing", nbPicts=str(len(brackets)) if not args.no_bracket else '1')
ET.SubElement(camera, "sensor", coef="1.0", ratio="4:3")
lens = ET.SubElement(header, "lens", type='rectilinear')
ET.SubElement(lens, "focal").text = str(focal)
mosaic = ET.SubElement(header, "mosaic")
ET.SubElement(mosaic, "nbPicts", pitch="1", yaw="1")
ET.SubElement(mosaic, "overlap", minimum="0.5", pitch="1.0", yaw="0.5")
shoot = ET.SubElement(root, "shoot")

#add each shot to document
count = 0
for shot in shots:
    count += 1
    bracket = str(brackets[shot['bracket']]) if not args.no_bracket else '1'
    pict = ET.SubElement(shoot, "pict", id=str(count), bracket=bracket)
    ET.SubElement(pict, "time").text = shot['time'].strftime(panDateFormat)
    ET.SubElement(pict, "position", pitch=shot['pitch'], roll=shot['roll'], yaw=shot['yaw'])

#write xml to file
ET.ElementTree(root).write(args.output, encoding='utf8', method='xml')

#success
print count, 'shots processed'

"""

'EXIF DateTimeOriginal'

'EXIF ExposureMode'
'EXIF ExposureIndex'
'EXIF ExposureProgram'
'EXIF ExposureTime'

'EXIF MeteringMode'
'EXIF ShutterSpeedValue'

'GPS GPSAltitude'
'GPS GPSLatitude'
'GPS GPSLongitude'

'{http://ns.adobe.com/tiff/1.0/}Make'
'{http://ns.adobe.com/tiff/1.0/}Model'

'{http://www.dji.com/drone-dji/1.0/}GimbalRollDegree'
'{http://www.dji.com/drone-dji/1.0/}FlightRollDegree'

'{http://www.dji.com/drone-dji/1.0/}GimbalPitchDegree'
'{http://www.dji.com/drone-dji/1.0/}FlightPitchDegree'

'{http://www.dji.com/drone-dji/1.0/}GimbalYawDegree'
'{http://www.dji.com/drone-dji/1.0/}FlightYawDegree'

"""

# pylint: disable=C0103
