"""
Visualize used datasets from files and from TFRecord files
"""
import os
import cv2
import json
import numpy as np

import xml.etree.ElementTree as ET
from scipy.io import loadmat
from random import shuffle

from modules.dataset_utils import init_dataset
from modules.config import CFG


def hands_from_files():
    """Visualize hand detection dataset from files
    """
    # subset to viz
    subset = 'train'
    # get dataset files
    with open(CFG.HANDS_DATASET_PATH['split'][subset], 'rt') as file:
        files = file.readlines()
    files = [file.rstrip('\n') for file in files]
    # for every file
    for file in files:
        # read 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)
        root = ET.parse(annot_path).getroot()
        # get objects annotations
        boxes = list()
        objects = root.findall('object')
        for obj in objects:
            # hand objects
            if obj.find('name').text != 'hand':
                continue
            # get bbox points
            points = list()
            for point in obj.find('bndbox'):
                points.append(int(point.text))
            boxes.append(points)
            # visualize annotations
            cv2.rectangle(img, (points[0], points[1]), (points[2], points[3]), (255, 0, 0), 2)
        cv2.imshow('dataset', img)
        if cv2.waitKey(0) == 27:
            break
    cv2.destroyAllWindows()


def face_from_files():
    """Visualize face detection dataset from files
    """
    # subset to viz
    subset = 'train'
    main_dir_path = CFG.FACES_DATASET_PATH['images'][subset]
    # load annotation file
    annotations = loadmat(CFG.FACES_DATASET_PATH['annotations'][subset])
    # for every data sample
    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')
            assert(os.path.isfile(img_path))
            img = cv2.imread(img_path)
            # draw boxes
            for box in bboxes[0]:
                cv2.rectangle(img, (box[0], box[1]), (box[0]+box[2], box[1]+box[3]), (0, 0, 255), 2)
            cv2.imshow('dataset', img)
            key = cv2.waitKey(0)
            if key == 27:
                cv2.destroyAllWindows()
                return
            elif key == 13:
                break
    cv2.destroyAllWindows()


def combined_from_files():
    """Visualize combined hand and face detection dataset from files
    """
    # subset to viz
    subset = 'train'
    # load annotation
    with open(CFG.COMBINED_DATASET_PATH[subset], 'rt') as file:
        data = json.load(file)
    # shuffle samples
    keys = list(data.keys())
    shuffle(keys)
    # for every dataset entry
    for key in keys:
        # read image
        img_path = os.path.join(CFG.BASE_PATH, data[key]['img_path'])
        img = cv2.imread(img_path)
        # draw boxes
        boxes = data[key]['boxes']
        classes = data[key]['classes']
        for box, cls in zip(boxes, classes):
            color = (191, 111, 39) if cls == 0 else (71, 58, 240)
            pt1 = (int((box[0] - box[2] / 2)*img.shape[1]), int((box[1] - box[3] / 2)*img.shape[0]))
            pt2 = (int((box[0] + box[2] / 2)*img.shape[1]), int((box[1] + box[3] / 2)*img.shape[0]))
            cv2.rectangle(img, pt1, pt2, color, 2)
        cv2.imshow('dataset', img)
        if cv2.waitKey(0) == 27:
            break
    cv2.destroyAllWindows()


def from_tfrecords(dataset):
    """Visualize a TFRecord dataset

    Args:
        dataset (tf.data.Dataset): dataset to visualize
    """
    # for every entry in dataset
    for x, y in dataset:
        # norm image to <0, 255>
        img = (x.numpy()[0, ...] * 255).astype(np.uint8)
        boxes = y[0, :, :4].numpy()
        classes = y[0, :, 4].numpy()
        # mask padded boxes and classes
        mask = np.logical_not(np.any(boxes == 0, axis=1))
        boxes = boxes[mask, :]
        classes = classes[mask]
        # draw boxes
        for box, cls in zip(boxes, classes):
            x, y, w, h = box
            pt1 = (int((x-w/2)*img.shape[1]), int((y-h/2)*img.shape[0]))
            pt2 = (int((x+w/2)*img.shape[1]), int((y+h/2)*img.shape[0]))
            color = (191, 111, 39) if cls == 0 else (71, 58, 240)
            cv2.rectangle(img, pt1, pt2, color, 2)
            cv2.putText(img, f'{np.sqrt(w*h):.3f}', pt1, cv2.FONT_HERSHEY_DUPLEX, 0.75, color, 2)
        # visualize
        cv2.imshow('dataset', img)
        key = cv2.waitKey(0)
        if key == 27:
            break
    cv2.destroyAllWindows()


def hands_from_tfrecords():
    """Visualize hand detection dataset from TFRecrod files
    """
    dataset = init_dataset(
        CFG.HANDS_TFRECORDS_PATH['train'],
        1,
        CFG.HANDS_PARAMS.ANCHORS,
        CFG.HANDS_PARAMS.IMG_SHAPE,
        CFG.HANDS_PARAMS.MAX_DETECTIONS,
        'yolov5s',
        augment=True,
        shuffle=False,
        testing=True)
    from_tfrecords(dataset)


def face_from_tfrecords():
    """Visualize face detection dataset from TFRecrod files
    """
    dataset = init_dataset(
        CFG.FACES_TFRECORDS_PATH['train'],
        1,
        CFG.FACES_PARAMS.ANCHORS,
        CFG.FACES_PARAMS.IMG_SHAPE,
        CFG.FACES_PARAMS.MAX_DETECTIONS,
        'yolov5s',
        testing=True)
    from_tfrecords(dataset)


def combined_from_tfrecords():
    """Visualize combined hand and face detection dataset from TFRecrod files
    """
    dataset = init_dataset(
        CFG.COMBINED_TFRECORDS_PATH['test'],
        1,
        CFG.COMBINED_PARAMS.YOLO_ANCHORS,
        CFG.COMBINED_PARAMS.IMG_SHAPE,
        CFG.COMBINED_PARAMS.MAX_DETECTIONS,
        'yolov5s',
        shuffle=True,
        testing=True)
    from_tfrecords(dataset)


if __name__ == '__main__':
    hands_from_files()
    # face_from_files()
    # combined_from_files()
    # hands_from_tfrecords()
    # face_from_tfrecords()
    # combined_from_tfrecords()
