"""Read TFRecords and provide data for the video part of AVLetters.

The TFRecords can be generated by using `conver_TFRecord.py`.
We get the `dataset` instance and then load data in batches.
"""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import os
import fnmatch
import tensorflow as tf

from data import dataset_utils

slim = tf.contrib.slim


_FILE_PATTERN = 'lips_%s_*.tfrecord'

_ITEMS_TO_DESCRIPTIONS = {
    'lips': 'The images for the video of lipreading',
    'label': 'A single integer representing the label',
}


def get_split_lips(split_name,
                   tfrecord_dir,
                   file_pattern=None,
                   reader=None,
                   num_frames=12):
    """Gets a dataset tuple for reading lip movement videos.

    Args:
        split_name: A 'train'/'validation' split name.
        tfrecord_dir: The base directory of the dataset sources.
        file_pattern: The file pattern to use when matching the
            dataset sources. It is assumed that the pattern contains
            a '%s' string so that the split name can be inserted.
        reader: The TensorFlow reader type.
        num_frames: The number of frames contained in each video sample
            (this is determined when we create the TFRecord files).

    Returns:
        A `Dataset` namedtuple.

    Raises:
        ValueError: if `split_name` is not a valid train/validation split.
    """
    if split_name not in ['train', 'validation']:
        raise ValueError(
            'The split_name %s is not recognized.' % (split_name)
            + 'Please input either train or validation as the split_name')

    if not file_pattern:
        file_pattern = _FILE_PATTERN

    # Compute the number of samples contained in the dataset.
    num_samples = 0
    tfrecords_to_count = [
        os.path.join(tfrecord_dir, file)
        for file in os.listdir(tfrecord_dir)
        if fnmatch.fnmatch(file, file_pattern % split_name)]
    for tfrecord_file in tfrecords_to_count:
        for record in tf.python_io.tf_record_iterator(tfrecord_file):
            num_samples += 1

    file_pattern = os.path.join(tfrecord_dir, file_pattern % split_name)

    if reader is None:
        reader = tf.TFRecordReader

    # Create the keys_to_features dictionary for the decoder
    keys_to_features = {
        'video/data': tf.FixedLenFeature((60, 80, num_frames), tf.float32),
        'video/label': tf.FixedLenFeature(
          (), tf.int64, default_value=tf.zeros((), dtype=tf.int64)),
    }

    items_to_handlers = {
        'video': slim.tfexample_decoder.Tensor(
            'video/data', shape=(60, 80, num_frames, 1)),
        'label': slim.tfexample_decoder.Tensor('video/label'),
    }

    decoder = slim.tfexample_decoder.TFExampleDecoder(
        keys_to_features, items_to_handlers)

    labels_to_names = None
    num_classes = None

    if dataset_utils.has_labels(tfrecord_dir):
        labels_to_names = dataset_utils.read_label_file(tfrecord_dir)
        num_classes = len(labels_to_names)

    return slim.dataset.Dataset(
        data_sources=file_pattern,
        reader=reader,
        decoder=decoder,
        num_samples=num_samples,
        items_to_descriptions=_ITEMS_TO_DESCRIPTIONS,
        num_classes=num_classes,
        labels_to_names=labels_to_names)


def load_batch_lips(dataset,
                    batch_size=32,
                    common_queue_capacity=800,
                    common_queue_min=400,
                    shuffle=True,
                    is_training=True):
    """Loads a single batch of data.

    Args:
        dataset: The dataset to load.
        batch_size: The number of videos in the batch.
        common_queue_capacity, common_queue_min: Decide the shuffle degree.
        shuffle: Whether to shuffle or not.
        is_training: If `is_training` is `True` apply data augmentation.

    Returns:
        videos: A `Tensor` of size [batch_size, height, width, time_frames, 1].
        labels: A `Tensor` of size [batch_size], whose values range between
            0 and `dataset.num_classes`.
    """
    data_provider = slim.dataset_data_provider.DatasetDataProvider(
        dataset, common_queue_capacity=common_queue_capacity,
        common_queue_min=common_queue_min, shuffle=shuffle)
    video, label = data_provider.get(['video', 'label'])

    if is_training:

        transformed_images = []

        # Random brightness, random contrast and random cropping but of
        # the same operations are applied to all the frames of a video
        delta = tf.random_uniform((), -1, 1)
        contrast_factor = tf.random_uniform((), 0.2, 1.8)
        bbox_begin, bbox_end, _ = tf.image.sample_distorted_bounding_box(
            [60, 80, 1], [[[0, 0, 1, 1]]], area_range=[0.8, 1])

        for i in range(video.get_shape()[2]):
            image = video[:, :, i, :]
            image = tf.image.adjust_brightness(image, delta)
            image = tf.image.adjust_contrast(image, contrast_factor)
            image = tf.slice(image, bbox_begin, bbox_end)
            image.set_shape([None, None, 1])
            image = tf.image.resize_images(image, [60, 80])
            transformed_images.append(tf.expand_dims(image, 2))
        video = tf.concat(transformed_images, 2)

    # Batch it up.
    videos, labels = tf.train.batch(
        [video, label],
        batch_size=batch_size,
        num_threads=1,
        capacity=2*batch_size)

    return videos, labels
