#!/home/enigma/.virtualenvs/gpx_clean/bin/python

import sys

import numpy as np
import gpxpy


def clean_segment(segment, badness=4):
    """Lazily strip out anything not in IQR"""
    new_points = []
    speeds = []
    for index, point in enumerate(segment.points[1:]):
        prev = segment.points[index-1]
        if point.time == prev.time:
            continue
        speeds.append(point.distance_2d(prev) / (point.time - prev.time).total_seconds())
    # Calculate IQR
    if not speeds:
        return
    speeds.sort()
    speed_mean = np.mean(speeds)
    speed_mu = np.std(speeds)
    if speed_mean * 3.6 < 5:
        speed_mean = 5 / 3.6
    def Z(speed):
        return (speed - speed_mean) / speed_mu
    on_bad = False
    new_points = [segment.points[0]]
    bad_set = []
    last_good = None
    for index, point in enumerate(segment.points[1:]):
        if len(bad_set) > 10:
            new_points.extend(bad_set)
            bad_set = []
            on_bad = False
        if on_bad:
            prev = bad_set[0]
        else:
            prev = segment.points[index-1]
        speed = point.distance_2d(prev) / (point.time - prev.time).total_seconds()
        this_bad = Z(speed)
        if on_bad:
            if this_bad >= badness:
                # See if this large shift puts us back on track
                new_speed = last_good.distance_2d(point) / (last_good.time - point.time).total_seconds()
                if this_bad < badness:
                    # We're going to go back, so discard this bad set
                    bad_set = []
                    on_bad = False
            else:
                bad_set.append(point)
        else:
            if this_bad >= badness:
                bad_set.append(point)
                on_bad = True
                last_good = prev
            else:
                new_points.append(point)
    #print(f"Starting with {len(segment.points)} pts, now have {len(new_points)}")
    segment.points = new_points


def clean_track(track):
    for segment in track.segments:
        clean_segment(segment)


def show_moving_data(pre, md):
    speed = md.moving_distance / md.moving_time * 3.6
    print(f"{pre}: Moved at {speed:.2f} km/h")


def clean_gpx(gpx):
    for track in gpx.tracks:
        clean_track(track)


def clean_file(filename):
    with open(filename, 'r') as infile:
        gpx = gpxpy.parse(infile)
        md = gpx.get_moving_data()
        show_moving_data("Before clean", md)
        clean_gpx(gpx)
        md = gpx.get_moving_data()
        show_moving_data("After clean", md)
    with open(f"{filename[:-4]}_cleaned.gpx", "w") as outfile:
        outfile.write(gpx.to_xml())


if __name__ == "__main__":
    for filename in sys.argv[1:]:
        clean_file(filename)