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

# ---------------- Load CSV ----------------
CSV_PATH = "best_models_summary.csv"
df = pd.read_csv(CSV_PATH)

# ---------------- Prediction Function ----------------
def predict_features(uploaded_file, dataset_type, metric_priority):

    try:
        features = np.load(uploaded_file)  # user uploads .npy feature
    except Exception as e:
        return f"❌ Error loading feature file: {e}"

    # Ensure correct shape
    if len(features.shape) == 1:
        features = features.reshape(1, -1)

    # ---------------- Select Best Model ----------------
    subset = df[df['dataset'] == dataset_type]

    if metric_priority != "best_overall":
        subset = subset.sort_values(by=metric_priority, ascending=False)
    else:
        subset = subset.sort_values(by='f1_score', ascending=False)

    model_row = subset.iloc[0]

    model_path = model_row['model_path']
    model_name = model_row['model_name']

    # ---------------- Load Model ----------------
    try:
        model = joblib.load(model_path)
    except Exception as e:
        return f"❌ Error loading model: {e}"

    # ---------------- Prediction ----------------
    try:
        pred = model.predict(features)[0]
    except Exception as e:
        return f"❌ Prediction error: {e}"

    # ---------------- Output ----------------
    result_text = f"""
### 🔥 Prediction Result

**Model Used:** {model_name}  
**Dataset:** {dataset_type}  
**Predicted Class:** {pred}  

**Metrics:**  
Accuracy: {model_row['accuracy']:.4f}  
Precision: {model_row['precision']:.4f}  
Recall: {model_row['recall']:.4f}  
F1 Score: {model_row['f1_score']:.4f}  
"""

    return result_text


# ---------------- Gradio UI ----------------
with gr.Blocks(css="""
    h1 {color: #ff6f61; text-align: center;}
    .gr-button {background-color: #ff6f61; color:white; font-weight:bold;}
""") as demo:

    gr.Markdown("<h1>🔥 Wildfire Feature Classification 🔥</h1>")
    gr.Markdown("Upload a feature vector (.npy file) extracted using Xception.")

    with gr.Row():
        uploaded_file = gr.File(label="Upload Feature (.npy)", file_types=['.npy'])
        dataset_type = gr.Dropdown(choices=['satellite','uav'], label="Dataset Type")
        metric_priority = gr.Dropdown(
            choices=['accuracy','precision','recall','f1_score','best_overall'],
            label="Metric Priority"
        )

    output_text = gr.Markdown()

    predict_btn = gr.Button("Predict 🔥")

    predict_btn.click(
        fn=predict_features,
        inputs=[uploaded_file, dataset_type, metric_priority],
        outputs=output_text
    )

# Launch
demo.launch()