import gradio as gr
import pandas as pd
import json
from datetime import datetime
import os

# Define weights for composite score matching leaderboard.py
WEIGHTS = {
    "em": 0.20,
    "f1": 0.15,
    "rouge_l": 0.15,
    "bert_score": 0.15,
    "recall_at_1": 0.10,
    "mrr": 0.05,
    "ndcg": 0.05,
    "hallucination_inv": 0.05,
    "citation_accuracy": 0.05,
    "calibration_inv": 0.05,
}

LEADERBOARD_FILE = "leaderboard_state.json"

def _nan_safe(v, fallback=0.0):
    if v is None:
        return fallback
    try:
        f = float(v)
        import math
        return fallback if math.isnan(f) else f
    except (TypeError, ValueError):
        return fallback

def compute_composite(overall):
    em = _nan_safe(overall.get("em"))
    f1 = _nan_safe(overall.get("f1"))
    rouge = _nan_safe(overall.get("rouge_l"))
    bert = _nan_safe(overall.get("bert_score"))
    r1 = _nan_safe(overall.get("recall_at_1"))
    mrr_v = _nan_safe(overall.get("mrr"))
    ndcg_v = _nan_safe(overall.get("ndcg"))
    halluc = _nan_safe(overall.get("hallucination"), fallback=0.0)
    cit_acc = _nan_safe(overall.get("citation_accuracy"))
    calib = _nan_safe(overall.get("calibration_error"), fallback=0.0)

    score = (
        WEIGHTS["em"] * em
        + WEIGHTS["f1"] * f1
        + WEIGHTS["rouge_l"] * rouge
        + WEIGHTS["bert_score"] * bert
        + WEIGHTS["recall_at_1"] * r1
        + WEIGHTS["mrr"] * mrr_v
        + WEIGHTS["ndcg"] * ndcg_v
        + WEIGHTS["hallucination_inv"] * max(0.0, 100.0 - halluc)
        + WEIGHTS["citation_accuracy"] * cit_acc
        + WEIGHTS["calibration_inv"] * max(0.0, 100.0 - calib)
    )
    return round(score, 3)

def load_leaderboard():
    if os.path.exists(LEADERBOARD_FILE):
        with open(LEADERBOARD_FILE, "r") as f:
            return json.load(f)
    return []

def save_leaderboard(entries):
    with open(LEADERBOARD_FILE, "w") as f:
        json.dump(entries, f, indent=2)

def get_leaderboard_df():
    entries = load_leaderboard()
    if not entries:
        return pd.DataFrame(columns=["Rank", "Model", "Composite Score", "EM", "F1", "ROUGE-L", "BERTScore", "Recall@1", "Hallucination", "Citation Acc."])
    
    # Sort by composite score
    entries.sort(key=lambda x: x.get("composite", 0), reverse=True)
    
    data = []
    for rank, e in enumerate(entries, 1):
        ov = e.get("overall", {})
        data.append({
            "Rank": rank,
            "Model": e["model_name"],
            "Composite Score": f"{e.get('composite', 0):.2f}",
            "EM": f"{_nan_safe(ov.get('em')):.2f}",
            "F1": f"{_nan_safe(ov.get('f1')):.2f}",
            "ROUGE-L": f"{_nan_safe(ov.get('rouge_l')):.2f}",
            "BERTScore": f"{_nan_safe(ov.get('bert_score')):.2f}",
            "Recall@1": f"{_nan_safe(ov.get('recall_at_1')):.2f}",
            "Hallucination": f"{_nan_safe(ov.get('hallucination')):.2f}",
            "Citation Acc.": f"{_nan_safe(ov.get('citation_accuracy')):.2f}"
        })
    
    return pd.DataFrame(data)

def submit_model(model_name, report_file):
    if not model_name:
        return "Error: Model name is required.", get_leaderboard_df()
    if report_file is None:
        return "Error: Please upload a JSON report file.", get_leaderboard_df()
    
    try:
        content = json.loads(report_file.decode("utf-8") if isinstance(report_file, bytes) else open(report_file.name).read())
    except Exception as e:
        return f"Error reading JSON: {str(e)}", get_leaderboard_df()
        
    overall = content.get("overall", {})
    composite = compute_composite(overall)
    
    entry = {
        "model_name": model_name,
        "submitted_at": datetime.now().isoformat(),
        "composite": composite,
        "overall": overall,
        "by_task": content.get("by_task", {}),
        "by_language": content.get("by_language", {}),
        "by_subset": content.get("by_subset", {})
    }
    
    entries = load_leaderboard()
    # Remove existing entry for same model
    entries = [e for e in entries if e["model_name"] != model_name]
    entries.append(entry)
    save_leaderboard(entries)
    
    return f"Successfully added {model_name} with score {composite:.2f}!", get_leaderboard_df()

# Build Gradio UI
with gr.Blocks(theme=gr.themes.Soft(primary_hue="blue", neutral_hue="slate")) as app:
    gr.Markdown("# 🏛️ IndicConBench Official Leaderboard")
    gr.Markdown("Welcome to the **IndicConBench** Leaderboard, evaluating LLMs on Multilingual Constitutional Reasoning for India. Models are ranked based on a composite score combining reasoning, factual accuracy, retrieval, and hallucination resistance.")
    
    with gr.Tabs():
        with gr.TabItem("Leaderboard"):
            lb_df = get_leaderboard_df()
            leaderboard_table = gr.Dataframe(value=lb_df, interactive=False)
            
            refresh_btn = gr.Button("🔄 Refresh Leaderboard")
            refresh_btn.click(fn=get_leaderboard_df, inputs=[], outputs=[leaderboard_table])
            
        with gr.TabItem("Submit Model"):
            gr.Markdown("### Submit Your Evaluation Report")
            gr.Markdown("Upload the `report.json` generated by `evaluate.py` to add your model to the leaderboard.")
            
            with gr.Row():
                model_name_input = gr.Textbox(label="Model Name", placeholder="e.g., Llama-3-8B-Instruct")
                file_upload = gr.File(label="Upload evaluation report (.json)", file_types=[".json"])
            
            submit_btn = gr.Button("Submit to Leaderboard", variant="primary")
            status_out = gr.Markdown()
            
            submit_btn.click(
                fn=submit_model,
                inputs=[model_name_input, file_upload],
                outputs=[status_out, leaderboard_table]
            )

if __name__ == "__main__":
    app.launch()