import gradio as gr
import pandas as pd
import numpy as np

num_labels = 6
headers = ["Benign","GG1", "GG2", "GG3", "GG4", "GG5"]
default_weights = [[0,3,5,5,5,5],
                    [1,0,1,3,5,5],
                    [3,1,0,1,5,5],
                    [3,3,1,0,5,5],
                    [5,5,5,5,0,1],
                    [5,5,5,5,1,0]]
example_conf_mats = [
    pd.DataFrame([
                [80,0,0,0,0,0],
                [20,80,0,0,0,0],
                [0,20,80,20,0,0],
                [0,0,20,80,0,0],
                [0,0,0,0,80,20],
                [0,0,0,0,20,80]],columns=headers),
    pd.DataFrame([
                [80,0,0,0,0,0],
                [15,80,10,0,0,0],
                [5,15,80,15,0,0],
                [0,5,10,80,10,10],
                [0,0,0,5,80,10],
                [0,0,0,0,10,80]],columns=headers),
    pd.DataFrame([
                [80,10, 3, 0, 0, 0],
                [10,80, 7, 5, 3, 0],
                [ 7, 7,80,10, 5, 3],
                [ 3, 3, 7,80, 7, 7],
                [ 0, 0, 3, 5,80,10],
                [ 0, 0, 0, 0, 5,80]],columns=headers),
    pd.DataFrame([
                [80,0,0,0,10,10],
                [0,80,0,0,10,10],
                [0,0,80,0,0,0],
                [0,0,0,80,0,0],
                [10,10,10,10,80,0],
                [10,10,10,10,0,80]],columns=headers)                        
]

def submit_vals(*argv):
    argv = list(argv)
    weights = np.zeros((num_labels, num_labels))
    for i in range(num_labels):
        for j in range(num_labels):
            if i != j:
                weights[i][j] = argv.pop(0)
    weights_df = pd.DataFrame(weights, columns=headers)
    return weights_df

def get_acc(input_df, weights_df):
    input_df = input_df.astype(int)
    total = sum(sum(np.array(input_df)))
    diag_total = sum(np.diag(input_df))
    non_diag_total = total - diag_total
    accuracy = 100 * diag_total / total
    if non_diag_total == 0:
        severity = 0
    else:
        severity = sum(sum(np.array(input_df.multiply(weights_df)))) / non_diag_total
    return accuracy, severity

with gr.Blocks() as demo:
    with gr.Tab("Error severity matrix"):
        with gr.Row():
            with gr.Column():            
                sliders = []
                for i in range(num_labels):
                    for j in range(num_labels):
                        if i != j:
                            sliders.append(gr.Slider(1, 5, value=default_weights[i][j], step=1, label="Impact of misclassifying "+ headers[j] + " as " + headers[i]))          
            with gr.Column():
                output_err_mat = gr.Dataframe(value = default_weights, datatype = "number", row_count = (num_labels, "fixed"), col_count=(num_labels,"fixed"), label="Error Severity Matrix", interactive=0, headers=headers)
                refresh_btn = gr.Button("Refresh")  
            refresh_btn.click(submit_vals, inputs=sliders, outputs=output_err_mat)
    with gr.Tab("Calculate accuracy and Error Severity"):
        with gr.Row():
            with gr.Column():
                conf_df = gr.Dataframe(datatype = "number", row_count = (num_labels, "fixed"), col_count=(num_labels,"fixed"), label="Confusion Matrix", interactive=1, headers=headers)
                submit_btn = gr.Button("Submit")
                examples = gr.Examples(examples=example_conf_mats, inputs=[conf_df])
            with gr.Column():
                outputs = [gr.Textbox(label="Accuracy"), gr.Textbox(label="Error Severity")]
            submit_btn.click(fn=get_acc, inputs=[conf_df,output_err_mat], outputs=outputs)
demo.launch()