import gradio as gr
from matplotlib import cm 
import matplotlib.pyplot as plt
from mpl_toolkits.axes_grid1 import make_axes_locatable
import numpy as np
import onnxruntime as ort
from PIL import Image
from scipy import special

# model_path = 'chlab/planet_detection_models/'
model_path = './models/'

# plotting a prameters
labels = 20
ticks = 14
legends = 14
text = 14
titles = 22
lw = 3
ps = 200
cmap = 'magma'

def normalize_array(x: list):

    '''Makes array between 0 and 1'''
    
    x = np.array(x)
    
    return (x - np.min(x)) / np.max(x - np.min(x))

def load_model(model: str, activation: bool=True):
    
    if activation:
        model += '_w_activation'
    
    options = ort.SessionOptions()
    options.intra_op_num_threads = 1
    options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
    provider = "CPUExecutionProvider"
    ort_session = ort.InferenceSession(model_path + '%s.onnx' % (model), options, providers=[provider])
    # ort_session = ORTModel.load_model(model_path + '%s.onnx' % (model))
    
    return ort_session

def get_activations(intermediate_model, image: list,
                    layer=None, vmax=2.5, sub_mean=True):
    
    '''Gets activations for a given input image'''
    
    # run model
    input_name = intermediate_model.get_inputs()[0].name
    outputs = intermediate_model.run(None, {input_name: image})
    
    # get activations
    output_1 = outputs[1]
    output_2 = outputs[2]
    
    # get prediction
    output = outputs[0][0]
    output = special.softmax(output)
    
    # sum over velocity channels
    in_image = np.sum(image[0, :, :, :], axis=0)
    in_image = normalize_array(in_image)

    if layer is None:
        # sum over all velocity channels
        activation_1 = np.sum(output_1[0, :, :, :], axis=0)
        activation_2 = np.sum(output_2[0, :, :, :], axis=0)
    else:
        # select a single channel
        activation_1 = output_1[0, layer, :, :]
        activation_2 = output_2[0, layer, :, :]
    
    if sub_mean:
        # y = |x - <x>|
        activation_1 -= np.mean(activation_1)
        activation_1 = np.abs(activation_1)
        
        activation_2 -= np.mean(activation_2)
        activation_2 = np.abs(activation_2)
    
    return output, in_image, activation_1, activation_2


def predict_and_analyze(model_name, num_channels, dim, image):
    
    '''
    Loads a model with activations, passes through image and shows activations
    
    The image must be a numpy array of shape (C, W, W) or (1, C, W, W) 
    '''
    
    model_name = model_name.lower()
    num_channels = int(num_channels)
    W = int(dim)

    print("Loading data")
    image = np.load(image.name, allow_pickle=True)
    image = image.astype(np.float32)
    
    if len(image.shape) != 4:
        image = image[np.newaxis, :, :, :]
        
    assert image.shape == (1, num_channels, W, W), "Data is the wrong shape"
    
    model_name += '_%i' % (num_channels)
    
    print("Loading model")
    model = load_model(model_name, activation=True)
    print("Model loaded")
    
    print("Looking at activations")
    output, input_image, activation_1, activation_2 = get_activations(model, image, sub_mean=True)
    print("Activations and predictions finished")
    
    if output[0] < output[1]:
        output = 'Planet predicted with %.3f percent confidence' % (100*output[1])
    else:
        output = 'No planet predicted with %.3f percent confidence' % (100*output[0])
        
    input_image = normalize_array(input_image)
    activation_1 = normalize_array(activation_1)
    activation_2 = normalize_array(activation_2)
    
    # convert input image to RGB
    input_pil_image = Image.fromarray(np.uint8(cm.magma(input_image)*255))
    
    print("Plotting")
    
    origin = 'lower'
    
    ##### Make the activation figure ######
    plt.rcParams['xtick.labelsize'] = ticks
    plt.rcParams['ytick.labelsize'] = ticks
    
    fig, axs = plt.subplots(nrows=1, ncols=2, figsize=(27, 12))
    
    ax1, ax2 = axs[0], axs[1]
    
    # im0 = ax0.imshow(input_image, cmap=cmap,
    #                  origin=origin)    
    im1 = ax1.imshow(activation_1, cmap=cmap,
                     origin=origin)
    im2 = ax2.imshow(activation_2, cmap=cmap, 
                     origin=origin) 
    
    ims = [im1, im2]
    
    for (i, ax) in enumerate(axs):
        divider = make_axes_locatable(ax)
        cax = divider.append_axes('right', size='5%', pad=0.05)
        fig.colorbar(ims[i], cax=cax, orientation='vertical')
        
    # ax0.set_title('Input', fontsize=titles)
    ax1.set_title('Activation 1', fontsize=titles)
    ax2.set_title('Activation 2', fontsize=titles)
    
    
    ##### make the figure for the input image #####
    plt.rcParams['xtick.labelsize'] = ticks
    plt.rcParams['ytick.labelsize'] = ticks
    
    input_fig, ax = plt.subplots(nrows=1, ncols=1, figsize=(9, 8))
    
    im0 = ax.imshow(input_image, cmap=cmap,
                     origin=origin)    
    
    divider = make_axes_locatable(ax)
    cax = divider.append_axes('right', size='5%', pad=0.05)
    fig.colorbar(im0, cax=cax, orientation='vertical')
        
    ax.set_title('Input', fontsize=titles)
    
    print("Sending to Hugging Face")
    
    return output, input_fig, fig


if __name__ == "__main__":

    demo = gr.Interface(
        fn=predict_and_analyze,
        inputs=[gr.Dropdown(["EfficientNet", "RegNet"], 
                            value="EfficientNet",
                            label="Model Selection",
                            show_label=True), 
                gr.Dropdown(["45", "61", "75"], 
                            value="61",
                            label="Number of Velocity Channels",
                            show_label=True), 
                gr.Dropdown(["600"], 
                            value="600",
                            label="Image Dimensions",
                            show_label=True), 
                gr.File(label="Input Data", show_label=True)],
        outputs=[gr.Textbox(lines=1, label="Prediction", show_label=True), 
                # gr.Image(label="Input Image", show_label=True), 
                gr.Plot(label="Input Image", show_label=True), 
                #  gr.Image(label="Activation 1", show_label=True), 
                #  gr.Image(label="Actication 2", show_label=True)],
                gr.Plot(label="Activations", show_label=True) 
                ],
        title="Kinematic Planet Detector"
    )
    demo.launch()