#!/usr/bin/env python3

# stdlib
import os
import argparse
from os.path import exists
import numpy as np

import matplotlib


# 3rd-party
import numpy
import rospkg
import rosbag
from pytictoc import TicToc

# local packages
from atom_core.utilities import atomPrintOK, atomError, atomStartupPrint
from atom_core.system import execute
from colorama import Style, Fore
from graphviz import Digraph
import networkx as nx
import graphviz
from urdf_parser_py.urdf import URDF
from matplotlib import cm

from atom_core.config_visualization import *
import atom_core.config_io
import atom_core.drawing


if __name__ == "__main__":
    # Parse command line arguments
    ap = argparse.ArgumentParser()
    ap.add_argument("-n", "--name", help='package name', type=str, required=True)
    ap.add_argument("-cgt", "--collect_ground_truth", action="store_true",
                    help='Assume transformations (\\tf and \\tf_static) in the bag file to be perfect, and collect them as ground truth for ATOM datasets. Useful for evaluating the calibration against ground truth.')
    ap.add_argument("-cfg", "--config_file", help='Specify if you want to configure the calibration package with a specific configuration file. If this flag is not given, the standard config.yml ill be used.',
                    type=str, required=False, default=None)
    args = vars(ap.parse_args())

    # --------------------------------------------------------------------------
    # Initial setup
    # --------------------------------------------------------------------------
    tictoc = TicToc()
    tictoc.tic()

    package_name = os.path.basename(args['name'])
    atomStartupPrint('View calibration config ' + Fore.BLUE + package_name + Style.RESET_ALL)

    rospack = rospkg.RosPack()
    atom_calibration_path = rospack.get_path('atom_calibration')

    if not package_name in rospack.list():  # Check if package is under $ROS_PACKAGE_PATH, abort if not
        atomError('ROS package ' + Fore.BLUE + package_name + Style.RESET_ALL + ' not found under ROS. Are you sure the package in under a directory listed in $ROS_PACKAGE_PATH? Can you run:\n\n' +
                  Fore.BLUE + 'roscd ' + package_name + Style.RESET_ALL + '\n\nPlease fix this before running your package configuration.')

    package_path = rospack.get_path(package_name)  # full path to the package, including its name.
    package_base_path = os.path.dirname(package_path)  # parent path where the package is located

    # --------------------------------------------------------------------------
    # Read the config.yml file
    # --------------------------------------------------------------------------
    if args['config_file'] is None:
        args['config_file'] = package_path + '/calibration/config.yml'
    else:
        args['config_file'] = package_path + '/calibration/' + args['config_file']
        if not exists(args['config_file']):
            args['config_file'] = package_path + '/calibration/config.yml'

    print('Loading config_file ' + Fore.BLUE + str(args['config_file']) + Style.RESET_ALL)
    config = atom_core.config_io.loadConfig(args['config_file'])

    # Sensors colormap. Access with:  color_map_sensors[idx, :]
    cm_sensors = cm.Set3(numpy.linspace(0, 1, len(config['sensors'].keys())))

    # --------------------------------------------------------------------------
    # Read the bag file
    # --------------------------------------------------------------------------
    bag_file, _, bag_file_rel = atom_core.config_io.uriReader(config['bag_file'])
    print('Loading bagfile ' + bag_file + ' ... ', end='')
    bag = rosbag.Bag(bag_file)
    atomPrintOK()
    bag_info = bag.get_type_and_topic_info()
    bag_types = bag_info[0]
    bag_topics = bag_info[1]

    # Get initial stamp to compute mission time
    for topic, msg, stamp in bag.read_messages():
        bag_initial_stamp = stamp
        break

    # --------------------------------------------------------------------------
    # Setup the description file
    # --------------------------------------------------------------------------
    description_file, _, _ = atom_core.config_io.uriReader(config['description_file'])
    description_file_out_initial_estimate = package_path + '/urdf/initial_estimate.urdf.xacro'
    execute('cp ' + description_file + ' ' + description_file_out_initial_estimate,
            verbose=False)  # Copy the xacro to the initial_estimate file

    # Check the description file
    urdf_file = '/tmp/description.urdf'
    if os.path.exists(urdf_file):
        # print('Deleting temporary file ' + urdf_file)
        os.remove(urdf_file)

    print('Parsing description file ' + Fore.BLUE + description_file + Style.RESET_ALL)
    xacro_cmd = 'xacro ' + description_file + ' -o ' + urdf_file
    execute(xacro_cmd, verbose=True)  # create tmp urdf file

    if not os.path.exists(urdf_file):
        atomError('Could not parse description file ' + Fore.BLUE + description_file + Style.RESET_ALL + '\nYou must manually run command:\n' +
                  Fore.BLUE + xacro_cmd + Style.RESET_ALL + '\nand fix the problem before configuring your calibration package.')

    description = URDF.from_xml_file(urdf_file)  # read the urdf file

    # --------------------------------------------------------------------------
    # Create a tf graph to support the creation of the graphviz object
    # --------------------------------------------------------------------------
    nx_graph = createNxGraph(args, description, config, bag)

    # --------------------------------------------------------------------------
    # Create dot calibration graph (used for printing out a summary)
    # --------------------------------------------------------------------------
    dot_graph = createDotGraph(nx_graph, config)

    # dot_graph.view()
    dot_graph.render(filename='summary', directory=package_path + '/calibration', cleanup=True)