import tensorflow as tf
from tensorflow import keras
from tensorflow.keras.layers import *
import numpy as np
import tensorflow.keras.initializers as initer
import matplotlib.pyplot as plt
from tqdm import tqdm
import pickle
"""
# Model
"""
class InstanceNormalization(Layer):
    def __init__(self, axis=(1, 2), epsilon=1e-6):
        super().__init__()
        self.epsilon = epsilon
        self.axis = axis
        self.beta, self.gamma = None, None

    def build(self, input_shape):
        shape = [1 for _ in range(len(input_shape))]
        shape[-1] = input_shape[-1]
        self.gamma = self.add_weight(
            name='gamma',
            shape=shape,
            initializer='ones')

        self.beta = self.add_weight(
            name='beta',
            shape=shape,
            initializer='zeros')

    def call(self, x, *args, **kwargs):
        mean = tf.math.reduce_mean(x, axis=self.axis, keepdims=True)
        x -= mean
        variance = tf.reduce_mean(tf.math.square(x), axis=self.axis, keepdims=True)
        x *= tf.math.rsqrt(variance + self.epsilon)
        return x * self.gamma + self.beta
class AdaNorm(Layer):
    def __init__(self, axis=(1, 2), epsilon=1e-6):
        super().__init__()
        # NHWC
        self.axis = axis
        self.epsilon = epsilon

    def call(self, x, **kwargs):
        mean = tf.math.reduce_mean(x, axis=self.axis, keepdims=True)
        x -= mean
        variance = tf.reduce_mean(tf.math.square(x), axis=self.axis, keepdims=True)
        x *= tf.math.rsqrt(variance + self.epsilon)
        return x
class AdaMod(Layer):
    def __init__(self):
        super().__init__()
        self.y = None

    def call(self, inputs, **kwargs):
        x, w = inputs
        y = self.y(w)
        o = (y[:, 0] + 1) * x + y[:, 1]
        return o

    def build(self, input_shape):
        x_shape, w_shape = input_shape
        self.y = keras.Sequential([
            Dense(x_shape[-1]*2, input_shape=w_shape[1:], kernel_initializer=initer.HeNormal()),
            Reshape([2, 1, 1, -1]),
        ])  # [2, h, w, c] per feature map
class AddNoise(Layer):
    def __init__(self):
        super().__init__()
        self.s = None
        self.x_shape = None

    def call(self, inputs, **kwargs):
        x, noise = inputs
        noise_ = noise[:, :self.x_shape[1], :self.x_shape[2], :]
        return self.s * noise_ + x

    def build(self, input_shape):
        self.x_shape, _ = input_shape
        self.s = self.add_weight(name="noise_scale", shape=[1, 1, self.x_shape[-1]],
                                 initializer=initer.RandomNormal(0, 0.05))
class Map(Layer):
    def __init__(self, size, num_layers, norm=None):
        super().__init__()
        self.size = size
        self.num_layers = num_layers
        self.norm_name = norm
        self.f = None

    def call(self, inputs, **kwargs):
        w = self.f(inputs)
        return w

    def build(self, input_shape):
        self.f = keras.Sequential()
        for i in range(self.num_layers):
            if i == 0:
                self.f.add(Dense(self.size, input_shape=input_shape[1:], kernel_initializer=initer.HeNormal()))
                continue
            self.f.add(LeakyReLU(0.2))
            if self.norm_name is not None:
                if self.norm_name.lower() == "batch":
                    self.f.add(BatchNormalization())    # batch norm also increase model collapse
                elif self.norm_name.lower() == "instance":
                    self.f.add(InstanceNormalization(axis=(1,)))        # instance norm increase model collapse

            self.f.add(Dense(self.size, kernel_initializer=initer.HeNormal()))
class Style(Layer):
    def __init__(self, filters, upsampling=True):
        super().__init__()
        self.filters = filters
        self.upsampling = upsampling
        self.ada_mod, self.ada_norm, self.add_noise, self.up, self.conv, self.conv_expend = None, None, None, None, None, None

    def call(self, inputs, **kwargs):
        x, w, noise = inputs
        # x = self.conv_expend(x)     #TODO: may help for styling
        x = self.ada_mod((x, w))
        if self.up is not None:
            x = self.up(x)
        x = self.conv(x)
        x = LeakyReLU(0.2)(x)
        x = self.add_noise((x, noise))
        x = self.ada_norm(x)
        return x

    def build(self, input_shape):
        self.ada_mod = AdaMod()
        self.ada_norm = AdaNorm()
        if self.upsampling:
            self.up = UpSampling2D((2, 2), interpolation="bilinear")
        self.add_noise = AddNoise()
        # self.conv_expend = Conv2D(self.filters*2, 1, 1, kernel_initializer=initer.HeNormal())
        self.conv = Conv2D(self.filters, 3, 1, "same", kernel_initializer=initer.HeNormal())
def get_generator(latent_dim, img_shape , num_layers ,base=256):
    
    n_style_block = 0
    
    const_size = _size = 4
    
    while _size <= img_shape[1]:
        n_style_block += 1
        _size *= 2

    z = keras.Input((n_style_block, latent_dim,), name="z") #(m,blocks,dim)
    noise_ = keras.Input((img_shape[0], img_shape[1]), name="noise") #(m,h,w)
    ones = keras.Input((1,), name="ones")#(m,1)

    w = Map(size=base, num_layers=num_layers)(z)
    noise = tf.expand_dims(noise_, axis=-1) #(m,h,w,1), pixel-wise noise
    const = keras.Sequential([
        Dense(const_size * const_size * base, use_bias=False, name="const", kernel_initializer=initer.HeNormal()),
        Reshape((const_size, const_size, -1)),
    ], name="const")(ones)

    x = AddNoise()((const, noise))
    x = AdaNorm()(x)
    
    
    for i in range(n_style_block):
        x = Style(base, upsampling=False if i == 0 else True)((x, w[:, i], noise))
        
        
    o = Conv2D(img_shape[-1], 7, 1, "same", activation=keras.activations.tanh)(x)

    g = keras.Model([ones, z, noise_], o, name="generator")
    return g, n_style_block
def get_discriminator(img_shape):
    def add_block(filters, do_norm=True, padding="same"):
        model.add(Conv2D(filters, 4, strides=2, padding=padding))
        if do_norm: 
            model.add(InstanceNormalization())
        model.add(LeakyReLU(alpha=0.2))

    model = keras.Sequential([Input(img_shape)], name="d")
    # [n, 128, 128, 3]
    # model.add(GaussianNoise(0.02))
    add_block(32, do_norm=False)   # -> 64^2
    add_block(64)                   # -> 32^2
    add_block(128)                  # -> 16^2
    add_block(256)                  # -> 8^2
    add_block(512, padding="valid")  # -> 4^2
    model.add(Flatten())
    # model.add(GlobalAveragePooling2D())
    model.add(Dense(256))
    model.add(Dense(1))
    return model
class Trainer(keras.Model):
    def __init__(self,img_shape,latent_dim,num_layers=5,lr=2e-4,
                 beta1=0.5, beta2=0.99, lambda_=10, wgan=2):
        '''
        img_shape:(h,w,c)
        z_dim: dim
        '''
        super().__init__()
        self.img_shape=img_shape
        self.latent_dim=latent_dim
        self.wgan=wgan
        self.lambda_=lambda_
        
        self.g,self.n_blocks=get_generator(latent_dim,img_shape,num_layers=num_layers)
        self.d=get_discriminator(img_shape)
        
    
        self.opt = keras.optimizers.Adam(lr, beta_1=beta1, beta_2=beta2)

        
        
    def call(self, inputs, training=None, mask=None):
        '''
        inputs:[ones,z,noise]
        '''
        if isinstance(inputs[0], np.ndarray):
            inputs = [tf.convert_to_tensor(i) for i in inputs]
        return self.g(inputs, training=training)
    
    
    def get_inputs(self, n):
        if np.random.rand() < 0.5:
            available_z = [tf.random.normal((n, 1, self.latent_dim)) for _ in range(2)]
            z = tf.concat([available_z[np.random.randint(0, len(available_z))] for _ in range(self.n_blocks)], axis=1)
        else:
            z = tf.repeat(tf.random.normal((n, 1, self.latent_dim)), self.n_blocks, axis=1)
        noise = tf.random.normal((n, self.img_shape[0], self.img_shape[1]))
        return [tf.ones((n, 1)), z, noise]
    
    @staticmethod
    def w_distance(real, fake):
        # the distance of two data distributions
        return tf.reduce_mean(real) - tf.reduce_mean(fake)
    
    def gp(self, real_img, fake_img):
        e = tf.random.uniform((len(real_img), 1, 1, 1), 0, 1)
        noise_img = e * real_img + (1. - e) * fake_img  # extend distribution space
        with tf.GradientTape() as tape:
            tape.watch(noise_img)
            o = self.d(noise_img)
        g = tape.gradient(o, noise_img)  # image gradients
        g_norm2 = tf.sqrt(tf.reduce_sum(tf.square(g), axis=[1, 2, 3]))  # norm2 penalty
        gp = tf.square(g_norm2 - 1.)
        return tf.reduce_mean(gp)
    
    def train_d(self, img):
        n = len(img)

        with tf.GradientTape() as tape:
            gimg = self.call(self.get_inputs(n), training=False) #(n,h,w,3)
            gp = self.gp(img, gimg)
            pred_fake = self.d(gimg, training=True)
            pred_real = self.d(img, training=True)
            w_distance = -self.w_distance(pred_real, pred_fake)  # maximize W distance
            gp_loss = self.lambda_ * gp
            loss = gp_loss + w_distance
        grads = tape.gradient(loss, self.d.trainable_variables)
        self.opt.apply_gradients(zip(grads, self.d.trainable_variables))
        return gp, w_distance

    def train_g(self, n):
        with tf.GradientTape() as tape:
            gimg = self.call(self.get_inputs(n), training=True)
            pred_fake = self.d(gimg, training=False)
            w_distance = tf.reduce_mean(-pred_fake)  # minimize W distance
        grads = tape.gradient(w_distance, self.g.trainable_variables)
        self.opt.apply_gradients(zip(grads, self.g.trainable_variables))
        return w_distance
    
    @tf.function
    def step(self, img):
        gw = self.train_g(len(img)*2) #this is just a trick, not following original training method
        for _ in range(self.wgan):
            dgp, dw = self.train_d(img)
        return gw, dgp, dw 
    
    '''
    gw: generator loss
    dw: discriminator loss
    dgp: gradient penalty
    '''
"""
# Training
"""
batch_size=32
epochs=4
latent_dim=128
lr=2e-4
b1=0
b2=0.99
w=2
image_size=128
lam=10
num_layers=5
def show(model):
    global z1, z2
    n = 7
    if "z1" not in globals():
        z1 = np.random.normal(0, 1, size=(n, 1, model.latent_dim))
    if "z2" not in globals():
        z2 = np.random.normal(0, 1, size=(n, 1, model.latent_dim))
    n_z1 = 3
    assert n_z1 < model.n_blocks - 1
    noise = np.random.normal(0, 1, [len(z1), model.img_shape[0], model.img_shape[1]])
    
    #mixing style
    inputs = [
        np.ones((len(z1)*n, 1)),
        np.concatenate(
            (z1.repeat(n, axis=0).repeat(n_z1, axis=1),
             np.repeat(np.concatenate([z2 for _ in range(n)], axis=0), model.n_blocks - n_z1, axis=1)),
            axis=1
        ),
        noise.repeat(n, axis=0),
    ]
    
    #marginal style
    z1_inputs = [np.ones((len(z1), 1)), z1.repeat(model.n_blocks, axis=1), noise]
    z2_inputs = [np.ones((len(z2), 1)), z2.repeat(model.n_blocks, axis=1), noise]

    imgs = model.predict(inputs)
    z1_imgs = model.predict(z1_inputs)
    z2_imgs = model.predict(z2_inputs)
    imgs = np.concatenate([z2_imgs, imgs], axis=0)
    rest_imgs = np.concatenate([np.ones([1, 128, 128, 3], dtype=np.float32), z1_imgs], axis=0)
    for i in range(len(rest_imgs)):
        imgs = np.concatenate([imgs[:i * (n+1)], rest_imgs[i:i + 1], imgs[i * (n+1):]], axis=0)
    imgs = (imgs + 1) / 2

    nc, nr = n+1, n+1
    f = plt.figure(0, (nc*1.5, nr*1.5))
    for c in range(nc):
        for r in range(nr):
            i = r * nc + c
            plt.subplot(nr, nc, i + 1)
            plt.imshow(imgs[i])
            plt.axis("off")

    plt.tight_layout()
    plt.show()
"""
* Train
"""
pth='../input/celeba-dataset/img_align_celeba'


ds=tf.keras.preprocessing.image_dataset_from_directory(
    directory=pth,
    image_size=(image_size,image_size),
    batch_size=batch_size,
    seed=0,
    shuffle=True
)
trainer=Trainer((image_size,image_size,3),latent_dim,num_layers,lr,b1,b2,lam,w)
_=trainer.step(tf.random.normal((32,128,128,3))) #to prevent nontype error
try:
    trainer.load_weights('../input/weights/StyleGAN.pt')
    print('load weights successfully')
except:
    print('no weights ')
track={'gw':[],'dgp':[],'dw':[]}
def main(ds):
    aug=Rescaling(1/127.5,offset=-1)
    ds=ds.map(lambda x,y : (aug(x),y))
    ckpt = tf.train.Checkpoint(trainer=trainer)
    ckpt_manager = tf.train.CheckpointManager(ckpt,'./ckpt', max_to_keep=1)
    i=0
    for epoch in range(epochs):
        show(trainer)
        trainer.save_weights('./weights/StyleGAN.pt')
        i+=1
            
        loop=tqdm(ds)
        for x,y in loop:
            gw, dgp, dw = trainer.step(x)
            
            track['gw'].append(gw)
            track['dgp'].append(dgp)
            track['dw'].append(dw)
            
            loop.set_postfix(loss=f'epoch : {epoch}, gw:{gw}, dgp:{dgp}, dw:{dw}')
main(ds)
with open('./track.pickle','wb') as file:
    pickle.dump(track,file)