import gradio as gr
import pandas as pd
import numpy as np
import joblib
import tensorflow as tf


# LOAD MODEL & SCALER
model = tf.keras.models.load_model("mlp_malware.keras")
scaler = joblib.load("scaler.pkl")

feature_names = [
    "filesize",
    "E_file",
    "E_text",
    "E_data",
    "AddressOfEntryPoint",
    "NumberOfSections",
    "SizeOfInitializedData",
    "SizeOfImage",
    "SizeOfOptionalHeader",
    "SizeOfCode",
    "DirectoryEntryImportSize",
    "ImageBase",
    "CheckSum",
    "Magic",
    "MinorLinkerVersion",
    "MajorSubsystemVersion",
    "e_lfanew",
    "sus_sections",
    "PointerToSymbolTable",
    "SectionsLength",
    "SizeOfStackReserve",
    "MajorOperatingSystemVersion",
    "non_sus_sections",
    "Characteristics",
    "NumberOfSymbols",
    "BaseOfData",
    "MajorImageVersion",
    "FH_char5",
    "FH_char8",
    "OH_DLLchar5"
]

# PREDICTION FUNCTION
def predict_malware(*inputs):
    # inputs → DataFrame
    X = pd.DataFrame([inputs], columns=feature_names)

    # scale
    X_scaled = scaler.transform(X)

    # predict
    prob = model.predict(X_scaled, verbose=0)[0][0]
    pred = int(prob >= 0.5)

    label = "Malware" if pred == 1 else "Benign"
    
    # Return dict for API consistency
    return {
        "prediction_label": label,
        "probability_malware": float(prob),
        "prediction": pred
    }

# UI
inputs = [
    gr.Number(label=feat, value=0.0)
    for feat in feature_names
]

outputs = gr.JSON(label="Prediction Result")

app = gr.Interface(
    fn=predict_malware,
    inputs=inputs,
    outputs=outputs,
    title="MLP-based Malware Detection",
    description="Malware detection using MLP neural network"
)

if __name__ == "__main__":
    app.launch()