#!/usr/bin/python3

"""
Copyright 2015, 2017 Lukas Toggenburger; https://github.com/ltog/photomapping-utils

This program is free software: you can redistribute it and/or modify
it under the terms of the GNU Affero General Public License as published by
the Free Software Foundation, either version 3 of the License, or
(at your option) any later version.

This program is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
GNU Affero General Public License for more details.

You should have received a copy of the GNU Affero General Public License
along with this program.  If not, see <http://www.gnu.org/licenses/>.
"""

# based on:
# - https://docs.python.org/2/library/xml.etree.elementtree.html#module-xml.etree.ElementTree

import sys
import os.path
import pprint # for debugging/printing dictionaries https://docs.python.org/2/library/pprint.html
import argparse
import re
import dateutil.parser
import datetime
from lxml import etree as ET

def main(args):
    for infile in args.infiles:
        if not os.path.exists(infile):
            sys.exit("ERROR: File '" + infile + "' does not exist. Aborting...")

    if os.path.exists(args.outfile):
        sys.exit("ERROR: File '" + args.outfile + "' exists. Aborting...")

    tree = [] # holds input data, one array-entry per input file

    for i in range(len(args.infiles)):
        tree.append(ET.parse(args.infiles[i]))

    data = {} # stores all data before writing new xml file

    coords = []

    for i in range(len(args.infiles)):
        coords.append(xml2dict(tree[i]))
        coords[i] = interpolate_coords(coords[i])
        insert_into_data(coords[i], data)

    # debugging
    #print("data = ")
    #pp = pprint.PrettyPrinter(indent=4)
    #pp.pprint(data)

    write_xml(args.outfile, data)

def write_xml(outfile: str, data):
    # create root element
    root = ET.Element("gpx")
    root.set("version", "1.0")
    root.set("xmlns", "http://www.topografix.com/GPX/1/0")

    # create further elements
    bounds = ET.SubElement(root, "bounds")
    trk    = ET.SubElement(root, "trk")
    trkseg = ET.SubElement(trk, "trkseg")

    # create new document tree
    doctree = ET.ElementTree(root)

    # write track points
    for time in sorted(data):
        lat = str(data[time]["lat"])
        lon = str(data[time]["lon"])

        trkpt = ET.SubElement(trkseg, "trkpt")
        trkpt.set("lat", lat)
        trkpt.set("lon", lon)

        time_element = ET.SubElement(trkpt, "time")
        time_element.text = time

        if data[time]["ele_count"] > 0:
            ele = str(data[time]["ele"])
            ele_element  = ET.SubElement(trkpt, "ele")
            ele_element.text  = ele

    # write xml file
    doctree.write(args.outfile, xml_declaration=True, encoding="UTF-8", pretty_print=True)

def xml2dict(tree):
    namespace = get_namespace(tree)

    coords = {}

    # iterate over each <trkpt>
    for el in tree.findall('.//{0}trkpt'.format(namespace)): # according to http://stackoverflow.com/a/1319417
        time  = el.find(".//{0}time".format(namespace)).text
        lat   = float(el.attrib['lat'])
        lon   = float(el.attrib['lon'])

        if time in coords:
            continue

        coords[time] = {"lat": lat, "lon": lon}

        ele_element = el.find(".//{0}ele".format(namespace))
        if ele_element != None: # the current trackpoint has an <ele> tag
            coords[time]["ele"] = float(ele_element.text)

    return coords

def interpolate_coords(coords):
    for time in sorted(coords):
        # is there a set of coordinates for the next second?
        if (datestring_add_secs(time, 1) in coords):
            continue # no need for interpolation
        
        for i in range(2, args.max_interpolate+1):
            upper = datestring_add_secs(time, i)
            if (upper in coords):
                coords = interpolate_between(coords, time, upper, i)
                break
        
    return coords

def interpolate_between(coords, first: str, last: str, diff):
    for i in range(1, diff):
        cur_datestring = datestring_add_secs(first, i)
        cur_lat = ((diff-i)/diff)*coords[first]["lat"] + (i/diff)*coords[last]["lat"]
        cur_lon = ((diff-i)/diff)*coords[first]["lon"] + (i/diff)*coords[last]["lon"]

        coords[cur_datestring] = {"lat": cur_lat, "lon": cur_lon}

        if ("ele" in coords[first] and "ele" in coords[last]):
            cur_ele = ((diff-i)/diff)*coords[first]["ele"] + (i/diff)*coords[last]["ele"]
            coords[cur_datestring]["ele"] = cur_ele

    return coords

def string2date(date: str):
    return dateutil.parser.parse(date) # https://stackoverflow.com/a/3908349

def date2string(date):
    return date.strftime("%Y-%m-%dT%H:%M:%SZ") # https://stackoverflow.com/a/7999977

def datestring_add_secs(date: str, summand):
    later = string2date(date)+datetime.timedelta(seconds=summand) # https://stackoverflow.com/a/100345
    return date2string(later)

def insert_into_data(coords, data):
    for time in sorted(coords):
        if time in data: # we already have encountered another trackpoint with this timestamp
            count     = data[time]["count"]
            ele_count = data[time]["ele_count"]

            data[time]["lat"]   = (count*float(data[time]["lat"])+coords[time]["lat"]) / (count+1)
            data[time]["lon"]   = (count*float(data[time]["lon"])+coords[time]["lon"]) / (count+1)
            data[time]["count"] += 1
 
            if "ele" in coords[time]: # the current trackpoint has an <ele> tag
                if ele_count > 0:
                    data[time]["ele"] = (ele_count*float(data[time]["ele"])+coords[time]["ele"]) / (ele_count+1)
                else:
                    data[time]["ele"] = coords[time]["ele"]

                data[time]["ele_count"] = ele_count + 1

        else: # first trackpoint with this timestamp
            data[time] = {"lat": coords[time]["lat"], "lon": coords[time]["lon"], "count": 1}

            if "ele" in coords[time]: #the current trackpoint has an <ele> tag
                data[time]["ele"]       = coords[time]["ele"]
                data[time]["ele_count"] = 1
            else:
                data[time]["ele_count"] = 0
                

# PRE: assume only one namespace is in use
def get_namespace(tree):
    # see https://stackoverflow.com/questions/9513540/python-elementtree-get-the-namespace-string-of-an-element
    # see https://stackoverflow.com/a/1319575
    # see https://stackoverflow.com/a/20104763

    #namespace = tree.getroot().get("version") # crashes when looking for attribute 'xmlns', instead of 'version'

    re_match = re.match('\{.*\}', tree.getroot().tag)
    if re_match:
        return re_match.group(0)
    else:
        return ''


if __name__ == '__main__':

    parser = argparse.ArgumentParser()
    parser.add_argument("-i", "--infiles", help="input files in gpx format, use like this: '-i file1.gpx file2.gpx ...'", nargs="+", required=True)
    parser.add_argument("-o", "--outfile", help="the output file", required=True)
    parser.add_argument("-m", "--max_interpolate", help="interpolate time differences up to this many seconds in a given gpx file (example: max_interpolate=3 would interpolate the values 12 and 13 between the second-values 11 and 14", default=3, type=int)
    args = parser.parse_args()

    main(args)
