import gradio as gr
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
from io import BytesIO
from PIL import Image
from datasets.exceptions import DatasetNotFoundError

from src.dataloading import get_leaderboard_datasets
from src.similarity import load_data_and_compute_similarities

# Set matplotlib backend for non-GUI environments
plt.switch_backend('Agg')

def create_heatmap(selected_models, selected_dataset, selected_metric):
    if not selected_models or not selected_dataset:
        return None
    
    # Sort models and get short names
    similarities = load_data_and_compute_similarities(selected_models, selected_dataset, selected_metric)

    # Check if similarity matrix contains NaN rows
    failed_models = []
    for i in range(len(similarities)):
        if np.isnan(similarities[i]).all():
            failed_models.append(selected_models[i])

    if failed_models:
        gr.Warning(f"Failed to load data for models: {', '.join(failed_models)}")

    # Create figure and heatmap using seaborn
    plt.figure(figsize=(8, 6))
    ax = sns.heatmap(
        similarities,
        annot=True,
        fmt=".2f",
        cmap="viridis",
        vmin=0,
        vmax=1,
        xticklabels=selected_models,
        yticklabels=selected_models
    )
    
    # Customize plot
    plt.title(f"{selected_metric} for {selected_dataset}", fontsize=16)
    plt.xlabel("Models", fontsize=14)
    plt.ylabel("Models", fontsize=14)
    plt.xticks(rotation=45, ha='right')
    plt.yticks(rotation=0)
    plt.tight_layout()

    # Save to buffer
    buf = BytesIO()
    plt.savefig(buf, format="png", dpi=100, bbox_inches="tight")
    plt.close()
    
    # Convert to PIL Image
    buf.seek(0)
    img = Image.open(buf).convert("RGB")
    return img

def validate_inputs(selected_models, selected_dataset):
    if not selected_models:
        raise gr.Error("Please select at least one model!")
    if not selected_dataset:
        raise gr.Error("Please select a dataset!")
    

def update_datasets_based_on_models(selected_models, current_dataset):
    try:
        available_datasets = get_leaderboard_datasets(selected_models) if selected_models else []
        if current_dataset in available_datasets:
            valid_dataset = current_dataset
        elif "mmlu_pro" in available_datasets:
            valid_dataset = "mmlu_pro"
        else:
            valid_dataset = None
        return gr.update(
            choices=available_datasets,
            value=valid_dataset
        )
    except DatasetNotFoundError as e:
        # Extract model name from error message
        model_name = e.args[0].split("'")[1]
        model_name = model_name.split("/")[-1].replace("__", "/").replace("_details", "")
        
        # Display a shorter warning
        gr.Warning(f"Data for '{model_name}' is gated or unavailable.")
        return gr.update(choices=[], value=None)
    
custom_css = """
.image-container img {
    width: 80% !important;  /* Make it 80% of the parent container */
    height: auto !important; /* Maintain aspect ratio */
    max-width: 800px; /* Optional: Set a max limit */
    display: block;
    margin: auto; /* Center the image */
}
"""