"""
DeepLabCut2.2 Toolbox (deeplabcut.org)
© A. & M. Mathis Labs
https://github.com/AlexEMG/DeepLabCut
Please see AUTHORS for contributors.

https://github.com/AlexEMG/DeepLabCut/blob/master/AUTHORS
Licensed under GNU Lesser General Public License v3.0
"""

import os
from pathlib import Path

import deeplabcut
from deeplabcut.utils import auxiliaryfunctions


def load_demo_data(config, createtrainingset=True):
    """
  Loads the demo data. Make sure that you are in the same directory where you have downloaded or cloned the deeplabcut.

  Parameter
    ----------
    config : string
        Full path of the config.yaml file of the provided demo dataset as a string.

    Example
    --------
    >>> deeplabcut.load_demo_data('config.yaml')
    --------
  """
    config = Path(config).resolve()
    config = str(config)

    transform_data(config)
    if createtrainingset:
        print("Loaded, now creating training data...")
        deeplabcut.create_training_dataset(config, num_shuffles=1)


def transform_data(config):
    """
    This function adds the full path to labeling dataset.
    It also adds the correct path to the video file in the config file.
    """

    cfg = auxiliaryfunctions.read_config(config)
    project_path = str(Path(config).parents[0])

    cfg["project_path"] = project_path
    if "Reaching" in project_path:
        video_file = os.path.join(project_path, "videos", "reachingvideo1.avi")
    elif "openfield" in project_path:
        video_file = os.path.join(project_path, "videos", "m4s1.mp4")
    else:
        print("This is not an offical demo dataset.")

    if "WILL BE AUTOMATICALLY UPDATED BY DEMO CODE" in cfg["video_sets"].keys():
        cfg["video_sets"][str(video_file)] = cfg["video_sets"].pop(
            "WILL BE AUTOMATICALLY UPDATED BY DEMO CODE"
        )

    auxiliaryfunctions.write_config(config, cfg)
