import glob
import os
from typing import List, Tuple

import cv2
import h5py
import numpy as np
import scipy.io as sio
from tqdm import tqdm

from mmhuman3d.core.conventions.keypoints_mapping import convert_kps
from mmhuman3d.data.data_structures.human_data import HumanData
from .base_converter import BaseModeConverter
from .builder import DATA_CONVERTERS


@DATA_CONVERTERS.register_module()
class MpiInf3dhpHybrikOptConverter(BaseModeConverter):
    """MPI-INF-3DHP dataset `Monocular 3D Human Pose Estimation In The Wild
    Using Improved CNN Supervision' 3DC`2017 More details can be found in the
    `paper.

    <https://arxiv.org/pdf/1611.09813.pdf>`__.

    Args:
        modes (list): 'test' or 'train' for accepted modes
        extract_img (bool): Store True to extract images into a separate
        folder. Default: False.
    """
    ACCEPTED_MODES = ['test', 'train']

    def __init__(self, modes: List = [], extract_img: bool = False) -> None:
        super(MpiInf3dhpHybrikOptConverter, self).__init__(modes)
        self.extract_img = extract_img

    def extract_keypoints(
        self, keypoints2d: np.ndarray, keypoints3d: np.ndarray,
        num_keypoints: int
    ) -> Tuple[bool, np.ndarray, np.ndarray, List[float]]:
        """Check keypoints validiy and add confidence and bbox."""

        bbox_xyxy = [
            min(keypoints2d[:, 0]),
            min(keypoints2d[:, 1]),
            max(keypoints2d[:, 0]),
            max(keypoints2d[:, 1])
        ]
        bbox_xyxy = self._bbox_expand(bbox_xyxy, scale_factor=1.2)
        bbox_xywh = self._xyxy2xywh(bbox_xyxy)

        # check that all joints are visible
        h, w = 2048, 2048
        x_in = np.logical_and(keypoints2d[:, 0] < w, keypoints2d[:, 0] >= 0)
        y_in = np.logical_and(keypoints2d[:, 1] < h, keypoints2d[:, 1] >= 0)
        ok_pts = np.logical_and(x_in, y_in)

        valid = True
        if np.sum(ok_pts) < num_keypoints:
            valid = False

        # add confidence column
        keypoints2d = np.hstack([keypoints2d, np.ones((num_keypoints, 1))])
        keypoints3d = np.hstack([keypoints3d, np.ones((num_keypoints, 1))])

        return valid, keypoints2d, keypoints3d, bbox_xywh

    def convert_by_mode(self, dataset_path: str, out_path: str,
                        mode: str) -> dict:
        """
        Args:
            dataset_path (str): Path to directory where raw images and
            annotations are stored.
            out_path (str): Path to directory to save preprocessed npz file
            mode (str): Mode in accepted modes

        Returns:
            dict:
                A dict containing keys image_path, bbox_xywh, keypoints2d,
                keypoints2d_mask, keypoints3d, keypoints3d_mask stored in
                HumanData() format
        """
        # use HumanData to store all data
        human_data = HumanData()

        image_path_, bbox_xywh_, keypoints2d_, keypoints3d_ = [], [], [], []

        # training data
        if mode == 'train':
            user_list = range(1, 9)
            # user_list = range(7, 8)
            seq_list = range(1, 3)
            vid_list = list(range(4)) + list(range(4, 9))
            # vid_list = [4]
            counter = 0

            for user_i in tqdm(user_list, desc='user list'):
                for seq_i in seq_list:
                    seq_path = os.path.join(dataset_path[:-10], 'S' + str(user_i),
                                            'Seq' + str(seq_i))
                    # mat file with annotations
                    # annot_file = os.path.join(seq_path, 'annot.mat')
                    # annot2 = sio.loadmat(annot_file)['annot2']
                    # annot3 = sio.loadmat(annot_file)['annot3']

                    for j, vid_i in tqdm(enumerate(vid_list), desc='vid list'):

                        # image folder
                        imgs_path = os.path.join(seq_path,
                                                 'video_' + str(vid_i))

                        # extract frames from video file
                        if self.extract_img:

                            # if doesn't exist
                            if not os.path.isdir(imgs_path):
                                os.makedirs(imgs_path)

                            # video file
                            vid_file = os.path.join(
                                seq_path, 'imageSequence',
                                'video_' + str(vid_i) + '.avi')
                            vidcap = cv2.VideoCapture(vid_file)

                            # process video
                            frame = 0
                            while 1:
                                # extract all frames
                                success, image = vidcap.read()
                                if not success:
                                    break
                                frame += 1
                                # image name
                                imgname = os.path.join(imgs_path,
                                                       '%06d.jpg' % frame)
                                # save image
                                cv2.imwrite(imgname, image)

                        #
                        # # per frame
                        # pattern = os.path.join(imgs_path, '*.jpg')
                        # img_list = glob.glob(pattern)
                        # for i, img_i in enumerate(sorted(img_list)):
                        #
                        #     # for each image we store the relevant annotations
                        #     img_name = img_i.split('/')[-1]
                        #     image_path = os.path.join('S' + str(user_i),
                        #                               'Seq' + str(seq_i),
                        #                               'video_' + str(vid_i),
                        #                               img_name)
                        #
                        #     # 2D keypoints
                        #     keypoints2d = np.reshape(annot2[vid_i][0][i],
                        #                              (28, 2))
                        #
                        #     # 3D keypoints
                        #     keypoints3d = np.reshape(annot3[vid_i][0][i],
                        #                              (28, 3)) / 1000
                        #     keypoints3d = keypoints3d - keypoints3d[
                        #         4]  # 4 is the root
                        #     valid, keypoints2d, keypoints3d, bbox_xywh = \
                        #         self.extract_keypoints(
                        #             keypoints2d, keypoints3d, 28)
                        #
                        #     if not valid:
                        #         continue
                        #
                        #     # because of the dataset size,
                        #     # we only keep every 10th frame
                        #     counter += 1
                        #     if counter % 10 != 1:
                        #         continue
