<reponame>Darkar25/HyperGAN
import tensorflow as tf
import hyperchamber as hc
from hypergan.ops.tensorflow.activations import minmaxzero

from hypergan.losses.base_loss import BaseLoss

class BoundaryEquilibriumLoss(BaseLoss):
    def required(self):
        return "type use_k reduce k_lambda gamma initial_k".split()

    # boundary equilibrium gan
    def began(self, d_real, d_fake):
        gan = self.gan
        config = self.config
        ops = self.gan.ops

        a,b,c = config.labels or [0,1,1]

        d_real = config.reduce(d_real)
        d_fake = config.reduce(d_fake)

        k = tf.get_variable(gan.ops.generate_scope()+'k', [1], initializer=tf.constant_initializer(config.initial_k), dtype=config.dtype)

        if config.type == 'wgan':
            l_x = d_real
            l_dg =-d_fake
            g_loss = d_fake
        elif config.type == 'least-squares':
            l_x = tf.square(d_real-b)
            l_dg = tf.square(d_fake - a)
            g_loss = tf.square(d_fake - c)
        else:
            print("No loss defined.  Get ready to crash")

        if config.use_k:
            d_loss = l_x+k*l_dg
        else:
            d_loss = l_x+l_dg

        gamma = config.gamma or 0.5
        gamma_d_real = gamma*d_real

        ### VERIFY FROM HERE
        k_loss = gamma_d_real - g_loss
        clip = k + config.k_lambda * k_loss
        clip = tf.clip_by_value(clip, 0, 1)
        clip = tf.reduce_mean(clip, axis=0)
        update_k = tf.assign(k, tf.reshape(clip, [1]))
        measure = self.gan.ops.squash(l_x + tf.abs(k_loss))

        d_loss = ops.reshape(d_loss, [])
        g_loss = ops.reshape(g_loss, [])

        return [k, update_k, measure, d_loss, g_loss]


    def _create(self, d_real, d_fake):
        gan = self.gan
        config = self.config

        x = gan.inputs.x
        k, update_k, measure, d_loss, g_loss = self.began(d_real, d_fake)

        self.metrics = {
            'k': k,
            'update_k': update_k, #side effect, this actually trains k
            'measure': measure
        }

        return [d_loss, g_loss]

import tensorflow as tf
import hyperchamber as hc
from hypergan.ops.tensorflow.activations import minmaxzero

from hypergan.losses.base_loss import BaseLoss

class BoundaryEquilibriumLoss(BaseLoss):
    def required(self):
        return "type use_k reduce k_lambda gamma initial_k".split()

    # boundary equilibrium gan
    def began(self, d_real, d_fake):
        gan = self.gan
        config = self.config
        ops = self.gan.ops

        a,b,c = config.labels or [0,1,1]

        d_real = config.reduce(d_real)
        d_fake = config.reduce(d_fake)

        k = tf.get_variable(gan.ops.generate_scope()+'k', [1], initializer=tf.constant_initializer(config.initial_k), dtype=config.dtype)

        if config.type == 'wgan':
            l_x = d_real
            l_dg =-d_fake
            g_loss = d_fake
        elif config.type == 'least-squares':
            l_x = tf.square(d_real-b)
            l_dg = tf.square(d_fake - a)
            g_loss = tf.square(d_fake - c)
        else:
            print("No loss defined.  Get ready to crash")

        if config.use_k:
            d_loss = l_x+k*l_dg
        else:
            d_loss = l_x+l_dg

        gamma = config.gamma or 0.5
        gamma_d_real = gamma*d_real

        ### VERIFY FROM HERE
        k_loss = gamma_d_real - g_loss
        clip = k + config.k_lambda * k_loss
        clip = tf.clip_by_value(clip, 0, 1)
        clip = tf.reduce_mean(clip, axis=0)
        update_k = tf.assign(k, tf.reshape(clip, [1]))
        measure = self.gan.ops.squash(l_x + tf.abs(k_loss))

        d_loss = ops.reshape(d_loss, [])
        g_loss = ops.reshape(g_loss, [])

        return [k, update_k, measure, d_loss, g_loss]


    def _create(self, d_real, d_fake):
        gan = self.gan
        config = self.config

        x = gan.inputs.x
        k, update_k, measure, d_loss, g_loss = self.began(d_real, d_fake)

        self.metrics = {
            'k': k,
            'update_k': update_k, #side effect, this actually trains k
            'measure': measure
        }

        return [d_loss, g_loss]
# boundary equilibrium gan ### VERIFY FROM HERE #side effect, this actually trains k