from albumentations.augmentations.utils import P
from PythonDataTransferMQRPC import PythonDataTransferMQRPC
import albumentations as A
import EgLib as El
import tensorflow as tf
import numpy as np
from sklearn.model_selection import train_test_split
from PIL import Image
from tensorflow.keras.utils import Sequence
import time

class DataGenerator(Sequence):
    def __init__(self, x_set, y_set, batch_size):
        self.x, self.y = x_set, y_set
        self.batch_size = batch_size

    def __len__(self):
        return int(np.ceil(len(self.x) / float(self.batch_size)))

    def __getitem__(self, idx):
        batch_x = self.x[idx * self.batch_size:(idx + 1) * self.batch_size]
        batch_y = self.y[idx * self.batch_size:(idx + 1) * self.batch_size]
        return batch_x, batch_y

transform = A.Compose([

    A.HorizontalFlip(p=0.5),
    A.VerticalFlip(p=0.5),
    A.Rotate(p=0.8),

])
classes = 25

models = {}



class HandlerClass:

    def __init__(self):
        self.model = None
        self.active_model = None
        self.val_dataset = None

        self.train_dataset = None

        self.Y_test = None
        self.Y_train = None
        self.X_train = None
        self.X_test = None
        self.const={}

    def activate_model(self, name):

        self.model = models[name]
        self.active_model = name

        

        return 0

    def start_learning(self, path):
        print("start learning")
        target_size = self.get_accept_size()
        images, labels = El.load_hawkeye(path, target_size=(target_size[1], target_size[2]))

        images_augmented = []
        labels_augmented = []

        for _ in range(15):
            for i, l in zip(images, labels):
                transformed = transform(image=i, mask=l)
                images_augmented.append(transformed["image"])
                labels_augmented.append(transformed["mask"])
        images, labels = np.array(images_augmented), np.array(labels_augmented)
        labels = np.squeeze(labels[:, :, :, 0])
        labels = tf.one_hot(labels, classes)

        images_normalized = images / 255.
        self.X_train, self.X_test, self.Y_train, self.Y_test = train_test_split(images_normalized, labels.numpy(),
                                                                                test_size=0.33, random_state=42)

        self.train_dataset = DataGenerator(self.X_train, self.Y_train, 16)
        self.val_dataset = DataGenerator(self.X_test, self.Y_test, 16)

        return 0

    def create_model(self, name, additional_layers, classes):
        input_layer = layers.Input(shape=(128,128,3))
        x = layers.Conv2D(16, 7,strides=1, padding='same',activation="relu")(input_layer)
        x = layers.MaxPooling2D()(x)
        x = residual_block(x,16,5)
        x = layers.MaxPooling2D()(x)

        for _ in range(additional_layers):
            x = residual_block(x,16,3)

        x = residual_block(x,16,3)
        x = layers.UpSampling2D()(x)
        x = residual_block(x,16,3)
        x = layers.UpSampling2D()(x)
        x = layers.Conv2D(9, 3, padding='same',activation="softmax")(x)


        model = Model(inputs=input_layer,outputs=x)
        model.summary()

        model.compile(loss="categorical_crossentropy",optimizer=tf.optimizers.Adam(learning_rate=0.01),metrics=["accuracy"])

    def save_model(self, path):
        self.model.save(path)
        return 0

    def load_model(self, path, name):
        models[name]=tf.keras.models.load_model(path)

        
        self.const[name]=None
            
        return 0

    def get_accept_size(self):
        return list(self.model.input_shape)

    def get_output_size(self):
        return list(self.model.layers[-1].output.shape[1:4])

    def predict(self, data):
        target_size = self.get_accept_size()
        image = np.reshape(np.frombuffer(data, dtype=np.uint8), (target_size[1], target_size[2], target_size[3])) / 255.
        if self.const[self.active_model] is not None:
            predictions = self.model.predict([np.array([image]),self.const])
        else:   
            predictions = self.model.predict(np.array([image]))
        
        

        if predictions[0].ndim==2:
            classed = predictions[0]
        else:
            classed = tf.argmax(predictions[0], axis=2)

        print(classed.shape)

        return classed.numpy().reshape(-1).tolist()

    def predict_anomaly(self,data):
        target_size = self.get_accept_size()
        image = np.reshape(np.frombuffer(data, dtype=np.uint8), (target_size[1], target_size[2], target_size[3])) / 255.
        inputs=[np.array([image]),np.array([self.const[self.active_model]])]
        
        predictions = self.model.predict(inputs)
        print(predictions[0][0][0].shape)
        print(predictions[1].shape)
        
        image=predictions[0][0][0]
                

        normalized_image=(image-image.min())/(image.max()-image.min())
        defect_probability=predictions[1][0][1]*100
        squized_image=(normalized_image*255).reshape(-1)
        
        squized_prop=np.array([defect_probability*100]).astype(int)
        
        res = np.concatenate((squized_prop,(normalized_image*255).astype(int).reshape(-1))).astype(np.uint8).tolist()
        print(len(res))
        return res

    def predict_raw(self, data):
        target_size = self.get_accept_size()
        image = np.reshape(np.frombuffer(data, dtype=np.uint8), (target_size[1], target_size[2], target_size[3])) / 255.

        # measure time
        strart_time=time.time()

        if self.const[self.active_model] is not None:
            predictions = self.model.predict([np.array([image]),self.const])
        else:   
            print("test")
            predictions = self.model(np.array([image])).numpy()

        print(time.time()-strart_time)

        if predictions[0].ndim==2:
            classed = predictions[0][:,:,np.newaxis].clip(0,255)
            return (classed).astype(np.uint8).reshape(-1).tolist()
        else:
            classed = predictions[0][:,:,:]

        print(classed.max())
        print(classed.min())
        return (classed*255).astype(np.uint8).reshape(-1).tolist()



    def continue_learning(self):
        print("continue learning")
        schedule = tf.keras.callbacks.LearningRateScheduler(
            lambda epoch: 0.001, verbose=0)

        model_history = self.model.fit(self.train_dataset, epochs=1,
                                       validation_data=self.val_dataset,
                                       callbacks=[schedule])
        tf.keras.backend.clear_session()
        return model_history.history["val_loss"][0]


data_transfer = PythonDataTransferMQRPC()

data_transfer.wait_for_data_and_process(HandlerClass())
