"""
Construct TFRecord files for hand, face and combined detection datasets
"""
import os
import cv2
import tensorflow as tf
import json
import numpy as np

from scipy.io import loadmat
from random import shuffle

from modules.dataset_utils import (
    check_bboxes,
    get_example,
    compute_anchor_boxes,
    preprocess_wider,
    preprocess_hands,
    read_hands_annotation,
    get_hands_files
)
from modules.config import CFG


def hands_dataset():
    """Construct hand detection TFRecords
    """
    # keep track of bbox heights and widths for anchor box computation
    data = {'width': [], 'height': []}
    # train, validation and testing subsets
    for subset in ['train', 'val', 'test']:
        # load list of dataset files
        files = get_hands_files(CFG.HANDS_DATASET_PATH['split'][subset])
        with tf.io.TFRecordWriter(CFG.HANDS_TFRECORDS_PATH[subset]) as writer:
            for file in files:
                # load image and annotation
                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)
                boxes = read_hands_annotation(annot_path)
                # preprocess data sample
                img, boxes = preprocess_hands(img, boxes, CFG.HANDS_PARAMS.IMG_SHAPE)
                classes = np.zeros((boxes.shape[0]), dtype=np.int32)
                boxes, classes = check_bboxes(boxes, classes, CFG.BBOX_MIN_AREA)
                if boxes is None or classes is None:
                    continue
                # construct example and write it
                example = get_example(img, boxes, classes)
                writer.write(example.SerializeToString())
                if subset == 'train':
                    data['width'] += boxes[:, 2].tolist()
                    data['height'] += boxes[:, 3].tolist()
    # calculate anchor boxes
    compute_anchor_boxes(data, CFG.HANDS_PARAMS.ANCHORS_PATH, visualize=False)
    compute_anchor_boxes(data, CFG.HANDS_PARAMS.EFFDET_ANCHORS_PATH, visualize=False, n_anchors=15)


def face_dataset():
    """Construct face detection TFRecords
    """
    # keep track of bbox heights and widths for anchor box computation
    data = {'width': [], 'height': []}
    # train and validation subsets, test subsets don't have public annotations
    for subset in ['train', 'val']:
        main_dir_path = CFG.FACES_DATASET_PATH['images'][subset]
        # load annotation file
        annotations = loadmat(CFG.FACES_DATASET_PATH['annotations'][subset])
        with tf.io.TFRecordWriter(CFG.FACES_TFRECORDS_PATH[subset]) as writer:
            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, bboxes in zip(file_names[0], bbox_list[0]):
                    # read image
                    img_path = os.path.join(sub_dir_path, file_name[0][0] + '.jpg')
                    img = cv2.imread(img_path)
                    # preprocess image and annotation
                    img, bboxes = preprocess_wider(img, bboxes, CFG.FACES_PARAMS.IMG_SHAPE)
                    classes = np.zeros((bboxes.shape[0]), dtype=np.int32)
                    bboxes, classes = check_bboxes(bboxes, classes, CFG.BBOX_MIN_AREA)
                    if bboxes is None or classes is None:
                        continue
                    # construct example and write it
                    example = get_example(img, bboxes, classes)
                    writer.write(example.SerializeToString())
                    if subset == 'train':
                        data['width'] += bboxes[:, 2].tolist()
                        data['height'] += bboxes[:, 3].tolist()
    # calculate anchor boxes
    compute_anchor_boxes(data, CFG.FACES_PARAMS.ANCHORS_PATH, visualize=False)
    compute_anchor_boxes(data, CFG.FACES_PARAMS.EFFDET_ANCHORS_PATH, visualize=False, n_anchors=15)


def combined_dataset():
    """Construct combined detection TFRecords
    """
    # keep track of bbox heights and widths for anchor box computation
    data = {'width': [], 'height': []}
    # train, validation and testing subsets
    for subset in ['train', 'val', 'test']:
        # load annotation file
        with open(CFG.COMBINED_DATASET_PATH[subset], 'rt') as file:
            dataset = json.load(file)
        # shuffle training examples
        if subset == 'train':
            keys = list(dataset.keys())
            shuffle(keys)
        else:
            keys = dataset.keys()
        with tf.io.TFRecordWriter(CFG.COMBINED_TFRECORDS_PATH[subset]) as writer:
            for key in keys:
                # load image
                img_path = os.path.join(CFG.BASE_PATH, dataset[key]['img_path'])
                img = cv2.imread(img_path)
                # load annotation
                bboxes = dataset[key]['boxes']
                classes = dataset[key]['classes']
                bboxes, classes = check_bboxes(bboxes, classes, CFG.BBOX_MIN_AREA)
                if bboxes is None or classes is None:
                    continue
                # construct example and write it
                example = get_example(img, bboxes, classes)
                writer.write(example.SerializeToString())
                if subset == 'train':
                    data['width'] += bboxes[:, 2].tolist()
                    data['height'] += bboxes[:, 3].tolist()
    # calculate anchor boxes
    compute_anchor_boxes(data, CFG.COMBINED_PARAMS.YOLO_ANCHORS_PATH, visualize=False)
    compute_anchor_boxes(data, CFG.COMBINED_PARAMS.EFFDET_ANCHORS_PATH, visualize=False, n_anchors=15)


if __name__ == '__main__':
    hands_dataset()
    face_dataset()
    # combined_dataset()
