#!/usr/bin/env python
# -*- coding: utf-8 -*-
# File: BEGAN.py
# Author: Yuxin Wu <ppwwyyxxc@gmail.com>

import os
import argparse

from tensorpack import *
from tensorpack.tfutils.summary import add_moving_summary
from tensorpack.utils.gpu import get_nr_gpu
from tensorpack.utils.globvars import globalns as G
from tensorpack.tfutils.scope_utils import auto_reuse_variable_scope
import tensorflow as tf

from GAN import GANModelDesc, GANTrainer, MultiGPUGANTrainer

"""
Boundary Equilibrium GAN.
See the docstring in DCGAN.py for usage.

A pretrained model on CelebA is at
https://drive.google.com/open?id=0B5uDfUQ1JTglUmgyZV8zQmNOTVU
"""


import DCGAN
G.BATCH = 32
G.Z_DIM = 64
NH = 64
NF = 64
GAMMA = 0.5


class Model(GANModelDesc):
    def _get_inputs(self):
        return [InputDesc(tf.float32, (None, G.SHAPE, G.SHAPE, 3), 'input')]

    @auto_reuse_variable_scope
    def decoder(self, z):
        l = FullyConnected('fc', z, NF * 8 * 8, nl=tf.identity)
        l = tf.reshape(l, [-1, 8, 8, NF])

        with argscope(Conv2D, nl=tf.nn.elu, kernel_shape=3, stride=1):
            l = (LinearWrap(l)
                 .Conv2D('conv1.1', NF)
                 .Conv2D('conv1.2', NF)
                 .tf.image.resize_nearest_neighbor([16, 16], align_corners=True)
                 .Conv2D('conv2.1', NF)
                 .Conv2D('conv2.2', NF)
                 .tf.image.resize_nearest_neighbor([32, 32], align_corners=True)
                 .Conv2D('conv3.1', NF)
                 .Conv2D('conv3.2', NF)
                 .tf.image.resize_nearest_neighbor([64, 64], align_corners=True)
                 .Conv2D('conv4.1', NF)
                 .Conv2D('conv4.2', NF)
                 .Conv2D('conv4.3', 3, nl=tf.identity)())
        return l

    @auto_reuse_variable_scope
    def encoder(self, imgs):
        with argscope(Conv2D, nl=tf.nn.elu, kernel_shape=3, stride=1):
            l = (LinearWrap(imgs)
                 .Conv2D('conv1.1', NF)
                 .Conv2D('conv1.2', NF)
                 .Conv2D('conv1.3', NF * 2)
                 .AvgPooling('pool1', 2)
                 # 32
                 .Conv2D('conv2.1', NF * 2)
                 .Conv2D('conv2.2', NF * 3)
                 .AvgPooling('pool2', 2)
                 # 16
                 .Conv2D('conv3.1', NF * 3)
                 .Conv2D('conv3.2', NF * 4)
                 .AvgPooling('pool3', 2)
                 # 8
                 .Conv2D('conv4.1', NF * 4)
                 .Conv2D('conv4.2', NF * 4)

                 .FullyConnected('fc', NH, nl=tf.identity)())
        return l

    def _build_graph(self, inputs):
        image_pos = inputs[0]
        image_pos = image_pos / 128.0 - 1

        z = tf.random_uniform([G.BATCH, G.Z_DIM], minval=-1, maxval=1, name='z_train')
        z = tf.placeholder_with_default(z, [None, G.Z_DIM], name='z')

        def summary_image(name, x):
            x = (x + 1.0) * 128.0
            x = tf.clip_by_value(x, 0, 255)
            tf.summary.image(name, tf.cast(x, tf.uint8), max_outputs=30)

        with argscope([Conv2D, FullyConnected],
                      W_init=tf.truncated_normal_initializer(stddev=0.02)):
            with tf.variable_scope('gen'):
                image_gen = self.decoder(z)

            with tf.variable_scope('discrim'):
                with tf.variable_scope('enc'):
                    hidden_pos = self.encoder(image_pos)
                    hidden_neg = self.encoder(image_gen)

                with tf.variable_scope('dec'):
                    recon_pos = self.decoder(hidden_pos)
                    recon_neg = self.decoder(hidden_neg)

        with tf.name_scope('viz'):
            summary_image('generated-samples', image_gen)
            summary_image('reconstruct-real', recon_pos)
            summary_image('reconstruct-fake', recon_neg)

        with tf.name_scope('losses'):
            L_pos = tf.reduce_mean(tf.abs(recon_pos - image_pos), name='loss_pos')
            L_neg = tf.reduce_mean(tf.abs(recon_neg - image_gen), name='loss_neg')

            eq = tf.subtract(GAMMA * L_pos, L_neg, name='equilibrium')
            measure = tf.add(L_pos, tf.abs(eq), name='measure')

            kt = tf.get_variable('kt', dtype=tf.float32, initializer=0.0)

            update_kt = kt.assign_add(1e-3 * eq)
            with tf.control_dependencies([update_kt]):
                self.d_loss = tf.subtract(L_pos, kt * L_neg, name='loss_D')
                self.g_loss = L_neg

        add_moving_summary(L_pos, L_neg, eq, measure, self.d_loss)
        tf.summary.scalar('kt-summary', kt)

        self.collect_variables()

    def _get_optimizer(self):
        lr = symbolic_functions.get_scalar_var('learning_rate', 1e-4, summary=True)
        opt = tf.train.AdamOptimizer(lr, beta1=0.5, beta2=0.9)
        return opt


if __name__ == '__main__':
    args = DCGAN.get_args()
    if args.sample:
        DCGAN.sample(Model(), args.load, 'gen/conv4.3/output')
    else:
        assert args.data
        logger.auto_set_dir()

        config = TrainConfig(
            model=Model(),
            dataflow=DCGAN.get_data(args.data),
            callbacks=[
                ModelSaver(),
                StatMonitorParamSetter(
                    'learning_rate', 'measure', lambda x: x * 0.5, 0, 10)
            ],
            steps_per_epoch=500,
            max_epoch=400,
            session_init=SaverRestore(args.load) if args.load else None,
            nr_tower=max(get_nr_gpu(), 1)
        )
        if config.nr_tower == 1:
            GANTrainer(config).train()
        else:
            MultiGPUGANTrainer(config).train()
