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

# Load the saved model and LabelEncoder
with open('xgb_model.pkl', 'rb') as f:
    model = pickle.load(f)

with open('le_fartoy.pkl', 'rb') as f:
    le_fartoy = pickle.load(f)

# Define the prediction function
def predict_accidents(fartoy_gruppe, year):
    try:
        # Encode the selected vessel group
        fartoy_encoded = le_fartoy.transform([fartoy_gruppe])[0]
        
        # Features must be in the exact order: ['FartøyGruppe', 'Ulykke_Year']
        features = np.array([[fartoy_encoded, year]])
        
        # Get model prediction
        prediction = model.predict(features)[0]
        
        # Round and ensure non-negative
        predicted_count = max(0, int(round(prediction)))
        return f"Predicted Number of Accidents: {predicted_count}"
    except Exception as e:
        return f"Error during prediction: {str(e)}"

# Get the list of vessel groups for the dropdown
vessel_groups = sorted(list(le_fartoy.classes_))

# Build Gradio UI with custom inputs (Dropdown and Slider)
interface = gr.Interface(
    fn=predict_accidents,
    inputs=[
        gr.Dropdown(choices=vessel_groups, label="FartøyGruppe (Vessel Group)", value=vessel_groups[0]),
        gr.Slider(minimum=1981, maximum=2032, step=1, label="Ulykke_Year (Accident Year)", value=2024)
    ],
    outputs=gr.Textbox(label="Prediction Output"),
    title="Maritime Accident Predictor (XGBoost)",
    description="Select a vessel category and year to predict the annual aggregated number of accidents using the optimized XGBoost Regressor model.",
    theme="soft"
)

if __name__ == '__main__':
    interface.launch()