import os.path
from collections import OrderedDict
from glob import glob
from typing import Generator, Union, Tuple

import cv2
import numpy as np
import pandas as pd

HandMask = np.ndarray
HandMaskMaybeWithFileName = Union[HandMask, Tuple[HandMask, str]]

_CSV_FILE_NAME = 'data.csv'
_CSV_COLUMNS = OrderedDict([('frame_number', int),
                            ('hand_id', int),
                            ('x', float),
                            ('y', float),
                            ('direction_x', float),
                            ('direction_y', float),
                            ('distance', float),
                            ('detected', bool),
                            ('frames_without_detection', int)])

_HAND_MASK_FILE_GLOB = 'hand_labels_*.png'


def load_coordinates(data_csv_path: str) -> pd.DataFrame:
    data_csv_path = os.path.abspath(data_csv_path)
    hand_data = pd.read_csv(data_csv_path,
                            sep=',',
                            header=None,
                            names=_CSV_COLUMNS.keys(),
                            dtype=_CSV_COLUMNS)

    return hand_data


def load_hand_masks(data_dir_path: str,
                    return_file_names: bool = False) -> Generator[HandMaskMaybeWithFileName, None, None]:
    data_dir_path = os.path.abspath(data_dir_path)

    mask_files = glob(os.path.join(data_dir_path, _HAND_MASK_FILE_GLOB))
    mask_files = sorted(mask_files)

    for mf in mask_files:
        hand_mask = cv2.imread(mf, cv2.IMREAD_UNCHANGED)

        if return_file_names:
            yield hand_mask, mf
        else:
            yield hand_mask


class RecordingHandTracking:
    def __init__(self, path: str):
        self._path = os.path.abspath(path)

    def hand_coordinates(self) -> pd.DataFrame:
        return load_coordinates(os.path.join(self._path, _CSV_FILE_NAME))

    def hand_mask_sequence(self, return_file_names: bool = False) -> Generator[HandMaskMaybeWithFileName, None, None]:
        return load_hand_masks(self._path, return_file_names)


        for mf in mask_files:
            hand_mask = cv2.imread(mf, cv2.IMREAD_UNCHANGED)

            if return_file_names:
                yield hand_mask, mf
            else:
                yield hand_mask