import numpy as np

import matplotlib.pyplot as plt

import tensorflow as tf
from tensorflow.keras import optimizers, callbacks

from tensorflow.keras.models import Model
from tensorflow.keras.layers import Input, Conv2D, BatchNormalization, Activation, MaxPool2D

from utils import data, plot, op

from loc import build, functional
from loc.models import AutoSegmentationBaseline

SEED = 2023
SHAPE = (224, 224, 3)
TEST_SIZE = 0.2
BATCH_SIZE = 64

DOWNLOAD = True
WORKING_DIR = './'
SUB_DIR = 'flowers_17'

tf.random.set_seed(SEED)

train_data, test_data = data.prepare_and_load_17flowers(download=DOWNLOAD, working_dir=WORKING_DIR, sub_dir=SUB_DIR,
                                                        shape2d=SHAPE[:2], test_size=TEST_SIZE,
                                                        train_batch_size=BATCH_SIZE, random_state=SEED)


def reduce2d(inputs, filters, strides):

    z = Conv2D(filters=filters[0], kernel_size=(3, 3), strides=strides[0], padding='same')(inputs)
    z = BatchNormalization()(z)
    z = Activation('relu')(z)

    z = Conv2D(filters=filters[1], kernel_size=(3, 3), strides=strides[0], padding='same')(z)
    z = BatchNormalization()(z)
    z = Activation('relu')(z)

    z = MaxPool2D(pool_size=2)(z)

    return z


def build_encoder(summary=True):

    inputs = Input(shape=SHAPE)

    # -------------------------------------------------------------------------------------

    z = reduce2d(inputs, [8, 8], [1, 2])
    z = reduce2d(z, [8, 8], [2, 2])

    # -------------------------------------------------------------------------------------

    latent = Conv2D(filters=1, kernel_size=(1, 1))(z)
    latent = Activation('sigmoid', name='latent')(latent)

    # -------------------------------------------------------------------------------------

    _model = Model(inputs=inputs, outputs=latent)

    # -------------------------------------------------------------------------------------

    if summary:

        _model.summary()

    return _model


def build_loc2d(g_shape, summary=True):

    inputs = Input(shape=SHAPE)
    latent = Input(shape=g_shape)

    outputs = build.LOC2D(padding='same', strides=1, name='loc', trainable=False)([inputs, latent])

    _model = Model(inputs=[inputs, latent], outputs=outputs)

    # -------------------------------------------------------------------------------------

    if summary:

        _model.summary()

    return _model


encoder = build_encoder(summary=False)
loc2d = build_loc2d(encoder.output_shape[1:], summary=True)

autoseg = AutoSegmentationBaseline(input_shape=SHAPE, encoder=encoder, loc2d=loc2d,
                                   rho=0.5, encoder_loss=None)

optimizer = optimizers.Adam(learning_rate=1e-3, beta_1=0.9, beta_2=0.999)

autoseg.summary()

autoseg.compile(optimizer=optimizer, metrics=[functional.partitions_reduction, functional.iou])

# autoseg.load_weights('loc_seg_17.h5')

early_stopping = callbacks.EarlyStopping(monitor='loss', min_delta=0.0, patience=1)

autoseg.fit(train_data, epochs=50, callbacks=[early_stopping])

autoseg.save('loc_seg_17.h5')

# ===============================================================================================================

accuracy_score, iou_score = functional.benchmark_evaluate(model=autoseg, tf_data=test_data)

print('Accuracy : {}'.format(accuracy_score))
print('IOU : {}'.format(iou_score))

# ===============================================================================================================

counts = functional.benchmark_clusters_count(model=autoseg, tf_data=test_data)
plot.clusters_count(counts, title=None)
plt.show()

# ===============================================================================================================

plt.style.use('classic')

sample_size = 4

sample_x, sample_y = test_data.unbatch().batch(sample_size).as_numpy_iterator().next()
sample_y = sample_y.squeeze()

sample_multi_class = autoseg.predict_on_batch(sample_x).squeeze().astype('int32')
sample_output = op.multiclass_to_bitmask(sample_multi_class)

_ = plot.show_images(sample_x, nrows=1, ncols=sample_size)

sample = np.concatenate([sample_y, sample_multi_class, sample_output], axis=0)

_ = plot.show_images(sample, nrows=3, ncols=sample_size)

plt.tight_layout()
plt.show()
