import os
import numpy as np
import cv2
import tensorflow as tf

from .estimation_utils import get_transformations


# data condition tags
WLFW_CLASS_MAP = [
    'pose',
    'expression',
    'illumination',
    'make-up',
    'occlusion',
    'blur',
]
# face features
FACE_FEATURE_MAP = {
    'image': tf.io.FixedLenFeature([], tf.string),
    'keypoints': tf.io.FixedLenFeature([196], tf.float32),
    'attribute': tf.io.FixedLenFeature([6], tf.int64),
}


def scale_bbox(bbox, scale, min_max=None):
    """Scale bounding box size by ratio

    Args:
        bbox (list): bounding box in format [x_min, y_min, x_max, y_max]
        scale (float): scale ratio
        min_max (tuple, optional): tuple of bounding box boundaries to clip with, if None - unused. Defaults to None.

    Returns:
        list: scaled bounding box in format [x_min, y_min, x_max, y_max]
    """
    # calculate bbox width, height and its centroid
    wh = bbox[2] - bbox[0], bbox[3] - bbox[1]
    xy = bbox[0] + wh[0] / 2, bbox[1] + wh[1] / 2
    # scale width and height
    wh = scale * wh[0], scale * wh[1]
    # convert back to [x_min, y_min, x_max, y_max] format
    bbox = [xy[0] - wh[0] / 2, xy[1] - wh[1] / 2,
            xy[0] + wh[0] / 2, xy[1] + wh[1] / 2]
    if min_max:
        # clip bbox coordinates
        clip_range = lambda x, min_val, max_val: max(min_val, min(x, max_val))
        bbox = [clip_range(bbox[0], *min_max[0]), clip_range(bbox[1], *min_max[1]),
                clip_range(bbox[2], *min_max[0]), clip_range(bbox[3], *min_max[1])]
    return bbox


def read_annotation(path):
    """Load annotation file and construct an annotation dictionary

    Args:
        path (str): path to annotation file

    Returns:
        dict: constructed annotation dictionary
    """
    # load annotation file
    with open(path, 'rt') as file:
        lines = file.readlines()
    data = {}
    # for every entry in annotation
    for i, line in enumerate(lines):
        split_line = line.split(' ')
        img_path = os.path.join(*split_line[-1].rstrip('\n').split('/'))
        # get face keypoints, reshape to (98, 2)
        keypoints = [float(val) for val in split_line[:196]]
        keypoints = np.array(keypoints, dtype=np.float32).reshape(-1, 2)
        # bbox [x_min, y_min, x_max, y_max] coordinates
        bbox = [float(val) for val in split_line[196:200]]
        # scene condition tags
        attributes = [float(val) for val in split_line[200:-1]]
        data[str(i)] = {'img_path': img_path, 'keypoints': keypoints, 'bbox': bbox, 'attr': attributes}
    return data


def data_from_annotation(annot, images_root_path, scale=None):
    """Get data from annotation entry in annotation dictionary

    Args:
        annot (dict): annotation entry
        images_root_path (str): root directory of dataset images
        scale (float, optional): scale ratio for bounding box extraction from images. None equals to 1.0 - the original bbox size.
                                 Values greater then 1.0 scale bbox' width and height read from annotation. Defaults to None.

    Returns:
        np.ndarray, list, np.ndarray, list: read image, bounding box, face keypoints, condition tags
    """
    path = annot['img_path']
    keypoints = annot['keypoints']
    # read image
    img_path = os.path.join(images_root_path, path)
    img = cv2.imread(img_path)
    bbox = annot['bbox']
    # scale bbox
    if scale:
        bbox = scale_bbox(bbox, scale, [[0, img.shape[1]], [0, img.shape[0]]])
    attributes = annot['attr']
    cls = attributes
    return img, bbox, keypoints, cls


def preprocess_image(img, bbox, out_shape, center_format=False, scale=1.0):
    """Preprocess image for model, cutting out ROI from original image

    Args:
        img (np.ndarray): input image
        bbox (list): face bounding box in [x_min, y_min, width, height] format
        out_shape (tuple): output shape of face cutout
        center_format (bool, optional): flag, when True - the bounding box is in [x_center, y_center, width, height]. Defaults to False.
        scale (float, optional): bounding box scale ratio. Defaults to 1.0.

    Returns:
        np.ndarray, np.ndarray, np.ndarray: output image, transformation from original image space to cutout space, inverse transformation matrix
    """
    # convert to center format if not already
    if center_format:
        c_x, c_y, w, h = bbox
    else:
        w, h = bbox[2] - bbox[0], bbox[3] - bbox[1]
        c_x, c_y = bbox[0] + w / 2, bbox[1] + h / 2
    # set aspect ratio to 1:1
    ar = out_shape[1] / out_shape[0]
    if w > ar * h:
        h = w / ar
    elif w < ar * h:
        w = h * ar
    # scale bbox
    w, h = w * scale, h * scale
    bbox = np.array([c_x, c_y, w, h], dtype=np.float32)
    # get transformations
    trans, inv_trans = get_transformations(bbox, out_shape)
    # transform image
    img_patch = cv2.warpAffine(img, trans, (int(out_shape[1]), int(out_shape[0])), flags=cv2.INTER_LANCZOS4)
    return img_patch, trans, inv_trans


def load_tfrecord(tfrecord_path):
    """Load TFrecord dataset file

    Args:
        tfrecord_path (str): path to tfrecord dataset

    Returns:
        (tf.data.TFRecordDatasetV2): parsable tfrecord dataset
    """
    files = tf.data.Dataset.list_files(tfrecord_path)
    dataset = files.flat_map(tf.data.TFRecordDataset)
    return dataset.map(lambda x: parse_tfrecord(x))


def parse_tfrecord(example, img_size):
    """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)

    Returns:
        (tf.Tensor, dict): parsed example: image and label dictionary {'keypoints': tf.Tensor, 'attribute': tf.Tensor}
    """
    x = tf.io.parse_single_example(example, FACE_FEATURE_MAP)
    # transform image to <0, 1> float and resize to target size
    x_train = tf.cast(tf.image.decode_jpeg(x['image'], channels=3), tf.float32)
    x_train = tf.image.resize(x_train, img_size) / 255
    # keypoints and attribute tags
    keypoints = tf.reshape(x['keypoints'], (-1, 2))
    attribute = tf.cast(x['attribute'], tf.int32)
    return x_train, {'keypoints': keypoints, 'attribute': attribute}


@tf.function
def transform_targets(y_train, img_size, heatmap_dims, sigma, testing=False):
    """Transforms keypoint labels to heatmaps

    Args:
        y_train (dict): label dictionary
        img_size (tuple): input image size
        heatmap_dims (tuple): size of generated heatmaps
        sigma (float): gaussian variance of generated heatmaps
        testing (bool, optional): flag if dataset is generated in testing mode, in testing mode keypoints
                                  and attribute tags are preserved in output dictionary. Defaults to False.

    Returns:
        dict: dictionary of labels
    """
    keypoints = y_train['keypoints']
    # scale keypoints to heatmap space
    scaled_keypoints = tf.reshape(keypoints, (-1, 98, 2)) * heatmap_dims / img_size
    scaled_keypoints = tf.reshape(scaled_keypoints, (-1, 1, 1, 98, 2))
    # generate heatmaps
    x, y = tf.meshgrid(tf.range(0, heatmap_dims[0], dtype=tf.float32),
                       tf.range(0, heatmap_dims[1], dtype=tf.float32))
    x = x[tf.newaxis, :, :, tf.newaxis]
    y = y[tf.newaxis, :, :, tf.newaxis]
    heatmaps = tf.exp(-(tf.square(x-scaled_keypoints[..., 0]) + tf.square(y-scaled_keypoints[..., 1])) / (2*sigma ** 2))
    heatmaps = tf.reshape(heatmaps, (-1, heatmap_dims[0], heatmap_dims[1], 98))
    # scale heatmaps to <0, 255>
    heatmaps = heatmaps * 255
    # if training output only heatmaps, otherwise output heatmaps and preserve keypoints and attribute tags in output dict
    if testing:
        # scale keypoints to <0, 1>
        y = {'keypoints': keypoints / img_size, 'attribute': y_train['attribute'], 'heatmaps': heatmaps}
    else:
        y = heatmaps
    return y


def init_dataset(tfrecord_path, batch_size, img_size, heatmap_dims, sigma, shuffle=False, prefetch=False, testing=False):
    """Load TFRecord dataset file as a generator like parser

    Args:
        tfrecord_path (str): path to TFRecord file
        batch_size (int): dataset batch size
        img_size (tuple): input image size in first two dimensions (K, L)
        heatmap_dims (tuple): size of generated heatmaps
        sigma (float): gaussian variance of generated heatmaps
        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 - keypoints are included in target dictionary. Defaults to False.

    Returns:
        tf.data.Dataset: tensorflow dataset, yielding x tensor and y_true dictionary
    """
    # load dataset
    dataset = tf.data.TFRecordDataset(tfrecord_path).map(lambda x: parse_tfrecord(x, img_size))
    # shuffle if specified
    if shuffle:
        dataset = dataset.shuffle(buffer_size=2048)
    # batch dataset
    dataset = dataset.batch(batch_size)
    # generate heatmaps
    dataset = dataset.map(lambda x, y: (x, transform_targets(y, img_size, heatmap_dims, sigma, testing=testing)), num_parallel_calls=tf.data.AUTOTUNE)
    # prefetch if specified
    if prefetch:
        dataset = dataset.prefetch(buffer_size=tf.data.AUTOTUNE)
    return dataset
