from __future__ import division
from __future__ import absolute_import
from __future__ import print_function
import hashlib
import simplejson
import random
import shutil
# import urllib
import os
from pprint import pprint

import math
import numpy as np
import sys

import psutil
import requests
from PIL import Image

from aetros.utils import get_option
from .keras_model_utils import ensure_dir
from .backend import invalid_json_values

from threading import Thread, Lock
from six.moves.queue import Queue, Empty
import six
from six.moves import range


def download_image(url, path):
    if os.path.exists(path):
        return True

    headers = {
        'User-Agent': 'Mozilla/5.0 (Macintosh; Intel Mac OS X 10_12_4) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/60.0.3112.101 Safari/537.36'
    }
    try:

        r = requests.get(url, stream=True, timeout=9, headers=headers)
        if r.status_code == 200:
            with open(path, 'wb') as f:
                r.raw.decode_content = True
                shutil.copyfileobj(r.raw, f)
                return True
        else:
            print(("Could not download image %s, response %d" % (url, r.status_code)))
    except Exception as e:
        if hasattr(e, 'message'):
            print(("Could not download image %s due to %s" % (url, e.message)))
        else:
            print(("Could not download image %s due to %s" % (url, repr(e))))

    return False


class ImageDownloaderWorker(Thread):

    def __init__(self, q, progress, dataset, max, images, controller):
        Thread.__init__(self)
        self.q = q
        self.progress = progress
        self.dataset = dataset
        self.max = max
        self.images = images
        self.controller = controller

    def run(self):
        try:
            while self.controller['running']:
                try:
                    self.handle(self.q.get(True, 4))
                    self.q.task_done()
                except Empty:
                    pass
        except KeyboardInterrupt:
            self.controller['running'] = False
            return

    def handle(self, message):
        image = message[0]

        if 'id' not in image:
            return

        local_image_path = message[2]

        try:
            ensure_dir(os.path.dirname(local_image_path))
        except Exception:
            pass

        if not os.path.isfile(local_image_path):
            if download_image(image['src'], local_image_path):
                try:
                    with open(local_image_path, "rb") as fp:
                        img = Image.open(fp)
                        img.load()
                    resize = bool(get_option(self.dataset['config'], 'resize', True))
                    if resize:
                        os.remove(local_image_path)
                        size = (int(get_option(self.dataset['config'], 'resizeWidth', 512)), int(
                            get_option(self.dataset['config'], 'resizeHeight', 512)))
                        quality = int(
                            float(get_option(self.dataset['config'], 'resizeCompression', 0.8)) * 100)
                        img.thumbnail(size, Image.ANTIALIAS)

                        local_image_path = os.path.splitext(local_image_path)[0] + '.jpg'
                        img.save(local_image_path, 'JPEG', quality=quality, optimize=True)

                except SystemExit:
                    self.controller['running'] = False
                except KeyboardInterrupt:
                    self.controller['running'] = False
                except Exception as e:
                    print(("No valid image found %s" % (local_image_path,)))
                    if os.path.exists(local_image_path):
                        os.remove(local_image_path)

        self.images[image['id']] = local_image_path
        self.progress.advance(1)

# class to read all images in the ram at once


class ImageReadWorker(Thread):

    def __init__(self, q, job_model, input_node, path, images, controller):
        Thread.__init__(self)
        self.q = q
        self.job_model = job_model
        self.input_node = input_node
        self.path = path
        self.images = images
        self.controller = controller

    def run(self):
        while self.controller['running']:
            self.handle(self.q.get())
            self.q.task_done()

    def handle(self, message):
        path, validation, category_dir = message

        try:
            if os.path.isfile(path):
                input = self.job_model.convert_file_to_input_node(path, self.input_node)
            else:
                return

            if input is not None and len(input) > 0:
                self.images.append([input, validation, category_dir])
        except IOError as e:
            print(('Could not open %s due to %s' % (path, e.message)))
            return


class InMemoryDataGenerator():

    def __init__(self, datagen, images, classes_count, batch_size):
        self.index = -1
        self.datagen = datagen
        self.images = images
        random.shuffle(self.images)
        self.lock = Lock()

        self.classes_count = classes_count
        self.batch_size = batch_size

    def __iter__(self):
        return self

    def next(self):
        batch_x = []
        batch_y = []

        for i in range(self.batch_size):

            with self.lock:
                self.index += 1
                # reset iterator if necessary
                if self.index == len(self.images):
                    self.index = 0
                    random.shuffle(self.images)

            image, class_idx = self.images[self.index]
            # we need to copy it, otherwise we'd operate on the same object again and again
            image = np.copy(image)

            if self.datagen is not None:
                image = self.datagen.random_transform(image)
                image = self.datagen.standardize(image)

            batch_x.append(image)

            y = np.zeros((self.classes_count,), dtype='float32')
            y[class_idx] = 1.
            batch_y.append(y)

        return np.array(batch_x), np.array(batch_y)

    def __next__(self):
        return self.next()


def read_images_in_memory(job_model, dataset, node, trainer):
    """
    Reads all images into memory and applies augmentation if enabled
    """
    concurrent = psutil.cpu_count()

    dataset_config = dataset['config']
    controller = {'running': True}
    q = Queue(concurrent)

    result = {
        'X_train': [],
        'Y_train': [],
        'X_test': [],
        'Y_test': []
    }

    images = []
    max = 0

    path = job_model.get_dataset_downloads_dir(dataset)
    if 'path' in dataset['config']:
        path = dataset['config']['path']

    classes_count = 0
    category_map = {}
    classes = []

    trainer.set_status('LOAD IMAGES INTO MEMORY')

    try:
        for i in range(concurrent):
            t = ImageReadWorker(q, job_model, node, path, images, controller)
            t.daemon = True
            t.start()

        for validation_or_training in ['validation', 'training']:
            if os.path.isdir(os.path.normpath(path + '/' + validation_or_training)):
                for category_name in os.listdir(os.path.normpath(path + '/' + validation_or_training)):
                    if os.path.isdir(os.path.normpath(path + '/' + validation_or_training + '/' + category_name)):

                        if category_name not in category_map:
                            category_map[category_name] = classes_count
                            if 'classes' in dataset_config and 'category_' in category_name:
                                category_idx = int(category_name.replace('category_', ''))
                                category_map[category_name] = category_idx
                                target_category = dataset_config['classes'][category_idx]
                                classes.append(target_category['title'] or 'Class %s' % (category_idx, ))
                            else:
                                classes.append(category_name)

                            classes_count += 1

                        for id in os.listdir(os.path.normpath(path + '/' + validation_or_training + '/' + category_name)):
                            file_path = os.path.join(path, validation_or_training, category_name, id)
                            q.put([file_path, validation_or_training == 'validation', category_name])
                            max += 1

        q.join()
        controller['running'] = False

        train_images = []
        test_images = []

        for v in images:
            image, validation, category_dir = v
            if validation is True:
                test_images.append([image, category_map[category_dir]])
            else:
                train_images.append([image, category_map[category_dir]])

        train_datagen = None
        augmentation = bool(get_option(dataset_config, 'augmentation', False))
        if augmentation:
            train_datagen = get_image_data_augmentor_from_dataset(dataset)

        train = InMemoryDataGenerator(train_datagen, train_images, classes_count, job_model.job['config']['batchSize'])

        test = InMemoryDataGenerator(None, test_images, classes_count, job_model.job['config']['batchSize'])

        nb_sample = len(train_images)
        trainer.set_info('Dataset size', {'training': nb_sample, 'validation': len(test_images)})
        trainer.set_generator_training_nb(nb_sample)
        trainer.set_generator_validation_nb(len(test_images))

        trainer.logger.info(("Found %d classes, %d images (%d in training [%saugmented], %d in validation). Read all images into memory from %s" %
               (classes_count, max, len(train_images), 'not ' if augmentation is False else '', len(test_images), path)))

        if classes_count == 0:
            trainer.logger.warning("Could not find any classes. Does the directory contains images?")
            sys.exit(1)

        trainer.output_size = classes_count
        trainer.set_info('classes', classes)
        trainer.classes = classes

        result['X_train'] = train
        result['Y_train'] = train
        result['X_test'] = test
        result['Y_test'] = test

        return result

    except KeyboardInterrupt:
        controller['running'] = False
        sys.exit(1)


def get_image_data_augmentor_from_dataset(dataset):
    from keras.preprocessing.image import ImageDataGenerator
    dataset_config = dataset['config']

    augShearRange = float(get_option(dataset_config, 'augShearRange', 0.1))
    augZoomRange = float(get_option(dataset_config, 'augZoomRange', 0.1))
    augHorizontalFlip = bool(get_option(dataset_config, 'augHorizontalFlip', False))
    augVerticalFlip = bool(get_option(dataset_config, 'augVerticalFlip', False))
    augRotationRange = float(get_option(dataset_config, 'augRotationRange', 0.2))

    return ImageDataGenerator(
        rotation_range=augRotationRange,
        shear_range=augShearRange,
        zoom_range=augZoomRange,
        horizontal_flip=augHorizontalFlip,
        vertical_flip=augVerticalFlip
    )


def read_images_keras_generator(job_model, dataset, node, trainer):
    from keras.preprocessing.image import ImageDataGenerator