#!/usr/bin/env python3

# dji_fpv_srt_geojson.py
# 2023-02-18
# by Gernot Walzl

# The DJI FPV drone records SRT files to track GPS coordinates, altitudes, etc.
# This script converts the recorded track to GeoJSON.

import argparse
import datetime
import json
import re


class DjiFpvSrtFileReader:

    def __init__(self):
        self._pattern_dt = re.compile(r'(\d+\-\d+\-\d+ \d+:\d+:\d+\.\d+)')
        self._pattern_var = re.compile(r'\[([\w+ *: *[a-zA-Z0-9.\-/]+)\]')
        self._frame_infos = []
        self._frame_time = 0.0

    def read(self, filename):
        self._frame_infos.clear()
        with open(filename, 'r') as file_srt:
            frame_datetime = None
            for line in file_srt:
                str_dts = self._pattern_dt.findall(line)
                if str_dts:
                    frame_datetime = datetime.datetime.strptime(
                        str_dts[0], '%Y-%m-%d %H:%M:%S.%f')
                    continue
                str_vars = self._pattern_var.findall(line)
                if str_vars:
                    infos = dict()
                    infos['datetime'] = frame_datetime
                    for str_var in str_vars:
                        (varname, value) = str_var.split(':')
                        varname = varname.strip()
                        str_value = value.strip()
                        if varname in ['iso', 'fnum', 'ev', 'ct', 'focal_len']:
                            infos[varname] = int(str_value)
                        elif varname in ['latitude', 'longitude', 'altitude']:
                            infos[varname] = float(str_value)
                        else:
                            infos[varname] = str_value
                    self._frame_infos.append(infos)
        self._frame_time = (
            self._frame_infos[-1]['datetime'] -
            self._frame_infos[0]['datetime']).total_seconds() / (
            len(self._frame_infos)-1)

    def get(self, varname, step=1):
        result = []
        for cnt in range(0, len(self._frame_infos), step):
            result.append(self._frame_infos[cnt][varname])
        return result

    def get_frame_time(self):
        return self._frame_time


class GeoJsonFileWriter():

    def _to_geolinestring(self, longitudes, latitudes):
        coords = []
        for cnt in range(0, len(longitudes)):
            coords.append([longitudes[cnt], latitudes[cnt]])
        geolinestring = {
            "type": "Feature",
            "geometry": {
                "type": "LineString",
                "coordinates": coords
            }
        }
        return geolinestring

    def write(self, filename, longitudes, latitudes):
        with open(filename, 'w') as file_json:
            geolinestring = self._to_geolinestring(longitudes, latitudes)
            file_json.write(json.dumps(geolinestring))


def main(inputfile, outputfile):
    srtfilereader = DjiFpvSrtFileReader()
    srtfilereader.read(inputfile)
    step = 6
    latitudes = srtfilereader.get("latitude", step)
    longitudes = srtfilereader.get("longitude", step)
    filewriter = GeoJsonFileWriter()
    filewriter.write(outputfile, longitudes, latitudes)


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument("inputfile", type=str)
    parser.add_argument("outputfile", type=str)
    args = parser.parse_args()
    main(args.inputfile, args.outputfile)