import gradio as gr
import pandas as pd
import joblib
from huggingface_hub import hf_hub_download

# Download the model from Hugging Face hub
model_filename = hf_hub_download(repo_id="poudel/fuel-burn-predictor", filename="fuel_burn_model.pkl")

# Load the model
model = joblib.load(model_filename)

# Define the prediction function
def predict_fuel_burn_kg(truck_id, kms, litros):
    # Map truck ID input to match the expected format
    truck_ids = {'Truck_ID_MTP3482': 0, 'Truck_ID_MTP5052': 1, 'Truck_ID_MTP5126': 2}
    truck_id_num = truck_ids.get(truck_id, -1)  # Convert Truck ID to numerical representation
    
    # Create a dataframe with the input data
    input_data = pd.DataFrame({
        'Truck_ID': [truck_id_num],
        'Kms': [kms],
        'Litros': [litros]
    })
    
    # Predict fuel burn in liters
    prediction_litros = model.predict(input_data)
    
    # Convert liters to kilograms (using diesel density of 0.835 kg/liter)
    density_kg_per_liter = 0.835
    prediction_kg = prediction_litros[0] * density_kg_per_liter
    
    return round(prediction_kg, 2)

# Create the Gradio interface
app = gr.Interface(
    fn=predict_fuel_burn_kg,
    inputs=[
        gr.Dropdown(['Truck_ID_MTP3482', 'Truck_ID_MTP5052', 'Truck_ID_MTP5126'], label="Truck ID"),
        gr.Number(label="Kilometers Driven"),
        gr.Number(label="Fuel Consumed (Liters)")
    ],
    outputs=gr.Number(label="Predicted Fuel Burn (kg)"),
    title="Truck Fuel Burn Predictor",
    description="Enter the truck ID, kilometers driven, and fuel consumed in liters to predict fuel burn in kilograms."
)

# Launch the Gradio app
app.launch()