"""
Deep Lab Cut
============

This module provides an interface between data collected during the reaching task and
forepaw coordinates extracted from videos of the task using Deep Lab Cut.

This depends on pandas, pytables and h5py.
"""


import glob
import json
import os
import pickle
import pandas as pd

import reach
from reach.utilities import cache


class Session(reach.Session):
    """
    This class adds methods to reach.Session that facilitate loading of Deep Lab Cut
    data and aligning this data with data collected during the reach task. This data can
    then be used for analysis or plotting.

    It is assumed that the JSON file storing data for a mouse contains a "video" field
    for each session to be analysed that provides the file name of the video from which
    Deep Lab Cut coordinates were extracted. This is used to identify the DLC data
    files, as the video's file name prefixes DLC's outputs.
    """

    def __init__(self, data, dlc_dir):
        self.data = data
        self.dlc_dir = dlc_dir
        self.coordinates = self.load_coordinates()

    @classmethod
    def init_all_from_file(cls, data_file, dlc_dir):
        """
        Generate a :class:`list` of :class:`Session` objects from data stored in a
        Training JSON.

        Parameters
        ----------
        data_file : :class:`str`
            Full path to file containing existing training data.

        dlc_dir : :class:`str`
            Directory containing data output by Deep Lab Cut for this mouse.

        """
        with open(data_file, "r") as fd:
            previous_data = json.load(fd)

        training_data = [cls(data=data, dlc_dir=dlc_dir) for data in previous_data]

        return training_data

    def load_coordinates(self):
        """
        Load coordinate data generated by DLC as a pandas DataFrame. This can be in H5
        or CSV format.
        """
        basename = os.path.splitext(self.data["video"])[0]
        h5_path = glob.glob(os.path.join(self.dlc_dir, basename + "*.h5"))[0]
        if os.path.isfile(h5_path):
            df = pd.read_hdf(h5_path, "df_with_missing")
            return df[df.keys()[0][0]]

        csv_path = glob.glob(os.path.join(self.dlc_dir, basename + "*.csv"))[0]
        if os.path.isfile(csv_path):
            df = pd.read_csv(csv_path, "df_with_missing")
            return df[df.keys()[0][0]]

        return None

    @cache
    def metadata(self):
        """
        Load DLC metadata from pickle file.
        """
        basename = os.path.splitext(self.data["video"])[0]
        pickle_path = glob.glob(os.path.join(self.dlc_dir, basename + "*.pickle"))[0]
        if os.path.isfile(pickle_path):
            with open(pickle_path, "br") as fd:
                return pickle.load(fd)

        return None


class Mouse(reach.Mouse):
    """
    This class is an extension of reach.Mouse that loads and manages reach.dlc.Session
    objects instead of reach.Session objects.
    """

    @classmethod
    def init_from_file(cls, data_dir, mouse_id, dlc_dir):
        """
        Initialise Mouse object using pre-existing training data stored within
        a training JSON.

        Parameters
        ----------

        data_dir : :class:`str`
            Directory containing training data.

        mouse_id : :class:`str`
            Mouse ID to pass to :class:`Mouse` instance. Will be used to find
            JSON if json_path is a folder.

        dlc_dir : :class:`str`
            Directory containing data output by Deep Lab Cut for this mouse.

        """
        data_file = os.path.join(data_dir, f"{mouse_id}.json")
        if not os.path.isfile(data_file):
            raise SystemError(f"Could not find data file {data_file}.")

        training_data = Session.init_all_from_file(data_file, dlc_dir)
        return cls(mouse_id=mouse_id, training_data=training_data)


class Cohort(reach.Cohort):  # pylint: disable=R0901
    """
    This class is an extension of reach.Cohort that manages reach.dlc.Mouse objects
    instead of reach.Mouse objects.
    """

    @classmethod
    def init_from_files(cls, data_dir, mouse_ids, dlc_dir):
        """
        Initialise the cohort of mice using training files stored within the same
        folder. This class extends reach.Cohort to include Deep Lab Cut data, and so
        needs to know the folder containing this data.

        Parameters
        ----------
        data_dir : :class:`str`
            Directory containing the training data files.

        mouse_ids : :class:`list` of :class:`str`\s
            IDs for the mice to be handled within the cohort.

        dlc_dir : :class:`str`
            Directory containing Deep Lab Cut data for all mice. This folder should
            contain a folder for each mouse containing its mouse ID, and these folders
            should each contain the data output by Deep Lab Cut in files that are
            referenced in the mouse's training data.

        """
        mice = []

        for mouse_id in mouse_ids:
            mouse_dlc_dirs = glob.glob(os.path.join(dlc_dir, f"*{mouse_id}*"))

            if len(mouse_dlc_dirs) > 1:
                print(f"Found more than one folder containing '{mouse_id}'")
                raise Exception("Unsure where to get deeplabcut data. Please merge.")

            mouse = Mouse.init_from_file(
                data_dir=data_dir, mouse_id=mouse_id, dlc_dir=mouse_dlc_dirs[0]
            )
            mice.append(mouse)

        return cls(mice, mouse_ids)
