#!/usr/bin/env python3

import json
import os

from PIL import Image
import numpy as np

import tensorflow as tf
from tensorflow.keras import Sequential
from tensorflow.keras import layers

from src.Places import MODELS_FOLDER

OP_SCALE = 'scale'
OP_NOISE = 'noise'
OP_NOISE_SCALE = 'noise_scale'

LEAKY_ALPHA = tf.constant(0.1)


def leaky_relu(x):
    return tf.where(tf.greater(0.0, x), tf.multiply(x, LEAKY_ALPHA), x)


def save_image_to(data, path):
    data = np.minimum(np.maximum(0., data[0]), 1.)
    data = np.uint8(np.round(data * 255.))
    image = Image.fromarray(data)
    image.save(path)


def load_weights(config):
    weights = np.asarray(config["weight"], dtype=np.float32).transpose(2, 3, 1, 0)
    bias = np.asarray(config["bias"], dtype=np.float32)
    return [weights, bias]


def create_conv2D_layer(config, activation=None):
    weights = load_weights(config)
    layer = layers.Conv2D(config["nOutputPlane"],
                          strides=(config["dH"], config["dW"]),
                          kernel_size=(config["kH"], config["kW"]),
                          activation=activation,
                          weights=weights)
    return layer


def create_conv2Dtranspose_layer(config):
    weights = load_weights(config)
    layer = layers.Conv2DTranspose(config["nOutputPlane"],
                                   strides=(config["dH"], config["dW"]),
                                   kernel_size=(config["kH"], config["kW"]),
                                   padding='same',
                                   weights=weights)
    return layer


def pad_image(img, padding):
    h, w = img.size
    size = (h + 2 * padding, w + 2 * padding)
    result = Image.new('RGB', size, (0, 0, 0))
    result.paste(img, (padding, padding))
    return result


class Waifu2x:

    def __init__(self, operation, noise_level=0):
        self._operation = operation
        self._noise_level = noise_level
        self.img = None

    def load_image(self, path):
        self.img = Image.open(path)
        if self.img.mode != 'RGB':
            # All images are either B/W (mode = 'L') or 'RGB'
            self.img = self.img.convert('RGB')

    def _get_model_path(self):
        if self._operation == OP_NOISE:
            model_name = 'vgg_7/art/noise%d_model.json' % self._noise_level
        elif self._operation == OP_SCALE:
            model_name = 'upconv_7/art/scale2.0x_model.json'
        elif self._operation == OP_NOISE_SCALE:
            model_name = 'upconv_7/art/noise%d_scale2.0x_model.json' % self._noise_level
        return os.path.join(MODELS_FOLDER, model_name)

    def _load_layers(self):
        model_path = self._get_model_path()
        decoder = json.JSONDecoder()
        with open(model_path, 'r') as hand:
            data = hand.read().strip()
        return decoder.decode(data)

    def _build_model(self):
        if self._operation == OP_NOISE:
            return self._build_vgg7()
        else:
            return self._build_upconv()

    def _build_vgg7(self):
        layers = self._load_layers()
        model = Sequential()
        for i in range(0, 6):
            model.add(create_conv2D_layer(layers[i], activation=leaky_relu))
        model.add(create_conv2D_layer(layers[6]))
        return model

    def _build_upconv(self):
        layers = self._load_layers()
        model = Sequential()
        for i in range(0, 6):
            model.add(create_conv2D_layer(layers[i], activation=leaky_relu))
        model.add(create_conv2Dtranspose_layer(layers[6]))
        return model

    def _get_input_tensor(self):
        if self._operation == OP_NOISE:
            padding = 7
        else:
            padding = 6
        img = pad_image(self.img, padding)
        data = np.asarray(img, dtype=np.float32) / 255.
        return np.expand_dims(data, axis=0)

    def run(self, input_path, output_path):
        self.load_image(input_path)
        model = self._build_model()
        input_data = self._get_input_tensor()
        result = model.predict(input_data)
        save_image_to(result, output_path)


def scale(input_path, output_path):
    waifu2x = Waifu2x(OP_SCALE)
    waifu2x.run(input_path, output_path)


def denoise(input_path, output_path, noise_level):
    waifu2x = Waifu2x(OP_NOISE, noise_level)
    waifu2x.run(input_path, output_path)


def denoise_scale(input_path, output_path, noise_level):
    waifu2x = Waifu2x(OP_NOISE_SCALE, noise_level)
    waifu2x.run(input_path, output_path)

tensorflow.multiply numpy.maximum tensorflow.constant numpy.expand_dims tensorflow.greater numpy.asarray tensorflow.keras.Sequential numpy.round
[(19, 'tensorflow.constant', 'tf.constant', (['(0.1)'], {}), True, 'import tensorflow as tf\n'), (29, 'PIL.Image.fromarray', 'Image.fromarray', (['data'], {}), False, 'from PIL import Image\n'), (35, 'numpy.asarray', 'np.asarray', (["config['bias']"], {'dtype': 'np.float32'}), True, 'import numpy as np\n'), (62, 'PIL.Image.new', 'Image.new', (['"""RGB"""', 'size', '(0, 0, 0)'], {}), False, 'from PIL import Image\n'), (23, 'tensorflow.greater', 'tf.greater', (['(0.0)', 'x'], {}), True, 'import tensorflow as tf\n'), (23, 'tensorflow.multiply', 'tf.multiply', (['x', 'LEAKY_ALPHA'], {}), True, 'import tensorflow as tf\n'), (27, 'numpy.maximum', 'np.maximum', (['(0.0)', 'data[0]'], {}), True, 'import numpy as np\n'), (28, 'numpy.round', 'np.round', (['(data * 255.0)'], {}), True, 'import numpy as np\n'), (75, 'PIL.Image.open', 'Image.open', (['path'], {}), False, 'from PIL import Image\n'), (87, 'os.path.join', 'os.path.join', (['MODELS_FOLDER', 'model_name'], {}), False, 'import os\n'), (91, 'json.JSONDecoder', 'json.JSONDecoder', ([], {}), False, 'import json\n'), (104, 'tensorflow.keras.Sequential', 'Sequential', ([], {}), False, 'from tensorflow.keras import Sequential\n'), (112, 'tensorflow.keras.Sequential', 'Sequential', ([], {}), False, 'from tensorflow.keras import Sequential\n'), (125, 'numpy.expand_dims', 'np.expand_dims', (['data'], {'axis': '(0)'}), True, 'import numpy as np\n'), (34, 'numpy.asarray', 'np.asarray', (["config['weight']"], {'dtype': 'np.float32'}), True, 'import numpy as np\n'), (124, 'numpy.asarray', 'np.asarray', (['img'], {'dtype': 'np.float32'}), True, 'import numpy as np\n')]