import gradio as gr
import joblib
import numpy as np

# Load your saved model
model = joblib.load('cybersecurity_model.joblib')

# Replace this list with your actual feature names in the order your model expects
FEATURE_NAMES = [
    'A_frequency', 'NS_frequency', 'CNAME_frequency', 'SOA_frequency', 'NULL_frequency',
    'PTR_frequency', 'HINFO_frequency', 'MX_frequency', 'TXT_frequency', 'AAAA_frequency',
    'SRV_frequency', 'OPT_frequency', 'rr_type', 'rr_count', 'rr_name_entropy',
    'rr_name_length', 'distinct_ns', 'distinct_ip', 'unique_country', 'unique_asn',
    'distinct_domains', 'reverse_dns', 'a_records', 'unique_ttl', 'ttl_mean',
    'ttl_variance', 'timestamp', 'FQDN_count', 'subdomain_length', 'upper',
    'lower', 'numeric', 'entropy'
]

def predict(*inputs):
    # Convert inputs to numpy array, reshape for prediction
    data = np.array(inputs).reshape(1, -1)
    
    # Predict class label
    prediction = model.predict(data)[0]
    
    # Optionally, return class probabilities:
    # proba = model.predict_proba(data)[0]
    # return prediction, proba
    
    return prediction

# Create Gradio inputs dynamically based on feature names
inputs = [gr.Number(label=feat) for feat in FEATURE_NAMES]

# Gradio Interface setup
iface = gr.Interface(
    fn=predict,
    inputs=inputs,
    outputs="text",
    title="Cybersecurity Attack Detection",
    description="Input the feature values and predict whether the network activity is benign or an attack."
)

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