import os
from keras.models import Model
from matplotlib import pyplot as plt
def generate_visualizations(model,input_image,root_dir):
        try:
            os.mkdir(root_dir)
        except OSError as error:
            pass
        no_conv_layers = 1
        plt.axis('off')
        plt.imshow(input_image)
        plt.savefig(os.path.join(root_dir,'input.jpg'),transparent=True)
        for layer in model.layers:
            if 'conv' in layer.name:
                try:
                    os.mkdir(os.path.join(root_dir,'conv{}'.format(no_conv_layers)))
                except OSError as error:
                    pass
                op = Model(inputs=model.inputs,outputs=model.layers[no_conv_layers-1].output)
                
                # Output images
                pred = op.predict(input_image[np.newaxis,...,np.newaxis])
                index = 0
                count = pred.shape[-1]
                while count != index:
                    plt.imshow(pred[0,:,:,index])
                    plt.savefig(os.path.join(root_dir,'conv{}'.format(no_conv_layers),'{}.jpg'.format(index+1)),transparent=True)
                    index += 1
                
                # Save weights
                weights,bias = layer.get_weights()
                file = open(os.path.join(root_dir,'meta.txt'), "a")
                file.write('{}\n{}\n{}\n'.format(weights.shape[0],weights.shape[2],weights.shape[3]))
                file.close()
                min_val,max_val = weights.min(), weights.max()
                weights = (weights - min_val)/(max_val-min_val+1e-8)
                depth = weights.shape[2]
                dimensions = (100,100)
                for index in range(weights.shape[-1]):
                    weight = weights[:,:,:,index]
                    for channel in range(depth):
                        final_image = cv.resize(weight[:,:,channel], dimensions, interpolation = cv.INTER_AREA)
                        plt.imshow(final_image,cmap='gray')
                        plt.savefig(os.path.join(root_dir,'conv{}'.format(no_conv_layers),'{}-{}.jpg'.format(index+1,channel+1)),transparent=True)
                
                no_conv_layers += 1
