#!/usr/bin/env python3
# -*- coding:utf-8 -*-
#############################################################
# File: network_models_LSTU.py
# Created Date: Tuesday February 25th 2020
# Author: Chen Xuanhong
# Email: chenxuanhongzju@outlook.com
# Last Modified:  Tuesday, 25th February 2020 9:57:06 pm
# Modified By: Chen Xuanhong
# Copyright (c) 2020 Shanghai Jiao Tong University
#############################################################


from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

from functools import partial

import tensorflow as tf
import tensorflow.contrib.slim as slim
import tflib as tl


conv = partial(slim.conv2d, activation_fn=None)
dconv = partial(slim.conv2d_transpose, activation_fn=None)
fc = partial(tl.flatten_fully_connected, activation_fn=None)
relu = tf.nn.relu
lrelu = tf.nn.leaky_relu
sigmoid = tf.nn.sigmoid
tanh = tf.nn.tanh
batch_norm = partial(slim.batch_norm, scale=True, updates_collections=None)
instance_norm = slim.instance_norm

MAX_DIM = 64 * 16


def Genc(x, dim=64, n_layers=5, multi_inputs=1, is_training=True):
    bn = partial(batch_norm, is_training=is_training)
    conv_bn_lrelu = partial(conv, normalizer_fn=bn, activation_fn=lrelu)

    with tf.variable_scope('Genc', reuse=tf.AUTO_REUSE):
        h, w = x.shape[1:3]
        z = x
        zs = []
        for i in range(n_layers):
            d = min(dim * 2**i, MAX_DIM)
            if multi_inputs > i and i > 0:
                z = tf.concat([z, tf.image.resize_bicubic(x, (h//(2**i), w//(2**i)))], 3)
            z = conv_bn_lrelu(z, d, 4, 2)
            zs.append(z)
        return zs

def LSTU(in_data, state, out_channel, is_training=True, kernel_size=3, norm='none', pass_state='lstate'):
    if norm == 'bn':
        norm_fn = partial(batch_norm, is_training=is_training)
    elif norm == 'in':
        norm_fn = instance_norm
    else:
        norm_fn = None
    gate = partial(conv, normalizer_fn=norm_fn, activation_fn=sigmoid)
    info = partial(conv, normalizer_fn=norm_fn, activation_fn=tanh)
    with tf.name_scope('ConvGRUCell'):
        state_      = dconv(state, out_channel, 4, 2)  # upsample and make `channel` identical to `out_channel`
        reset_gate  = gate(tf.concat([in_data, state_], axis=3), 1, kernel_size)
        update_gate = gate(tf.concat([in_data, state_], axis=3), 1, kernel_size)
        new_state   = reset_gate * state_
        new_info    = info(tf.concat([in_data, new_state], axis=3), out_channel, kernel_size)
        output      = (1-update_gate)*state_ + update_gate*new_info
        if pass_state == 'gru':
            return output, output
        elif pass_state == 'direct':
            return output, state_
        else: # 'stu'
            return output, new_state
        # state_hat = dconv(old_state, outdim, 4, 2)
        # tmp_concat= _concat(x, state_hat, None)
        # channelpool1=tf.concat([tf.reduce_max(tmp_concat,3,True), tf.reduce_mean(tmp_concat,3,True)], axis=3)
        # r_channel=conv(channelpool1,1,7,1,normalizer_fn=None,activation_fn=sigmoid)

        # new_state   = r_channel * state_hat
        # tmp_concat=  _concat(x, new_state, None)
        # hidden_info = conv(tmp_concat,outdim,3,1,normalizer_fn=None,activation_fn=tanh)

        # tmp_concat= _concat(x, state_hat, None)
        # channelpool2=tf.concat([tf.reduce_max(tmp_concat,3,True), tf.reduce_mean(tmp_concat,3,True)], axis=3)
        # z=conv(channelpool2,1,7,1,normalizer_fn=None,activation_fn=sigmoid)
        # output =z *hidden_info +(1-z)*state_hat
    # return output,new_state

def Gstu(zs, _a, dim=64, n_layers=1, inject_layers=0, is_training=True, kernel_size=3, norm='none', pass_state='stu'):
    def _concat(z, z_, _a):
        feats = [z]
        if z_ is not None:
            feats.append(z_)
        if _a is not None:
            _a = tf.reshape(_a, [-1, 1, 1, tl.shape(_a)[-1]])
            _a = tf.tile(_a, [1, tl.shape(z)[1], tl.shape(z)[2], 1])
            feats.append(_a)
        return tf.concat(feats, axis=3)
    
    with tf.variable_scope('Gstu', reuse=tf.AUTO_REUSE):
        zs_ = [zs[-1]]
        state = _concat(zs[-1], None, _a)
        for i in range(n_layers): # n_layers <= 4
            d = min(dim * 2**(n_layers - 1 - i), MAX_DIM)

            output = LSTU(zs[n_layers - 1 - i],state,d,is_training=is_training,
                                 kernel_size=kernel_size, norm=norm, pass_state=pass_state)
            zs_.insert(0, output[0])
            if inject_layers > i:
                state = _concat(output[1], None, _a)
            else:
                state = output[1]
        return zs_

def Gdec(zs, _a, dim=64, n_layers=5, shortcut_layers=1, inject_layers=0, is_training=True, one_more_conv=0):
    bn = partial(batch_norm, is_training=is_training)
    dconv_bn_relu = partial(dconv, normalizer_fn=bn, activation_fn=relu)

    shortcut_layers = min(shortcut_layers, n_layers - 1)
    inject_layers = min(inject_layers, n_layers - 1)

    def _concat(z, z_, _a):
        feats = [z]
        if z_ is not None:
            feats.append(z_)
        if _a is not None:
            _a = tf.reshape(_a, [-1, 1, 1, tl.shape(_a)[-1]])
            _a = tf.tile(_a, [1, tl.shape(z)[1], tl.shape(z)[2], 1])
            feats.append(_a)
        return tf.concat(feats, axis=3)

    with tf.variable_scope('Gdec', reuse=tf.AUTO_REUSE):
        z = _concat(zs[-1], None, _a)
        for i in range(n_layers):
            if i < n_layers - 1:
                d = min(dim * 2**(n_layers - 1 - i), MAX_DIM)
                z = dconv_bn_relu(z, d, 4, 2)
                if shortcut_layers > i:
                    z = _concat(z, zs[n_layers - 2 - i], None)
                if inject_layers > i:
                    z = _concat(z, None, _a)
            else:
                if one_more_conv: # add one more conv after the decoder
                    z = dconv_bn_relu(z, dim//4, 4, 2)
                    x = tf.nn.tanh(dconv(z, 3, one_more_conv))
                else:
                    x = z = tf.nn.tanh(dconv(z, 3, 4, 2))
        return x


def D(x, n_att, dim=64, fc_dim=MAX_DIM, n_layers=5):
    conv_in_lrelu = partial(conv, normalizer_fn=instance_norm, activation_fn=lrelu)

    with tf.variable_scope('D', reuse=tf.AUTO_REUSE):
        y = x
        for i in range(n_layers):
            d = min(dim * 2**i, MAX_DIM)
            y = conv_in_lrelu(y, d, 4, 2)

        logit_gan = lrelu(fc(y, fc_dim))
        logit_gan = fc(logit_gan, 1)

        logit_att = lrelu(fc(y, fc_dim))
        logit_att = fc(logit_att, n_att)

        return logit_gan, logit_att


def gradient_penalty(f, real, fake=None):
    def _interpolate(a, b=None):
        with tf.name_scope('interpolate'):
            if b is None:   # interpolation in DRAGAN
                beta = tf.random_uniform(shape=tf.shape(a), minval=0., maxval=1.)
                _, variance = tf.nn.moments(a, range(a.shape.ndims))
                b = a + 0.5 * tf.sqrt(variance) * beta
            shape = [tf.shape(a)[0]] + [1] * (a.shape.ndims - 1)
            alpha = tf.random_uniform(shape=shape, minval=0., maxval=1.)
            inter = a + alpha * (b - a)
            inter.set_shape(a.get_shape().as_list())
            return inter

    with tf.name_scope('gradient_penalty'):
        x = _interpolate(real, fake)
        pred = f(x)
        if isinstance(pred, tuple):
            pred = pred[0]
        grad = tf.gradients(pred, x)[0]
        norm = tf.norm(slim.flatten(grad), axis=1)
        gp = tf.reduce_mean((norm - 1.)**2)
        return gp

[(25, 'functools.partial', 'partial', False, 'from functools import partial\n'), (26, 'functools.partial', 'partial', False, 'from functools import partial\n'), (27, 'functools.partial', 'partial', False, 'from functools import partial\n'), (32, 'functools.partial', 'partial', False, 'from functools import partial\n'), (39, 'functools.partial', 'partial', False, 'from functools import partial\n'), (40, 'functools.partial', 'partial', False, 'from functools import partial\n'), (61, 'functools.partial', 'partial', False, 'from functools import partial\n'), (62, 'functools.partial', 'partial', False, 'from functools import partial\n'), (118, 'functools.partial', 'partial', False, 'from functools import partial\n'), (119, 'functools.partial', 'partial', False, 'from functools import partial\n'), (154, 'functools.partial', 'partial', False, 'from functools import partial\n'), (42, 'tensorflow.variable_scope', 'tf.variable_scope', True, 'import tensorflow as tf\n'), (56, 'functools.partial', 'partial', False, 'from functools import partial\n'), (63, 'tensorflow.name_scope', 'tf.name_scope', True, 'import tensorflow as tf\n'), (100, 'tensorflow.concat', 'tf.concat', True, 'import tensorflow as tf\n'), (102, 'tensorflow.variable_scope', 'tf.variable_scope', True, 'import tensorflow as tf\n'), (132, 'tensorflow.concat', 'tf.concat', True, 'import tensorflow as tf\n'), (134, 'tensorflow.variable_scope', 'tf.variable_scope', True, 'import tensorflow as tf\n'), (156, 'tensorflow.variable_scope', 'tf.variable_scope', True, 'import tensorflow as tf\n'), (184, 'tensorflow.name_scope', 'tf.name_scope', True, 'import tensorflow as tf\n'), (191, 'tensorflow.reduce_mean', 'tf.reduce_mean', True, 'import tensorflow as tf\n'), (65, 'tensorflow.concat', 'tf.concat', True, 'import tensorflow as tf\n'), (66, 'tensorflow.concat', 'tf.concat', True, 'import tensorflow as tf\n'), (68, 'tensorflow.concat', 'tf.concat', True, 'import tensorflow as tf\n'), (173, 'tensorflow.name_scope', 'tf.name_scope', True, 'import tensorflow as tf\n'), (179, 'tensorflow.random_uniform', 'tf.random_uniform', True, 'import tensorflow as tf\n'), (189, 'tensorflow.gradients', 'tf.gradients', True, 'import tensorflow as tf\n'), (190, 'tensorflow.contrib.slim.flatten', 'slim.flatten', True, 'import tensorflow.contrib.slim as slim\n'), (49, 'tensorflow.image.resize_bicubic', 'tf.image.resize_bicubic', True, 'import tensorflow as tf\n'), (97, 'tflib.shape', 'tl.shape', True, 'import tflib as tl\n'), (98, 'tflib.shape', 'tl.shape', True, 'import tflib as tl\n'), (98, 'tflib.shape', 'tl.shape', True, 'import tflib as tl\n'), (129, 'tflib.shape', 'tl.shape', True, 'import tflib as tl\n'), (130, 'tflib.shape', 'tl.shape', True, 'import tflib as tl\n'), (130, 'tflib.shape', 'tl.shape', True, 'import tflib as tl\n'), (175, 'tensorflow.shape', 'tf.shape', True, 'import tensorflow as tf\n'), (178, 'tensorflow.shape', 'tf.shape', True, 'import tensorflow as tf\n'), (177, 'tensorflow.sqrt', 'tf.sqrt', True, 'import tensorflow as tf\n')]