import os

from typing import Tuple

import numpy as np

import cv2

import tensorflow as tf
from tensorflow.keras import datasets

from skimage import color

from utils import scripts


def mnist(flatten=True) -> Tuple[Tuple[np.ndarray, np.ndarray], Tuple[np.ndarray, np.ndarray]]:

    (x_train, y_train), (x_test, y_test) = datasets.mnist.load_data()

    x_train = x_train.astype('float32')
    x_test = x_test.astype('float32')

    x_train /= 255.0
    x_test /= 255.0

    if flatten:

        x_train = x_train.reshape(-1, 784)
        x_test = x_test.reshape(-1, 784)

    return (x_train, y_train), (x_test, y_test)


def load_17flowers(download, working_dir, sub_dir, size=(224, 224),
                   inter_img=cv2.INTER_AREA, inter_seg=cv2.INTER_AREA):

    def parse_id(path):

        return os.path.splitext(path.split('/')[-1])[0].split('_')[-1]

    if download:

        url = 'https://www.robots.ox.ac.uk/~vgg/data/flowers/17/17flowers.tgz'
        fname = '17flowers.tgz'

        download_dir = tf.keras.utils.get_file(origin=url, fname=fname, extract=True, cache_dir=working_dir,
                                               cache_subdir=sub_dir)

        url = 'https://www.robots.ox.ac.uk/~vgg/data/flowers/17/trimaps.tgz'
        fname = 'trimaps.tgz'

        _ = tf.keras.utils.get_file(origin=url, fname=fname, extract=True, cache_dir=working_dir, cache_subdir=sub_dir)

        download_dir = os.path.split(download_dir)[0]

    else:

        download_dir = os.path.join(working_dir, sub_dir)

    print('dir : {}'.format(download_dir))

    seg_dir = os.path.join(download_dir, 'trimaps/')

    img_path = tf.io.gfile.glob(os.path.join(download_dir, 'jpg/*'))
    tmp_path = tf.io.gfile.glob(seg_dir + '*')

    valid_idx = scripts.parse_joint(img_path, tmp_path, parse_id)

    img_path = np.array(img_path)[valid_idx]

    def get_seg_path(i):

        return os.path.join(seg_dir, f'image_{parse_id(img_path[i])}.png')

    def read_example(i):

        seg_path = get_seg_path(i)

        img = cv2.imread(img_path[i], cv2.IMREAD_UNCHANGED)
        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)

        seg = cv2.imread(seg_path, cv2.IMREAD_UNCHANGED)
        seg = cv2.cvtColor(seg, cv2.COLOR_BGR2GRAY)

        img = cv2.resize(img, dsize=size, interpolation=inter_img)
        seg = cv2.resize(seg, dsize=size, interpolation=inter_seg)

        mask = (seg != 113).astype('bool')

        return img, mask

    return len(img_path), read_example


def load_102flowers(download, working_dir, sub_dir, size=(224, 224),
                    inter_img=cv2.INTER_AREA, inter_seg=cv2.INTER_AREA):

    def parse_id(path):

        return os.path.splitext(path.split('/')[-1])[0].split('_')[-1]

    if download:

        url = 'https://www.robots.ox.ac.uk/~vgg/data/flowers/102/102flowers.tgz'
        fname = '102flowers.tgz'

        download_dir = tf.keras.utils.get_file(origin=url, fname=fname, extract=True, cache_dir=working_dir,
                                               cache_subdir=sub_dir)

        url = 'https://www.robots.ox.ac.uk/~vgg/data/flowers/102/102segmentations.tgz'
        fname = '102segmentations.tgz'

        _ = tf.keras.utils.get_file(origin=url, fname=fname, extract=True, cache_dir=working_dir, cache_subdir=sub_dir)

        download_dir = os.path.split(download_dir)[0]

    else:

        download_dir = os.path.join(working_dir, sub_dir)

    print('dir : {}'.format(download_dir))

    seg_dir = os.path.join(download_dir, 'segmim/')

    img_path = tf.io.gfile.glob(os.path.join(download_dir, 'jpg/*'))
    tmp_path = tf.io.gfile.glob(seg_dir + '*')

    valid_idx = scripts.parse_joint(img_path, tmp_path, parse_id)

    img_path = np.array(img_path)[valid_idx]

    def get_seg_path(i):

        return os.path.join(seg_dir, f'segmim_{parse_id(img_path[i])}.jpg')

    def read_example(i):

        seg_path = get_seg_path(i)

        img = cv2.imread(img_path[i], cv2.IMREAD_UNCHANGED)
        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)

        seg = cv2.imread(seg_path, cv2.IMREAD_UNCHANGED)
        seg = cv2.cvtColor(seg, cv2.COLOR_BGR2RGB)

        img = cv2.resize(img, dsize=size, interpolation=inter_img)
        seg = cv2.resize(seg, dsize=size, interpolation=inter_seg)

        return img, seg

    return len(img_path), read_example


def prepare_and_load_17flowers(download, working_dir, sub_dir, shape2d=(224, 224), test_size=0.2,
                               train_batch_size=64, test_batch_size=32, random_state=None):

    m, read_example = load_17flowers(download=download, working_dir=working_dir, sub_dir=sub_dir, size=shape2d)

    print('No. examples : {}'.format(m))

    train_idx, test_idx = DatasetGenerator.train_test_split(num_examples=m, test_size=test_size,
                                                            random_state=random_state)

    print('Train-set size: {}'.format(len(train_idx)))
    print('Test-set size: {}'.format(len(test_idx)))

    def reader(i):

        image, mask = read_example(i)

        image = image.astype('float32')
        image /= 255.0

        return image, mask[..., None]

    data_gen = DatasetGenerator(reader=reader)
    data_gen.set_shape(shape=(*shape2d, 3))
    data_gen.set_shape(shape=(*shape2d, 1))
    data_gen.set_dtype(tf.float32)
    data_gen.set_dtype(tf.float32)

    train_data = data_gen.generate(train_idx).batch(train_batch_size).prefetch(tf.data.AUTOTUNE)
    test_data = data_gen.generate(test_idx).batch(test_batch_size).prefetch(tf.data.AUTOTUNE)

    return train_data, test_data


def prepare_and_load_102flowers(download, working_dir, sub_dir, shape2d=(224, 224), test_size=0.2,
                                train_batch_size=64, test_batch_size=32, random_state=None):

    m, read_example = load_102flowers(download=download, working_dir=working_dir, sub_dir=sub_dir, size=shape2d)

    print('No. examples : {}'.format(m))

    train_idx, test_idx = DatasetGenerator.train_test_split(num_examples=m, test_size=test_size,
                                                            random_state=random_state)

    print('Train-set size: {}'.format(len(train_idx)))
    print('Test-set size: {}'.format(len(test_idx)))

    def reader(i):

        image, seg = read_example(i)

        image = image.astype('float32')
        image /= 255.0

        mask = (color.rgb2gray(seg) > 0.14).astype('uint8')

        return image, mask[..., None]

    data_gen = DatasetGenerator(reader=reader)
    data_gen.set_shape(shape=(*shape2d, 3))
    data_gen.set_shape(shape=(*shape2d, 1))
    data_gen.set_dtype(tf.float32)
    data_gen.set_dtype(tf.float32)

    train_data = data_gen.generate(train_idx).batch(train_batch_size).prefetch(tf.data.AUTOTUNE)
    test_data = data_gen.generate(test_idx).batch(test_batch_size).prefetch(tf.data.AUTOTUNE)

    return train_data, test_data


class DatasetGenerator:

    def __init__(self, reader):

        self.reader = reader

        self._output_shapes = []
        self._output_types = []

    def set_shape(self, shape):

        self._output_shapes.append(tf.TensorShape(shape))

    def set_dtype(self, dtype):

        self._output_types.append(dtype)

    @property
    def output_shapes(self):

        return self._output_shapes

    @property
    def output_types(self):

        return self._output_types

    def generate(self, split_args):

        def py_generator():

            for args in split_args:

                x, y = self.reader(args)

                yield x, y

        tf_dataset = tf.data.Dataset.from_generator(py_generator, output_types=tuple(self.output_types),
                                                    output_shapes=tuple(self.output_shapes))

        return tf_dataset

    @staticmethod
    def train_test_split(num_examples, test_size=0.2, shuffle=False, random_state=None):

        random = np.random.RandomState(random_state)

        train_size = int(num_examples * (1.0 - test_size))

        indices = np.arange(num_examples)

        if shuffle:

            random.shuffle(indices)

        return indices[:train_size], indices[train_size:]
