"""Utilities for dataset manipulation used for hand, face and combined datasets
"""
import os
import cv2
import numpy as np
import tensorflow as tf
import xml.etree.ElementTree as ET

from sklearn.cluster import KMeans
import matplotlib.pyplot as plt


# tfrecord features
DATASET_FEATURES = {
    'image': tf.io.FixedLenFeature([], tf.string),
    'bbox/x_center': tf.io.VarLenFeature(tf.float32),
    'bbox/y_center': tf.io.VarLenFeature(tf.float32),
    'bbox/width': tf.io.VarLenFeature(tf.float32),
    'bbox/height': tf.io.VarLenFeature(tf.float32),
    'bbox/class': tf.io.VarLenFeature(tf.int64),
}


def check_bboxes(bboxes, classes=None, min_area=None):
    """Checks bounding box validity. Expects bbox format in [x_center, y_center, width, height] all normed to range <0, 1>.


    Args:
        bboxes (np.ndarray): bounding boxes [N, 4] where N is number of boxes. Can also be list of lists in the same format.
        classes (np.array, optional): classes to filter alongside bboxes. Defaults to None.
        min_area (_type_, optional): minimum valid bbox area. Defaults to None.

    Returns:
        tuple: filtered bboxes and filtered classes if passed
    """
    # convert to numpy object
    bboxes = np.array(bboxes, dtype=np.float32)
    # if empty return None
    if bboxes.size == 0:
        return None
    # init mask with ones -> all valid
    mask = np.ones((bboxes.shape[0]), dtype=np.bool)
    # any record in bbox < 0 or > 1 => bbox is invalid, update mask
    mask *= np.all(bboxes > 0, axis=1)
    mask *= np.all(bboxes < 1, axis=1)
    # if minimum bbox are is specified, check area and update mask
    if min_area is not None:
        mask *= ((bboxes[:, 2] * bboxes[:, 3]) > min_area)
    # apply mask
    bboxes = bboxes[mask, :]
    # if classes objects are passed
    if classes is not None:
        # convert to numpy object and apply mask
        classes = np.array(classes, dtype=np.int32)
        classes = classes[mask]
        if bboxes.size == 0:
            return None, None
        return bboxes, classes
    if bboxes.size == 0:
        return None
    return bboxes


def get_hands_files(path):
    """Load list of image filenames from annotation file

    Args:
        path (str): path to annotation file containing filenames

    Returns:
        list: list of read filenames
    """
    with open(path, 'rt') as file:
        files = file.readlines()
    files = [file.rstrip('\n') for file in files]
    return files


def read_hands_annotation(path):
    """Read hand detection annotation file

    Args:
        path (str): path to annotation file

    Returns:
        list: list of bounding boxes read in annotation
    """
    # annotation is in XML format
    root = ET.parse(path).getroot()
    boxes = list()
    # search for object tag
    objects = root.findall('object')
    for obj in objects:
        # check if object is hand
        if obj.find('name').text != 'hand':
            continue
        # read bbox points
        points = list()
        for point in obj.find('bndbox'):
            points.append(int(point.text))
        boxes.append(points)
    return boxes


def preprocess_hands(img, boxes, img_size):
    """Preprocess hand detection dataset example. Resize image to target size and transform bounding boxes
       to target format from native hand detection format.

    Args:
        img (numpy.ndarray): BGR image in np.uint8 type
        boxes (numpy.ndarray): array of bounding boxes [N, 4] where N is number of boxes. Boxes are expected in native
                               hand detection dataset format of [x_min, y_min, x_max, y_max] in pixels.
        img_size (tuple): target image size in first two dimensions (K, L)

    Returns:
        tuple: resized image and transformed bounding boxes to [x_center, y_center, width, height] format
    """
    bboxes = list()
    # convert boxes to target format
    for box in boxes:
        x_min, y_min, x_max, y_max = box
        # calculate target bbox params and norm to <0, 1> range in image plane
        w, h = x_max - x_min, y_max - y_min
        x, y = (x_min + w / 2) / img.shape[1], (y_min + h / 2) / img.shape[0]
        w, h = w / img.shape[1], h / img.shape[0]
        bboxes.append([x, y, w, h])
    bboxes = np.array(bboxes, dtype=np.float32)
    # resize image
    img = cv2.resize(img, img_size, cv2.INTER_LANCZOS4)
    return img, bboxes


def preprocess_wider(img, bboxes, img_size):
    """Preprocess face detection dataset example. Resize image to target size and transform bounding boxes
       to target format from native face detection format.

    Args:
        img (numpy.ndarray): BGR image in np.uint8 type
        boxes (numpy.ndarray): array of bounding boxes [N, 4] where N is number of boxes. Boxes are expected in native
                               face detection dataset format of [x_min, y_min, width, height] in pixels.
        img_size (tuple): target image size in first two dimensions (K, L)

    Returns:
        tuple: resized image and transformed bounding boxes to [x_center, y_center, width, height] format
    """
    # convert x_min and y_min to centered coordinates
    bboxes = np.array(list(bboxes), dtype=np.float32).reshape(-1, 4)
    bboxes[..., :2] += bboxes[..., 2:] / 2
    bboxes[..., :2] = bboxes[..., :2]
    # norm to <0, 1> range in image plane
    bboxes[..., :2] /= img.shape[1::-1]
    bboxes[..., 2:] /= img.shape[1::-1]
    # resize image
    img = cv2.resize(img, img_size, cv2.INTER_LANCZOS4)
    return img, bboxes


def get_example(img, bboxes, classes):
    """Form tfrecord example from image, boxes and classes

    Args:
        img (np.ndarray): BGR image resized to target image size
        bboxes (np.ndarray): array of N bounding boxes [N, 4] in [x_center, y_center, width, height] format
        classes (np.array): list like iterable of classes in the same order as bboxes

    Returns:
        tf.train.Example: Tfrecord example object
    """
    # convert to RGB and decode to jpg, the conversion has BGR->RGB integrated,
    # hence example is written in BGR format
    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
    img_encoded = cv2.imencode('.jpg', img)[1].tobytes()
    # form lists of object parameters - bbox params and class
    x, y, w, h, c = list(), list(), list(), list(), list()
    for box, cls in zip(bboxes, classes):
        x.append(box[0])
        y.append(box[1])
        w.append(box[2])
        h.append(box[3])
        c.append(int(cls))
    # construct example
    example = tf.train.Example(features=tf.train.Features(feature={
        'image':
            tf.train.Feature(bytes_list=tf.train.BytesList(value=[img_encoded])),
        'bbox/x_center':
            tf.train.Feature(float_list=tf.train.FloatList(value=x)),
        'bbox/y_center':
            tf.train.Feature(float_list=tf.train.FloatList(value=y)),
        'bbox/width':
            tf.train.Feature(float_list=tf.train.FloatList(value=w)),
        'bbox/height':
            tf.train.Feature(float_list=tf.train.FloatList(value=h)),
        'bbox/class':
            tf.train.Feature(int64_list=tf.train.Int64List(value=c)),
    }))
    return example


def compute_anchor_boxes(data, path, n_anchors=9, visualize=False):
    """Computes anchor boxes from input bounding boxes and saves them to file

    Args:
        data (dict): dictionary in format of {'width': list, 'height': list}, where width and height are
                     normed bbox widths and height in <0, 1> range.
        path (str): target path of anchor file
        n_anchors (int, optional): number of anchors to compute. Defaults to 9.
        visualize (bool, optional): flag, wether or not to visualize anchor results. Defaults to False.
    """
    # stack width and height features to matrix [N, 2]
    w = np.array(data['width'], dtype=np.float32).reshape((-1, 1))
    h = np.array(data['height'], dtype=np.float32).reshape((-1, 1))
    x = np.concatenate((w, h), axis=1)
    # construct KMeans object and fit feature matrix
    kmeans = KMeans(n_clusters=n_anchors)
    kmeans.fit(x)
    # KMeans centroids are anchor boxes' widths and heights
    anchors = kmeans.cluster_centers_.tolist()
    # sort anchors by size from heighest to lowest
    anchors.sort(key=lambda x: x[0]*x[1], reverse=True)
    # save anchors to file
    if os.path.isfile(path):
        os.remove(path)
    with open(path, 'w') as file:
        for anchor in anchors:
            file.write(f'{anchor[0]} {anchor[1]}\n')
    # if viz flag is passed
    if visualize:
        plt.scatter(x[:, 0], x[:, 1], c=kmeans.predict(x), s=2, cmap='viridis')
        plt.scatter(kmeans.cluster_centers_[:, 0], kmeans.cluster_centers_[:, 1], c='red', s=50)
        plt.show()


def parse_tfrecord(example, img_shape, max_detections):
    """Parse TFrecord serialized dataset files

    Args:
        example (tf.data.TFRecordDatasetV2): parsable tfrecord dataset
        img_shape (tuple): target image size in first two dimensions (K, L)
        max_detections (int): max detections accross the dataset

    Returns:
        (tf.Tensor, tf.Tensor): parsed example and label tf.Tensors
    """
    # parse tfrecord
    x = tf.io.parse_single_example(example, DATASET_FEATURES)
    # decode image and resize it to square
    x_train = tf.image.decode_jpeg(x['image'], channels=3)
    x_train = tf.image.resize(x_train, img_shape) / 255
    # stack bounding boxes and labels to single matrix
    y_train = tf.stack([tf.sparse.to_dense(x['bbox/x_center']),
                        tf.sparse.to_dense(x['bbox/y_center']),
                        tf.sparse.to_dense(x['bbox/width']),
                        tf.sparse.to_dense(x['bbox/height']),
                        tf.cast(tf.sparse.to_dense(x['bbox/class']), tf.float32)], axis=1)
    # add padding to y_train to unite dimensions
    paddings = [[0, max_detections - tf.shape(y_train)[0]], [0, 0]]
    y_train = tf.pad(y_train, paddings)
    return x_train, y_train


def augmentation(x_train, y_train):
    """Data augmentation - random flip along vertical axis

    Args:
        x_train (tf.Tensor): tensor of batched BGR images [B, K, L, 3]
        y_train (tf.tensor): tensor of batched labels [B, N, 5]

    Returns:
        tf.Tensor, tf.Tensor: augmented examples and labels
    """
    # flip coefficient
    random = tf.random.uniform((), 0, 1, dtype=tf.float32, seed=42)
    if random >= 0.5:
        # flip images
        x_train = tf.image.flip_left_right(x_train)
        # flip x coordinates of labels
        y_train = tf.constant([1., 0., 0., 0., 0.]) + tf.constant([-1., 1., 1., 1., 1.]) * y_train
    return x_train, y_train


@tf.function
def assign_anchors(y_true, grid_size, anchor_idxs):
    """Assigns labels to anchors to target grid tiles

    Args:
        y_true (tf.Tensor): tensor of labels [B, N, 6]
        grid_size (int): size of a grid tile in pixels
        anchor_idxs (tuple): anchor indices used in the current grid

    Returns:
        tf.Tensor: tensor containing labels in target space
    """
    batches = tf.shape(y_true)[0]
    y_true_out = tf.zeros((batches, grid_size, grid_size, 3, 6))

    anchor_idxs = tf.cast(anchor_idxs, tf.int32)
    indexes = tf.TensorArray(tf.int32, 1, dynamic_size=True)
    updates = tf.TensorArray(tf.float32, 1, dynamic_size=True)
    idx = 0
    # iterate over batches
    for i in tf.range(batches):
        # iterate over targets
        for j in tf.range(tf.shape(y_true)[1]):
            # ignore invalid targets
            if tf.equal(y_true[i][j][2], 0) or tf.equal(y_true[i][j][3], 0):
                continue
            # if target anchor is in specified anchors
            anchor_eq = tf.equal(anchor_idxs, tf.cast(y_true[i][j][5], tf.int32))
            if tf.reduce_any(anchor_eq):
                box = y_true[i][j][:4]
                anchor_idx = tf.cast(tf.where(anchor_eq), tf.int32)[0][0]
                # position on grid
                grid_xy = tf.cast(tf.math.floor(box[:2] / tf.cast(1 / grid_size, tf.float32)), tf.int32)
                # save box, label, position on grid and anchor index
                indexes = indexes.write(idx, [i, grid_xy[1], grid_xy[0], anchor_idx])
                # grid[y][x][anchor] = (x, y, w, h, obj, class)
                updates = updates.write(idx, [box[0], box[1], box[2], box[3], 1., y_true[i][j][4]])
                idx += 1
    return tf.tensor_scatter_nd_update(y_true_out, indexes.stack(), updates.stack())


def transform_targets(y_train, anchors, grid_size, anchor_indices):
    """Transform labels to detector target space

    Args:
        y_train (tf.Tensor): Target tensor [B, N, 5]
        anchors (tf.Tensor): Target anchor boxes [M, 3, 2]
        grid_size (int): smallest size of a grid tile
        anchor_indices (list): list of tuples of anchor box indices grouped by corresponding grid size

    Returns:
        tuple: tuple of targets assigned to corresponding grids
    """
    grid_size = tf.cast(grid_size, tf.int32)
    # compute anchors and box areas
    anchors = tf.reshape(tf.cast(anchors, tf.float32), (-1, 2))
    anchor_area = anchors[..., 0] * anchors[..., 1]
    box_wh = tf.tile(tf.expand_dims(y_train[..., 2:4], -2), (1, 1, tf.shape(anchors)[0], 1))
    box_area = box_wh[..., 0] * box_wh[..., 1]
    # compute intersection over union between boxes and anchors
    intersection = tf.minimum(box_wh[..., 0], anchors[..., 0]) * tf.minimum(box_wh[..., 1], anchors[..., 1])
    iou = intersection / (box_area + anchor_area - intersection)
    # assign anchor indices to boxes based on best iou
    anchor_idx = tf.cast(tf.argmax(iou, axis=-1), tf.float32)[..., tf.newaxis]
    # concat anchor indices to label tensor
    y_train = tf.concat([y_train, anchor_idx], axis=-1)
    # iterate over grids, assign all labels to grid tiles and anchors
    y_outputs = []
    for anchor_idxs in anchor_indices:
        y_outputs.append(assign_anchors(y_train, grid_size, anchor_idxs))
        # next grid size
        grid_size *= 2
    return tuple(y_outputs)


def init_dataset(path, batch_size, anchors, img_shape, max_detections, model_cfg, augment=False, shuffle=False, prefetch=False, testing=False):
    """Load TFRecord dataset file as a generator like parser

    Args:
        path (str): path to TFRecord file
        batch_size (int): dataset batch size
        anchors (tf.Tensor): target anchors
        img_shape (tuple): input image size in first two dimensions (K, L)
        max_detections (int): max detections accross the dataset
        model_cfg (str): model configuration - 'yolov5' or 'efficientdet' based architectures
        augment (bool, optional): data augmentation flag. Defaults to False.
        shuffle (bool, optional): data shuffle flag. Defaults to False.
        prefetch (bool, optional): data prefetching flag. Defaults to False.
        testing (bool, optional): testing flag, if True, dataset yields labels as bounding boxes and labels.
                                  If False, targets are transformed to detector output space. Defaults to False.

    Returns:
        tf.data.Dataset: tensorflow dataset, yielding x and y_true tensors
    """
    # load tfrecord file
    dataset = tf.data.TFRecordDataset(path).map(lambda x: parse_tfrecord(x, img_shape, max_detections))
    # shuffle dataset if flag passed
    if shuffle:
        dataset = dataset.shuffle(buffer_size=2048)
    # batch dataset and transform x and y to trainable targets
    dataset = dataset.batch(batch_size)
    # augment dataset if flag passed
    if augment:
        dataset = dataset.map(lambda x, y: augmentation(x, y))
    # if training, transform labels to target space
    if not testing:
        # smallest grid size is based on the model architecture
        grid_size = img_shape[0] // 32 if 'yolov5' in model_cfg else img_shape[0] // 128
        # yolov5 models have 9 anchors, efficientdet 15 anchors
        anchor_indices = ((0, 1, 2), (3, 4, 5), (6, 7, 8)) if 'yolov5' in model_cfg else ((0, 1, 2), (3, 4, 5), (6, 7, 8), (9, 10, 11), (12, 13, 14))
        dataset = dataset.map(lambda x, y: (x, transform_targets(y, anchors, grid_size, anchor_indices)), num_parallel_calls=tf.data.AUTOTUNE)
    # prefetch dataset if flag passed
    if prefetch:
        dataset = dataset.prefetch(tf.data.AUTOTUNE)
    return dataset
