#!/usr/bin/env python3

# dji_mini_subtitles.py
# 2023-09-14
# by Gernot Walzl

# The DJI Mavic Mini drone records subtitles to track GPS coordinates, altitudes, etc.
#
# To extract flight data stored in subtitles, ffmpeg can be used:
# ffmpeg -ss 00:00:10 -i DJI_0100.MP4 DJI_0100.srt

import argparse
import re


class DjiMiniSrtFileReader:

    def __init__(self):
        self._pattern_var = re.compile(r'F/[0-9.]+|[A-Z.]+ [0-9.\-]+|GPS \(.*\)')
        self._frame_infos = []

    def read(self, filename):
        self._frame_infos.clear()
        with open(filename, 'r') as file_srt:
            for line in file_srt:
                str_vars = self._pattern_var.findall(line)
                if str_vars:
                    infos = {}
                    for str_var in str_vars:
                        if str_var.startswith('F/'):
                            infos['F'] = float(str_var[2:])
                        elif str_var.startswith('GPS'):
                            str_value = str_var[5:-1].split(', ')
                            value = []
                            value.append(float(str_value[0]))
                            value.append(float(str_value[1]))
                            value.append(int(str_value[2]))
                            infos['GPS'] = value
                        else:
                            (varname, str_value) = str_var.split(' ')
                            if varname in ['F', 'SS', 'EV', 'D', 'H', 'H.S', 'V.S']:
                                infos[varname] = float(str_value)
                            elif varname in ['ISO']:
                                infos[varname] = int(str_value)
                            else:
                                infos[varname] = str_value
                    self._frame_infos.append(infos)

    def get(self, varname):
        result = []
        for frame_info in self._frame_infos:
            result.append(frame_info[varname])
        return result


class VttFileWriter:

    def _format_time(self, float_time):
        minutes, remainder = divmod(float_time, 60)
        seconds, remainder = divmod(remainder, 1)
        milliseconds = remainder * 1000
        return '{:02d}:{:02d}.{:03d}'.format(
            int(minutes), int(seconds), int(milliseconds))

    def write(self, filename, subtitles):
        with open(filename, 'w') as file_vtt:
            file_vtt.write("WEBVTT\n")
            cnt = 0
            time_start = 0.0
            for subtitle in subtitles:
                cnt += 1
                file_vtt.write("\n")
                time_end = float(cnt)
                str_time_s = self._format_time(time_start)
                str_time_e = self._format_time(time_end)
                file_vtt.write("{} --> {}\n".format(str_time_s, str_time_e))
                time_start = time_end
                file_vtt.write("{}\n".format(subtitle))


def main(inputfile, outputfile):
    srtfilereader = DjiMiniSrtFileReader()
    srtfilereader.read(inputfile)
    heights = srtfilereader.get("H")
    ground_speeds = srtfilereader.get("H.S")
    vert_speeds = srtfilereader.get("V.S")

    subtitles = []
    for cnt in range(0, len(ground_speeds)):
        vert_speed = vert_speeds[cnt]
        vert_speed_ud = u'↑'
        if vert_speed < 0.0:
            vert_speed *= -1.0
            vert_speed_ud = u'↓'
        sub = u"h={:04.1f}m {}{:04.1f}m/s    v={:04.1f}m/s".format(
            heights[cnt], vert_speed_ud, vert_speed,
            ground_speeds[cnt])
        subtitles.append(sub)

    filewriter = VttFileWriter()
    filewriter.write(outputfile, subtitles)


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)