# import io
# import base64
from PIL import Image
from fastai.vision.all import load_learner
from binary2image import get_size, save_file, get_binary_data
import gradio as gr
import numpy as np
import pandas as pd
import plotly.express as px

from scipy import stats
import pickle

## Loading the models
entropy_classifier = pickle.load(open('entropy_tester_classifier.pkl', 'rb'))
model_NonObf = load_learner("model_non_obfuscated.pkl", cpu=True) # change to "model_NonObf.pkl"
model_Obf = load_learner("resnet18_obfuscated_samples.pkl", cpu=True) # change to "model_Shikata.pkl"

def entropy_tester(bin_data):
    entropy = stats.entropy(bin_data,base=2)
    pred = entropy_classifier.predict(
        pd.DataFrame(data = entropy.reshape(1,-1), columns=['Entropy'])
        )
    return pred[0]

def process_file(file):
    greyscale_data = get_binary_data(file.name)

    pred_entropy = entropy_tester(greyscale_data)

    bin_size = get_size(len(greyscale_data))
    save_file(".", "tempfile", greyscale_data, bin_size)

    converted_filename = "tempfile.png"

    # Make prediction depending on the type of obfuscation
    if pred_entropy == "NonObf":
        prediction, _, probas = model_NonObf.predict(converted_filename)
    elif pred_entropy == "Obfuscated":
        prediction, _, probas = model_Obf.predict(converted_filename)

    message = f"Your file is {prediction}!!!"
    
    if pred_entropy == "NonObf": pred_entropy = "Non-Obfuscated"
    
    # Convert probas to percentages
    probas_percentage = [float(prob) * 100 for prob in probas]
    
    fig = px.bar(x=["Goodware", "Malware"], 
                 y=probas_percentage, labels={'x':'Type', 'y':'Probability (%)'},
                 height=300)

    
    return pred_entropy, message, fig, Image.open("tempfile.png")

# Define the layout using Blocks, Row, and Column
with gr.Blocks() as demo:
    with gr.Column(scale=1):
        file = gr.File(label="Upload Executable File")
    with gr.Row():
        with gr.Column():            
            text1 = gr.Textbox(label="Prediction")
            prob_dist_img = gr.Plot(label="Probability Distribution")
        with gr.Column():
            text0 = gr.Textbox(label="Type of Obfuscation")
            converted_img = gr.Image(label="Converted Image", height=300, width=300)
    
    button = gr.Button(value="Process File")
    button.click(process_file, inputs=[file], outputs=[text0, text1, prob_dist_img, converted_img])

demo.launch()