"""
Combine hand and face detection datasets into one, use trained individual hand a face detectors to add labels to train and validation set.
"""
import sys
import json
import os
import cv2
import numpy as np

from itertools import chain
from scipy.io import loadmat

from modules.yolov5 import YOLOv5
from modules.dataset_utils import (
    check_bboxes,
    get_hands_files,
    read_hands_annotation,
    preprocess_hands,
    preprocess_wider,
)
from modules.config import CFG


def preprocess_prediction(img, img_size):
    """Preprocess image for prediction

    Args:
        img (np.ndarray): input BGR image
        img_size (tuple): target image size in first two dimensions (K, L)

    Returns:
        np.ndarray: transformed image
    """
    # resize to target size and scale to <0, 1> range
    img = cv2.resize(img, img_size)
    img = img[np.newaxis, ...].astype(np.float32) / 255
    return img


def postprocess_prediction(prediction):
    """Postprocess detection outputs

    Args:
        prediction (list): prediction outputs - boxes, scores and classes

    Returns:
        tf.Tensor: detected boxes
    """
    boxes = prediction[0][0, :]
    scores = prediction[1][0, :]
    # filter boxes with zero scores
    mask = scores != 0
    boxes = boxes[mask, :]
    return boxes


def get_days_of_hands():
    """100 Days of Hands dataset generator

    Yields:
        str, np.ndarray, str: image file path, hand bounding boxes, sample subset
    """
    # iterate over all dataset samples
    for subset in ['train', 'val', 'test']:
        files = get_hands_files(CFG.HANDS_DATASET_PATH['split'][subset])
        for file in files:
            img_path = os.path.join(CFG.HANDS_DATASET_PATH['images'], file + '.jpg')
            annot_path = os.path.join(CFG.HANDS_DATASET_PATH['annotations'], file + '.xml')
            img = cv2.imread(img_path)
            # preprocess and check bbox validity
            boxes = read_hands_annotation(annot_path)
            img, boxes = preprocess_hands(img, boxes, CFG.HANDS_PARAMS.IMG_SHAPE)
            boxes = check_bboxes(boxes, min_area=CFG.BBOX_MIN_AREA)
            if boxes is None:
                continue
            yield img_path, boxes, subset
    yield None, None, None


def get_wider():
    """WIDER Face dataset generator

    Yields:
        str, np.ndarray, str: image file path, hand bounding boxes, sample subset
    """
    # iterate over all dataset samples
    for subset in ['train', 'val']:
        main_dir_path = CFG.FACES_DATASET_PATH['images'][subset]
        annotations = loadmat(CFG.FACES_DATASET_PATH['annotations'][subset])
        for sub_dir, file_names, bbox_list in zip(os.listdir(main_dir_path), annotations['file_list'], annotations['face_bbx_list']):
            sub_dir_path = os.path.join(main_dir_path, sub_dir)
            for file_name, boxes in zip(file_names[0], bbox_list[0]):
                img_path = os.path.join(sub_dir_path, file_name[0][0] + '.jpg')
                img = cv2.imread(img_path)
                # preprocess and check bbox validity
                img, boxes = preprocess_wider(img, boxes, CFG.FACES_PARAMS.IMG_SHAPE)
                boxes = check_bboxes(boxes, min_area=CFG.BBOX_MIN_AREA)
                if boxes is None:
                    continue
                yield img_path, boxes, 'test' if subset == 'val' else 'train'
    yield None, None, None


def hands_datasets(faces_model):
    """Loads hand detection dataset samples and labels them with face labels

    Args:
        faces_model (tf.keras.Model): face detection model

    Yields:
        str, np.ndarray, np.array, str: image file path, object bounding boxes, object classes, sample subset
    """
    for img_path, boxes, subset in get_days_of_hands():
        if img_path is None or boxes is None or subset is None:
            continue
        # native hand objects
        classes = np.zeros(boxes.shape[0], dtype=np.int32)
        # test subsets don't contain automatic labels
        if subset != 'test':
            # read image, preprocess it and make prediction
            img = cv2.imread(img_path)
            img = preprocess_prediction(img, CFG.HANDS_PARAMS.IMG_SHAPE)
            prediction = faces_model.predict(img)
            # postprocess prediction
            boxes_pred = postprocess_prediction(prediction)
            # face objects
            classes_pred = np.ones(boxes_pred.shape[0], dtype=np.int32)
            # combine objects
            boxes = np.concatenate((boxes, boxes_pred), axis=0)
            classes = np.concatenate((classes, classes_pred), axis=0)
        yield img_path, boxes, classes, subset
    yield None, None, None, None


def faces_datasets(hands_model):
    """Loads face detection dataset samples and labels them with hand labels

    Args:
        hands_model (tf.keras.Model): hand detection model

    Yields:
        str, np.ndarray, np.array, str: image file path, object bounding boxes, object classes, sample subset
    """
    for img_path, boxes, subset in get_wider():
        if img_path is None or boxes is None or subset is None:
            continue
        # native face objects
        classes = np.ones(boxes.shape[0], dtype=np.int32)
        # test subsets don't contain automatic labels
        if subset != 'test':
            # read image, preprocess it and make prediction
            img = cv2.imread(img_path)
            img = preprocess_prediction(img, CFG.FACES_PARAMS.IMG_SHAPE)
            prediction = hands_model.predict(img)
            # postprocess prediction
            boxes_pred = postprocess_prediction(prediction)
            # face objects
            classes_pred = np.zeros(boxes_pred.shape[0], dtype=np.int32)
            # combine objects
            boxes = np.concatenate((boxes, boxes_pred), axis=0)
            classes = np.concatenate((classes, classes_pred), axis=0)
        yield img_path, boxes, classes, subset
    yield None, None, None, None


def main(faces_model_cfg, hands_model_cfg):
    """Main script function

    Args:
        faces_model_cfg (str): face detection model configuration, see object detectors for valid configurations
        hands_model_cfg (str): hand detection model configuration, see object detectors for valid configurations
    """
    # load face model and load its weights
    faces_model = YOLOv5(
        (CFG.FACES_PARAMS.IMG_SHAPE[0], CFG.FACES_PARAMS.IMG_SHAPE[1], 3),
        1,
        CFG.FACES_PARAMS.ANCHORS,
        model=faces_model_cfg,
        training=False
    )
    faces_model.load_weights(CFG.FACES_WEIGHTS_PATH[faces_model_cfg]).expect_partial()
    # load hand model and load its weights
    hands_model = YOLOv5(
        (CFG.HANDS_PARAMS.IMG_SHAPE[0], CFG.HANDS_PARAMS.IMG_SHAPE[1], 3),
        1,
        CFG.HANDS_PARAMS.ANCHORS,
        model=hands_model_cfg,
        training=False
    )
    hands_model.load_weights(CFG.HANDS_WEIGHTS_PATH[hands_model_cfg]).expect_partial()

    # annotation dictionary
    annotation_dict = {
        'train': {},
        'val': {},
        'test': {}
    }
    # subset counters
    idx = {'train': 0, 'val': 0, 'test': 0}
    # chain dataset generators together
    datasets = chain(hands_datasets(faces_model), faces_datasets(hands_model))
    # for each dataset record
    for img_path, boxes, classes, subset in datasets:
        if img_path is None or boxes is None or classes is None or subset is None:
            continue
        i = idx[subset]
        # set combined annotation record
        annotation_dict[subset][f'{i:06d}'] = {
            'img_path': img_path.replace(os.path.join(CFG.BASE_PATH, ''), ''),
            'boxes': boxes.tolist(),
            'classes': classes.tolist(),
        }
        idx[subset] += 1
    # write annotation file
    for key in annotation_dict.keys():
        if os.path.isfile(CFG.COMBINED_DATASET_PATH[key]):
            os.remove(CFG.COMBINED_DATASET_PATH[key])
        with open(CFG.COMBINED_DATASET_PATH[key], 'wt') as file:
            json.dump(annotation_dict[key], file, indent=4)


if __name__ == '__main__':
    main('yolov5s', 'yolov5s')
