import os
import csv

import numpy as np
import tensorflow as tf

import common
from drawing import parse_drawing, rotate_strokes, \
    mirror_strokes, normalize_strokes, render_image


def _process_strokes(strokes, image_size, augmentation=False):
    if augmentation:
        # Rotation
        rotation_degrees = np.random.uniform(-20, 20)
        strokes = rotate_strokes(strokes, rotation_degrees)
        # Horizontal flip
        if np.random.uniform() <= 0.3:
            strokes = mirror_strokes(strokes)
    strokes = normalize_strokes(strokes)
    image = render_image(strokes, image_size)
    return image


def _process_drawing(drawing, image_size, augmentation=False):
    strokes = parse_drawing(drawing, normalize=False)
    image = _process_strokes(strokes, image_size, augmentation)
    return image


def _csv_generator(csv_file, image_size, augmentation=False, read_label=False):
    with open(csv_file, newline='') as f:
        reader = csv.DictReader(f)
        for example in reader:
            image = _process_drawing(example['drawing'], image_size, augmentation)
            if read_label:
                label = common.WORD2LABEL[example['word']]
                yield example['key_id'], image, label
            else:
                yield example['key_id'], image


def _bin_generator(bin_file, image_size, augmentation=False):
    with open(bin_file, 'rb') as f:
        for example in common.unpack_examples(f):
            image = _process_strokes(example['strokes'], image_size, augmentation)
            yield example['key_id'], image, example['label']


def _load_train_sample(random_seed, name, image_size, augmentation):
    load_bin = lambda bin_file: \
        tf.data.Dataset.from_generator(
            _bin_generator,
            (tf.string, tf.float32, tf.int32),
            args=(bin_file, image_size, augmentation))

    # Samples are generated by sample_train_simplified.py
    sample_dir = os.path.join(common.DATA_DIR,
        'train_simplified_sample', str(random_seed), name)
    bin_files = ['{}/{}'.format(sample_dir, bin_file)
                 for bin_file in os.listdir(sample_dir)
                 if bin_file.endswith('.bin')]

    dataset = tf.data.Dataset.from_tensor_slices(bin_files)
    if name == 'train':
        # For the training set, the partitions will be shuffled before
        # each epoch and the dataset repeats indefinitely.
        dataset = dataset.apply(
            tf.contrib.data.shuffle_and_repeat(
                buffer_size=len(bin_files), count=None))
    # Interleave each partition's records.
    dataset = dataset.apply(
        tf.contrib.data.parallel_interleave(
            load_bin, cycle_length=common.CPU_COUNT))

    return dataset


def load_train_val_datasets(image_size, batch_size, augmentation, random_seed):

    train_dataset = _load_train_sample(random_seed, 'train', image_size, augmentation) \
        .shuffle(buffer_size=3 * batch_size) \
        .batch(batch_size=batch_size, drop_remainder=True) \
        .prefetch(10)

    val_dataset = _load_train_sample(random_seed, 'val', image_size, False) \
        .batch(batch_size=batch_size, drop_remainder=True) \
        .prefetch(10)

    return train_dataset, val_dataset


def load_test_dataset(image_size, batch_size, augmentation=False):
    # Load the test_simplified dataset.
    csv_file = os.path.join(common.DATA_DIR, 'test_simplified.csv')
    dataset = tf.data.Dataset.from_generator(
            _csv_generator,
            (tf.string, tf.float32),
            args=(csv_file, image_size, augmentation, False)) \
        .batch(batch_size=batch_size) \
        .prefetch(1)

    return dataset
