# ISC License
#
# Copyright 2019 Simon Bruder
#
# Permission to use, copy, modify, and/or distribute this software for any
# purpose with or without fee is hereby granted, provided that the above
# copyright notice and this permission notice appear in all copies.
#
# THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES WITH
# REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF MERCHANTABILITY
# AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR ANY SPECIAL, DIRECT,
# INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES WHATSOEVER RESULTING FROM
# LOSS OF USE, DATA OR PROFITS, WHETHER IN AN ACTION OF CONTRACT, NEGLIGENCE OR
# OTHER TORTIOUS ACTION, ARISING OUT OF OR IN CONNECTION WITH THE USE OR
# PERFORMANCE OF THIS SOFTWARE.
#
# This script is based on “aae-export.py” by Martin Herkt (lachs0r): https://gist.github.com/torque/6453947

bl_info = {
    'name': 'Export: CSV Tracking Data',
    'description': 'Export motion tracking data to CSV file',
    'author': 'Simon Bruder',
    'version': (0, 1, 0),
    'blender': (2, 80, 0),
    'location': 'Movie Clip Editor > Footage > CSV Tracking Data Export',
    'warning': '',
    'category': 'Import-Export',
}

import bpy
import csv
import math
import mathutils

def write_files(prefix, context):
    scene = context.scene
    fps = scene.render.fps / scene.render.fps_base
    clip = context.edit_movieclip
    real_scale = clip.real_length_pixels / clip.real_length_meters

    for trackno, track in enumerate(clip.tracking.tracks):
        with open(f'{prefix}_{trackno}.csv', 'w') as f:
            writer = csv.DictWriter(
                f,
                fieldnames=[
                    'time (frames)',
                    'time (s)',
                    'x (px)',
                    'x (m)',
                    'y (px)',
                    'y (m)',
                    'rotation'
                ]
            )
            writer.writeheader()

            startrot = None

            for marker in [track.markers.find_frame(frameno) for frameno in range(1, clip.frame_duration + 1)]:
                if not marker or marker.mute:
                    continue

                coords = marker.co
                corners = marker.pattern_corners

                p1 = mathutils.Vector(corners[0])
                p2 = mathutils.Vector(corners[1])
                mid = (p1 + p2) / 2
                diff = mid - mathutils.Vector((0,0))

                rotation = math.atan2(diff[0], diff[1]) * 180 / math.pi

                if startrot == None:
                    startrot = rotation
                    rotation = 0
                else:
                    rotation -= startrot - 360

                x = coords[0] * clip.size[0]
                y = (1 - coords[1]) * clip.size[1]

                writer.writerow({
                    'time (frames)': marker.frame,
                    'time (s)': marker.frame / fps,
                    'x (px)': x,
                    'x (m)': x / real_scale,
                    'y (px)': y,
                    'y (m)': y / real_scale,
                    'rotation': rotation % 360
                })

    return {'FINISHED'}

from bpy_extras.io_utils import ExportHelper
from bpy.props import StringProperty

class ExportCSVKey(bpy.types.Operator, ExportHelper):
    bl_idname = 'export.csvkey'
    bl_label = 'Export tracking data to CSV file'
    filename_ext = ''

    filter_glob: StringProperty(default='*', options={'HIDDEN'})

    def execute(self, context):
        return write_files(self.filepath, context)

class ExportCSVKeyPanel(bpy.types.Panel):
    bl_idname = 'CLIP_PT_tracking_export'
    bl_label = 'CSV Tracking Data Export'
    bl_space_type = 'CLIP_EDITOR'
    bl_region_type = 'UI'
    bl_category = 'Footage'
    bl_options = {'DEFAULT_CLOSED'}

    def draw(self, context):
        layout = self.layout

        col = self.layout.column(align = True)
        col.prop(context.edit_movieclip, 'real_length_pixels')
        col.prop(context.edit_movieclip, 'real_length_meters')
        col.operator('export.csvkey', text='Export')

def register():
    bpy.utils.register_class(ExportCSVKey)
    bpy.utils.register_class(ExportCSVKeyPanel)
    bpy.types.MovieClip.real_length_pixels = bpy.props.IntProperty(name='Real length (px)', default=100)
    bpy.types.MovieClip.real_length_meters = bpy.props.FloatProperty(name='Real length (m)', default=0.1, precision=3)

def unregister():
    bpy.utils.unregister_class(ExportCSVKey)
    del bpy.types.MovieClip.real_length_pixels
    del bpy.types.MovieClip.real_length_meters

if __name__ == '__main__':
    register()