# coding: utf-8
'''
Copyright (C) 2024 Kurochkin Ilia @brothermechanic
The Original Code by: Alexander Shushpanov @shalenikol
Created by brothermechanic

    This program is free software: you can redistribute it and/or modify
    it under the terms of the GNU 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 General Public License for more details.

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

import collections
import logging
import os
from pathlib import Path
from rosbags.highlevel import AnyReader

logger = logging.getLogger(__name__)
logging.basicConfig(level=logging.INFO)


def append_frame(idx: int, frames, links, frame_data) -> int:
    ''' Append datalist for frame and increase frame index. '''
    frame_data['links'] = links
    frames.append(frame_data)
    return idx + 1


def gimbal_locks_filter(ros2bag_func):
    ''' Finding and fix gimbal locks. '''
    def wrapper(*args, **kwargs):
        # finding gimbal locks
        ros2bag_data = ros2bag_func(*args, **kwargs)
        gimbal_locks = {}
        for link in ros2bag_data['scene_links']:
            gimbal_data = collections.defaultdict(list)
            link_rotation = []
            for frame_data in ros2bag_data['frames_data']:
                link_rotation_pre = []
                for link_data in frame_data['links']:
                    if link_data['link'] == link:
                        link_rotation_pre = link_rotation
                        link_rotation = link_data['rot_wxyz']
                if link_rotation_pre:
                    idx = 0
                    for axis_pre, axis_cur in zip(link_rotation_pre, link_rotation):
                        if (round(axis_cur, 2) > 0 > round(axis_pre, 2)
                                or round(axis_cur, 2) < 0 < round(axis_pre, 2)):
                            gimbal_data[frame_data['id']].append(idx)
                            gimbal_locks[link] = gimbal_data
                            logger.info(
                                'Gimbal lock detected for %s link %s axis at %s frame',
                                link, idx, frame_data['id'])
                        idx += 1

        # fix gimbal locks
        for link in gimbal_locks:
            gimbal_frame = list(gimbal_locks[link])[0]
            if len(gimbal_locks[link]) == 2:
                gimbal_range = list(gimbal_locks[link])
            else:
                at_start = abs(ros2bag_data['frame_start'] - gimbal_frame)
                at_end = abs(ros2bag_data['frame_end'] - gimbal_frame)
                if at_start < at_end:
                    gimbal_range = [ros2bag_data['frame_start'], gimbal_frame]
                else:
                    gimbal_range = [gimbal_frame, ros2bag_data['frame_end']]
            for frame_data in ros2bag_data['frames_data'][gimbal_range[0]:gimbal_range[1]]:
                for link_data in frame_data['links']:
                    if link_data['link'] == link:
                        for axis in gimbal_locks[link][gimbal_frame]:
                            link_data['rot_wxyz'][axis] = (
                                -1 * link_data['rot_wxyz'][axis])
                            logger.debug(
                                'Gimbal lock fixed for %s link %s axis %s',
                                link, axis, frame_data['id'])
            logger.info('Gimbal lock fixed for %s link', link)

        return ros2bag_data

    return wrapper


@gimbal_locks_filter
def get_animation_data(ros2bag_path: str,
                       frame_start=0, frame_end=0, fps=30, target_topic='/tf') -> dict:
    ''' Get animation data from Ros2Bag database '''

    assert ros2bag_path.endswith('.db3'), (
        'Please, check Ros2Bag file format and extension!')

    scene_links = set()
    frames_data = []
    ros2bag_dir = os.path.split(ros2bag_path)[0]

    with AnyReader([Path(ros2bag_dir.replace('\\', '/'))]) as ros2bag:
        targets = [
            x for x in ros2bag.connections if x.topic == target_topic]
        # TODO Switch to timestamp instead of frame index
        idx = 0
        idx_pre = - 1
        for connection, timestamp, rawdata in ros2bag.messages(connections=targets):
            if idx != idx_pre:
                frame_data = {'id': idx, 'timestamp': timestamp}
                key_link = []
                links = []
                idx_pre = idx

            tf_msg = ros2bag.deserialize(rawdata, connection.msgtype)
            for t in tf_msg.transforms:
                c_key = t.header.frame_id + t.child_frame_id
                if key_link.count(c_key) > 0:
                    idx = append_frame(idx, frames_data, links, frame_data)
                else:
                    key_link.append(c_key)
                    scene_links.add(t.child_frame_id)
                    if t.header.frame_id != 'world':
                        scene_links.add(t.header.frame_id)
                    link_data = {}
                    link_data['parent'] = t.header.frame_id
                    link_data['link'] = t.child_frame_id
                    link_data['loc_xyz'] = [
                        t.transform.translation.x,
                        t.transform.translation.y,
                        t.transform.translation.z]
                    link_data['rot_wxyz'] = [
                        t.transform.rotation.w,
                        t.transform.rotation.x,
                        t.transform.rotation.y,
                        t.transform.rotation.z]
                    links.append(link_data)

            if frame_end and not idx < frame_end + 1:
                break

    assert len(frames_data) - 1 == frames_data[-1]['id'], (
        'Ros2Bag database has dublicated frames!')

    frame_end = frames_data[-1]['id']

    return {'path': ros2bag_path,
            'frame_start': frame_start,
            'frame_end': frame_end,
            'fps': fps,
            'scene_links': list(scene_links),
            'frames_data': frames_data}


if __name__ == '__main__':
    import json
    in_file = '/media/disk/robossembler/project/collab/138-rosbag-to-blender/cg/blender/scripts/addons/rosbag-importer/rosbag2/rosbag2_2024_02_15-18_15_42_0.db3'
    out_file = '/media/disk/robossembler/project/collab/138-rosbag-to-blender/cg/blender/scripts/addons/rosbag-importer/rosbag2/rosbag2_2024_02_15-18_15_42_0.json'
    with open(out_file, "w", encoding='utf-8') as fh:
        json.dump(get_animation_data(in_file), fh, indent=4)
