import gradio as gr
import pandas as pd
import numpy as np
from PIL import Image
import os
import time
import cloudpickle
import traceback

# =========================
# Configuration
# =========================
IMAGE_SIZE = (64, 64)
LEADERBOARD_PATH = "leaderboard.csv"
LABEL_FILE = "labels.csv"
TEST_DIR = "test_images"

# =========================
# Helpers
# =========================
def image_to_features(image: Image.Image) -> np.ndarray:
    image = image.resize(IMAGE_SIZE)
    return np.array(image.convert("L")).flatten()

def load_test_data():
    if not os.path.exists(LABEL_FILE) or not os.path.exists(TEST_DIR):
        raise FileNotFoundError("Missing labels.csv or test_images/ folder.")
    df = pd.read_csv(LABEL_FILE)
    X_test, y_test = [], []
    for _, row in df.iterrows():
        img_path = os.path.join(TEST_DIR, row["filename"])
        try:
            img = Image.open(img_path)
            X_test.append(image_to_features(img))
            y_test.append(row["label"])
        except Exception as e:
            print(f"❌ Error loading {img_path}: {e}")
    return np.array(X_test), np.array(y_test)

def load_leaderboard():
    if os.path.exists(LEADERBOARD_PATH):
        df = pd.read_csv(LEADERBOARD_PATH)
        return df.sort_values(by=["Accuracy", "Avg Time (ms)"], ascending=[False, True]).reset_index(drop=True)
    else:
        df = pd.DataFrame(columns=["Name", "Accuracy", "Avg Time (ms)"])
        df.to_csv(LEADERBOARD_PATH, index=False)
        return df

# =========================
# Evaluation Logic
# =========================
def evaluate_model(file, name):
    try:
        with open(file.name, "rb") as f:
            model = cloudpickle.load(f)

        X_test, y_test = load_test_data()
        start = time.time()
        y_pred = model.predict(X_test)
        elapsed = (time.time() - start) * 1000

        if len(y_pred) != len(y_test):
            return "❌ Model output length does not match test set.", load_leaderboard()

        accuracy = round(100.0 * (y_pred == y_test).mean(), 5)
        avg_time = round(elapsed / len(X_test), 5)

        leaderboard = load_leaderboard()

        # Remove existing entries for the same name
        leaderboard = leaderboard[leaderboard["Name"] != name]

        # Add new result
        new_entry = pd.DataFrame([{
            "Name": name,
            "Accuracy": accuracy,
            "Avg Time (ms)": avg_time
        }])
        leaderboard = pd.concat([leaderboard, new_entry], ignore_index=True)

        # Keep only the best score per name
        leaderboard = leaderboard.sort_values(by=["Accuracy", "Avg Time (ms)"], ascending=[False, True])
        leaderboard = leaderboard.drop_duplicates(subset=["Name"], keep="first").reset_index(drop=True)

        leaderboard.to_csv(LEADERBOARD_PATH, index=False)

        return "", leaderboard

    except Exception as e:
        return f"❌ Error:\n{traceback.format_exc()}", load_leaderboard()

# =========================
# UI
# =========================
with gr.Blocks(title="Olive Fly Classifier Leaderboard") as demo:
    gr.Markdown("## Olive Fly Classifier Leaderboard")
    gr.Markdown(
        "Upload your `.cloudpkl` model trained on Olive Fly images. "
        "We'll evaluate it and update the leaderboard."
    )

    with gr.Row():
        with gr.Column(scale=1):
            file_input = gr.File(label="Upload your model")
            name_input = gr.Text(label="Your name or team")
            submit_btn = gr.Button("Submit model")

    error_box = gr.Textbox(label="Output log", visible=True)

    leaderboard_table = gr.Dataframe(
        label="Leaderboard",
        interactive=False,
        wrap=True
    )

    submit_btn.click(
        fn=evaluate_model,
        inputs=[file_input, name_input],
        outputs=[error_box, leaderboard_table]
    )

    # Refresh leaderboard on app launch
    demo.load(fn=load_leaderboard, inputs=[], outputs=[leaderboard_table])

demo.launch()