#!/usr/bin/env python3

# System and standar imports
import argparse
import os
import os.path
import sys
from pathlib import Path
from glob import glob
from functools import partial

# ros imports
import rospy
from atom_core.naming import generateLabeledTopic
from geometry_msgs.msg import PointStamped
from sensor_msgs.msg import PointCloud2

# 3rd-party imports
from colorama import Back, Fore, Style
from pynput import keyboard

# Atom imports
from atom_calibration.dataset_playback.depth_manual_labeling import clickedPointsCallback, clickedPointsReset
from atom_calibration.dataset_playback.lidar3d_manual_labeling import *
from atom_calibration.dataset_playback.visualization import setupVisualization, visualizationFunction
from atom_core.dataset_io import (filterCollectionsFromDataset, filterSensorsFromDataset, loadResultsJSON,
                                  saveAtomDataset)

# own packages

# -------------------------------------------------------------------------------
# --- FUNCTIONS
# -------------------------------------------------------------------------------

# -------------------------------------------------------------------------------
# --- MAIN
# -------------------------------------------------------------------------------


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("-d", "--dataset",
                    help="Dataset folder.", type=str, required=True)
    ap.add_argument("-n", "--name", help="New name of dataset.",
                    required=True, type=str)

    # Roslaunch adds two arguments (__name and __log) that break our parser. Lets remove those.
    arglist = [x for x in sys.argv[1:] if not x.startswith('__')]
    args = vars(ap.parse_args(args=arglist))

    args['dataset'] = str(Path(args['dataset']))
    # ---------------------------------------
    # --- Verify that the dataset is ok
    # ---------------------------------------
    if not os.path.isdir(args['dataset']):
        sys.exit(Fore.RED + 'Dataset ' + args['dataset'] + ' does not exist.' + Style.RESET_ALL)

    # Get path of new dataset
    tmp = args['dataset'].split('/')
    new_dataset = '/'.join(tmp[:-1]) + '/' + args['name'] + '/'

    os.rename(args['dataset'], new_dataset)
    print('Dataset will be renamed to ' + Fore.BLUE + new_dataset + Style.RESET_ALL)

    # Get the json files in the dataset and change the name of dataset in the _metadata field.
    json_files = glob(new_dataset + '/*.json')
    json_files = [x for x in json_files if not 'annotation' in x]

    for json_file in json_files:
        # Loads a json file containing the detections. Returned json_file has path resolved by urireader.
        print('Loading json file ' + json_file)
        dataset, _ = loadResultsJSON(json_file)
        if '_metadata' in dataset:
            print('File contains metadata. Changing dataset_name ...')
            dataset['_metadata']['dataset_name'] = args['name']
            saveAtomDataset(json_file, dataset)


if __name__ == "__main__":
    main()