import os
import scipy.io
import scipy.misc
import imageio
from skimage.transform import resize
from operator import mul
from functools import reduce
import numpy as np
import tensorflow as tf
from tensorflow.python.framework import ops
ops.reset_default_graph()

#オリジナル画像とスタイル画像の場所を指定
sess = tf.Session()

#画像ファイル
original_image_file = 'images/icu.jpeg'
style_image_file = 'images/starry_night.jpeg'

#モデルのパラメーターとして、matファイルの場所、重み、学習率、イテレーション回数、中間画像を出力する頻度を設定
#重みは、オリジナル画像に対してスタイル画像の非常を高くするのに役だ立つ
vgg_path = 'imagenet-vgg-verydeep-19.mat'

# Default Arguments
original_image_weight = 5.0
style_image_weight = 500.0
regularization_weight = 100
learning_rate = 0.001
generations = 20(1800回生成することを推奨！！)
output_generations = 5

#SciPyを使って２つの画像を読み込み、オリジナル画像の大きさに合わせてスタイル画像の形状を変更する
#画像を読み込む
original_image = scipy.misc.imread(original_image_file)
style_image = scipy.misc.imread(style_image_file)

#スタイル画像の形状をオリジナル画像と同じにする
target_shape = original_image.shape
style_image = scipy.misc.imresize(style_image, target_shape[1] / style_image.shape[1])

#論文に基づき、層の順番を定義する（ここでは論文著者の命名規則を使用する）。
# VGG-19 Layer Setup
# From paper
vgg_layers = ['conv1_1', 'relu1_1',
              'conv1_2', 'relu1_2', 'pool1',
              'conv2_1', 'relu2_1',
              'conv2_2', 'relu2_2', 'pool2',
              'conv3_1', 'relu3_1',
              'conv3_2', 'relu3_2',
              'conv3_3', 'relu3_3',
              'conv3_4', 'relu3_4', 'pool3',
              'conv4_1', 'relu4_1',
              'conv4_2', 'relu4_2',
              'conv4_3', 'relu4_3',
              'conv4_4', 'relu4_4', 'pool4',
              'conv5_1', 'relu5_1',
              'conv5_2', 'relu5_2',
              'conv5_3', 'relu5_3',
              'conv5_4', 'relu5_4']
#matファイルからパラメータを抽出する関数
def extract_net_info(path_to_params):
    vgg_data = scipy.io.loadmat(path_to_params)
    normalization_matrix = vgg_data['normalization'][0][0][0]
    mat_mean = np.mean(normalization_matrix, axis=(0,1))
    network_weights = vgg_data['layers'][0]
    return mat_mean, network_weights
    
# 抽出した重みと層の定義に基づき、TensorFlowでCNNを再現するための関数
# 角層を順番に処理する際、適切な重みとバイアスを使用できるようにした上で、各層にあった関数を割り当てる
# Create the VGG-19 Network
def vgg_network(network_weights, init_image):
    network = {}
    image = init_image

    for i, layer in enumerate(vgg_layers):
        if layer[0] == 'c':
            weights, bias = network_weights[i][0][0][0][0]
            weights = np.transpose(weights, (1, 0, 2, 3))
            bias = bias.reshape(-1)
            conv_layer = tf.nn.conv2d(image, tf.constant(weights), (1, 1, 1, 1), 'SAME')
            image = tf.nn.bias_add(conv_layer, bias)
        elif layer[0] == 'r':
            image = tf.nn.relu(image)
        else:  # pooling
            image = tf.nn.max_pool(image, (1, 2, 2, 1), (1, 2, 2, 1), 'SAME')
        network[layer] = image
    return network
    
# どの層を適応させるか選択（オリジナル画像はrelu4_2のまま、スタイル画像はrelu_X層の出力を組み合わせてみる）
original_layer = 'relu4_2'
style_layers = ['relu1_1', 'relu2_1', 'relu3_1', 'relu4_1', 'relu5_1']

# extract_net()を実行して重みと平均を取得する。
# TensorFlowの画像演算は4つの次元を操作するため、先頭にサイズ１の次元を追加し画像の行列の形状を変更する。
normalization_mean, network_weights = extract_net_info(vgg_path)

shape = (1,) + original_image.shape
style_shape = (1,) + style_image.shape
original_features = {}
style_features = {}

# imageプレースホルダを設定し、このプレースホルダを使ってネットワークを形成する
image = tf.placeholder('float', shape=shape)
vgg_net = vgg_network(network_weights, image)

# オリジナル画像の行列を正規化し、ネットワークを通じて実行する
original_minus_mean = original_image - normalization_mean
original_norm = np.array([original_minus_mean])
original_features[original_layer] = sess.run(vgg_net[original_layer], feed_dict= {image: original_norm})

# 選択したスタイル画像の層ごとに同じ処理を繰り返す
image = tf.placeholder('float', shape=style_shape)
vgg_net = vgg_network(network_weights, image)
style_minus_mean = style_image - normalization_mean
style_norm = np.array([style_minus_mean])

for layer in style_layers:
    layer_output = sess.run(vgg_net[layer], feed_dict={image: style_norm})
    layer_output = np.reshape(layer_output, (-1, layer_output.shape[3]))
    style_gram_matrix = np.matmul(layer_output.T, layer_output) / layer_output.size
    style_features[layer] = style_gram_matrix
    
# オリジナル画像とスタイル画像を組み合わせる処理は、ランダムノイズを適用した上でネットワークで実行する
initial = tf.random_normal(shape) * 0.256
image = tf.Variable(initial)
vgg_net = vgg_network(network_weights, image)

# 最適化関数とトレーニングステップを設定し、モデルのすべての変数を初期化する
# 最適化関数を設定
optimizer = tf.train.AdamOptimizer(learning_rate)
train_step = optimizer.minimize(loss)

# 変数を初期化し、トレーニンングを開始
sess.run(tf.global_variables_initializer())

# トレーニングループを開始
# 定期的的に生成された画像を保存（選択した画像に応じて異なる可能性があるため、満足のいく画像が出力された時にアルゴリズムの実行を中止する）
for i in range(generations):
    sess.run(train_step)
    # Print update and save temporary output
    if (i+1) % output_generations == 0:
        print('Generation {} out of {}, loss: {}'.format(i + 1, generations,sess.run(loss)))
        image_eval = sess.run(image)
        best_image_add_mean = image_eval.reshape(shape[1:]) + normalization_mean
        output_file = 'temp_output_{}.jpg'.format(i)
        scipy.misc.imsave(output_file, best_image_add_mean)