# autopep8: off
import tensorflow as tf

physical_devices = tf.config.experimental.list_physical_devices('GPU')
print(physical_devices)
tf.config.experimental.set_memory_growth(physical_devices[0], True)

import colors
from aesthetic_loss import aesthetic_loss
from polymap import Polymap
from tensorflow.keras.layers import Dense, Conv2D, Reshape, ReLU, Conv2DTranspose, LeakyReLU, BatchNormalization, GaussianNoise
# autopep8: on

batch_size = 256


input_small = tf.random.normal(shape=[4, 2])
pm = Polymap(8, 3)
rl = ReLU()
d1 = Dense(256, activation="relu",
           kernel_regularizer=tf.keras.regularizers.l1_l2())
d2 = Dense(512, activation="relu",
           kernel_regularizer=tf.keras.regularizers.l1_l2())
d3 = Dense(1024, activation="relu",
           kernel_regularizer=tf.keras.regularizers.l1_l2())
rs = Reshape([2, 2, 1024//4])
dc1 = Conv2DTranspose(256, (5, 5),
                      padding="valid", use_bias=False)
dc2 = Conv2DTranspose(128, (3, 3),
                      padding="valid", use_bias=False)
dc3 = Conv2DTranspose(64, (3, 3),
                      padding="same", use_bias=False)
dc4 = Conv2D(2, (1, 1), activation="tanh", use_bias=False)
bn1 = BatchNormalization()
bn2 = BatchNormalization()
bn3 = BatchNormalization()

def gen_model():
    model = tf.keras.Sequential()

    model.add(pm)
    # model.add(rl)

    model.add(d1)
    model.add(d2)
    model.add(d3)

    model.add(rs)

    model.add(dc1)
    model.add(bn1)
    model.add(LeakyReLU())

    model.add(dc2)
    model.add(bn2)
    model.add(LeakyReLU())

    model.add(dc3)
    model.add(bn3)
    model.add(LeakyReLU())

    model.add(dc4)
    return model


optimizer = tf.keras.optimizers.Adam(learning_rate=0.001, epsilon=0.001)
train_loss = tf.keras.metrics.Mean(name='train_loss')

generator = gen_model()

rng = tf.random.Generator.from_seed(1)

@tf.function
def train_step(input):
    # input = tf.random.normal(shape=[batch_size, 2])
    # input = rng.normal(shape=[batch_size, 2])

    # print(input)

    with tf.GradientTape() as tape:
        # pics = gen(input)
        pics = generator(input, training=True)
        loss = aesthetic_loss(
            pics) + tf.reduce_sum([d1.losses, d2.losses, d3.losses])
        + 0.01 * tf.reduce_mean(tf.abs(pm.trainable_weights))
    # tv = sum([pm.trainable_weights,
    #           d1.trainable_weights,
    #           d2.trainable_weights,
    #           d3.trainable_weights,
    #          dc1.trainable_weights,
    #          dc2.trainable_weights,
    #          dc3.trainable_weights,], [])

    # print(tv)
    gradients = tape.gradient(loss, generator.trainable_weights)
    optimizer.apply_gradients(zip(gradients, generator.trainable_weights))

    train_loss(loss + tf.reduce_sum([d1.losses, d2.losses, d3.losses]))


EPOCHS = 200

for epoch in range(EPOCHS):
    train_loss.reset_states()

    input = rng.normal(shape=[batch_size, 2])
    # print(input)

    for i in range(100):
        train_step(input)

    print(
        f'Epoch {epoch + 1}, '
        f'Loss: {train_loss.result()}, '
    )

    input_small = rng.normal(shape=[4, 2])
    ps = generator(input_small, training=False)
    # print(ps[0])
    # print(pm.kernels)
    s0 = colors.vector_to_string(ps[0])
    s1 = colors.vector_to_string(ps[1])
    s2 = colors.vector_to_string(ps[2])
    s3 = colors.vector_to_string(ps[3])

    print("\n".join(" ".join(r)
          for r in zip(s0.split("\n"), s1.split("\n"), s2.split("\n"), s3.split("\n"))))