import argparse
import ast
import pickle
import os
import threading
import time

import gradio as gr
import numpy as np
import pandas as pd
from serve.model import model_config

def amend_model_name(config):
    if config.page_link:
        return '[' + config.model_name + '](' + config.page_link + ')'
    elif config.code_link:
        return '[' + config.model_name + '](' + config.page_link + ')'
    else:
        return config.model_name

def amend_rank(rank):
    if rank==1:
        return "🥇"
    elif rank==2:
        return "🥈"
    elif rank==3:
        return '🥉'
    else:
        return rank

def get_leaderboard_values(leaderboard_df):
    leaderboard_vals = []
    for i, row in leaderboard_df.iterrows():
        rank = i+1
        model_name = row["Method"]
        task = row["Task"]
        if model_name not in model_config.keys() or model_config[model_name].task != task:
            continue
        
        cfg = model_config[model_name]
        values = [amend_rank(rank), amend_model_name(cfg), task]
        values = values + [row.get(dim, np.NaN) for dim in leaderboard_df.columns[2:]]
        # values.append(round(np.sum([v for v in values[3:] if pd.notna(v)]), 4))
        values.append(cfg.organization if cfg.organization else "N/A")
        
        leaderboard_vals.append(values)
    return leaderboard_vals

def get_topk_ranks(df, k=3):
    ranks = {}
    for col_idx, col in enumerate(df.columns[2:]):  # skip "Model" “Task"
        topk = df[col].nlargest(k)
        for rank, idx in enumerate(topk.index):
            if idx not in ranks:
                ranks[idx] = {}
            ranks[idx][col_idx] = rank + 1  # 1-based rank
    for i in range(k): ranks[i][5] = i + 1
    return ranks  # dict: row -> {col: rank}

def build_leaderboard_tab(leaderboard_file: str, task: str = ""):
    if not isinstance(leaderboard_file, str):
        leaderboard_file = leaderboard_file.value
    if not isinstance(task, str):
        task = task.value

    df = pd.read_csv(leaderboard_file)
    if task in ["Text-to-3D only", "Image-to-3D only"]:
        df = df[df["Task"] == task.split()[0]]
        # df = df.drop(df[df["Task"]!=task.split()[0]].index)
    leaderboard_df = df.drop(df[df["Method"].isnull()].index)
    leaderboard_df = leaderboard_df.reset_index(drop=True)
    leaderboard_df = leaderboard_df.sort_values(by="Overall", ascending=False).reset_index(drop=True)

    leaderboard_vals = get_leaderboard_values(leaderboard_df)
    leaderboard = gr.Dataframe(
        headers = ['Rank', "🤖 Model", "🪧 Task" ] 
            + [f"{dim}" for dim in leaderboard_df.keys()[2:-1]] 
            + ["⭐ Overall", "🏛️ Orgnization"],     # "🔗 Links",
        datatype = ["number", "markdown", "markdown"] 
            + ["number"] * (len(leaderboard_df.columns) - 3) 
            + ["number", "markdown"],
        value = leaderboard_vals,
        height = 680,
        column_widths = [60, 120, 100] 
            + [120] * (len(leaderboard_df.columns) - 3) 
            + [120, 180],
        wrap = True,
        type="pandas",
    )
    return leaderboard