import pandas as pd
import gradio as gr
from collections import OrderedDict
import logging
import tempfile
import os
from huggingface_hub import (
    HfApi,
    hf_hub_download,
    get_safetensors_metadata,
    metadata_load,
)

from utils.misc import human_format, make_clickable_model

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)


EXCLUDED_MODELS = []  # For models that misbehave :)

K_EVALUATIONS = [1, 5, 10, 20, 50]
DIST_EVALUATIONS = [1_000_000, 500_000, 100_000, 10_000]

EXPECTED_KEY_TO_COLNAME = OrderedDict(
    [
        ("rank", "Rank"),  # Just for columns order
        ("model", "Model"),  # Just for columns order
        ("model_size", "Model Size (Million)"),  # Just for columns order
        ("conditioning", "Conditioning"),
        ("embedding_dim", "Embedding Dimension"),
    ]
    + [
        (f"recall_at_{K}|{D}", f"R@{K} +{human_format(D)} Dist.")
        for D in DIST_EVALUATIONS
        for K in K_EVALUATIONS
    ]
    + [
        ("n_dists", "Available Dists"),
    ],
)


def get_safetensors_nparams(modelId):
    try:
        safetensors = get_safetensors_metadata(modelId)
        num_parameters = sum(safetensors.parameter_count.values())
        return round(num_parameters / 1e6)
    except Exception:
        pass


def parse_model(m):
    readme_path = hf_hub_download(m.modelId, filename="README.md")
    meta = metadata_load(readme_path)

    if "model-index" not in meta:
        raise ValueError("Missing `model-index` in metadata")

    for result in meta["model-index"][0]["results"]:
        if result["dataset"]["type"] == "Slep/LAION-RVS-Fashion":
            break  # Found the right dataset

    # Get data from model-index / safetensors metadata
    d = {
        EXPECTED_KEY_TO_COLNAME["model"]: make_clickable_model(m.modelId),
        EXPECTED_KEY_TO_COLNAME["model_size"]: get_safetensors_nparams(m.modelId),
    }

    # Get data from exported results
    for metric in result["metrics"]:
        t = metric["type"]

        if t in EXPECTED_KEY_TO_COLNAME:
            d[EXPECTED_KEY_TO_COLNAME[t]] = metric["value"]

    return d


def get_data_from_hub():
    api = HfApi()
    models = api.list_models(filter="lrvsf-benchmark")

    df_list = []
    for m in models:
        if m.modelId in EXCLUDED_MODELS:
            continue

        try:
            parsed = parse_model(m)
            if parsed:
                df_list.append(parsed)
        except Exception as e:
            logging.warning(f"Failed to parse model {m.modelId} : {e}")

    return pd.DataFrame(df_list, columns=EXPECTED_KEY_TO_COLNAME.values())


def filter_dataframe(df, k_filter, d_filter, c_filter):
    # ===== FILTER COLUMNS
    # Fixed column positions
    selected_columns = [
        EXPECTED_KEY_TO_COLNAME["rank"],
        EXPECTED_KEY_TO_COLNAME["model"],
        EXPECTED_KEY_TO_COLNAME["conditioning"],
        EXPECTED_KEY_TO_COLNAME["model_size"],
        EXPECTED_KEY_TO_COLNAME["embedding_dim"],
    ]

    datatypes = ["number", "markdown", "number", "number"]

    for key, name in EXPECTED_KEY_TO_COLNAME.items():
        if name in selected_columns:
            # Already added, probably part of the initial columns
            continue

        if key.startswith("recall_at_"):
            # Process : recall_at_K|D -> recall_at_K , D -> K , D
            # Could be a regex... but simple enough
            recall_at_K, D = key.split("|")
            K = recall_at_K.split("_")[-1]

            if int(K) in k_filter and int(D) in d_filter:
                selected_columns.append(name)
                datatypes.append("str")  # Because of the ± std

    selected_columns.append(EXPECTED_KEY_TO_COLNAME["n_dists"])
    datatypes.append("number")

    df = df[selected_columns]

    # ===== FILTER ROWS
    if c_filter != "all":
        df = df[df[EXPECTED_KEY_TO_COLNAME["conditioning"]] == c_filter]

    return df[selected_columns], datatypes


def add_rank(df):
    main_metrics = df["R@1 +1M Dist."].str.split("±").str[0].astype(float)

    # Argsort is from smallest to largest so we reverse it
    df["Rank"] = df.shape[0] - main_metrics.argsort()
    return df


def save_current_leaderboard(df):
    filename = tempfile.NamedTemporaryFile(
        prefix="lrvsf_export_", suffix=".csv", delete=False
    ).name
    df.to_csv(filename, index=False)
    return filename


def load_lrvsf_models(k_filter, d_filter, c_filter, csv_file):
    # Remove previous tmpfile
    if csv_file:
        os.remove(csv_file)

    df = get_data_from_hub()
    df = add_rank(df)
    df, datatypes = filter_dataframe(df, k_filter, d_filter, c_filter)
    df = df.sort_values(by="Rank")

    filename = save_current_leaderboard(df)

    outputs = [
        gr.DataFrame(value=df, datatype=datatypes),
        gr.File(filename, label="CSV File"),
    ]

    return outputs


if __name__ == "__main__":
    with gr.Blocks() as demo:
        gr.Markdown(
            """
                # 👗 LAION - Referred Visual Search - Fashion : Leaderboard

                - To submit, refer to the [LAION-RVS-Fashion Benchmark repository](https://github.com/Simon-Lepage/LRVSF-Benchmark).
                - For details on the task and the dataset, refer to the [LRVSF paper](https://arxiv.org/abs/2306.02928).
                - To download the leaderboard as CSV, click on the file below the table.
            """
        )
        with gr.Row():
            k_filter = gr.CheckboxGroup(
                choices=K_EVALUATIONS, value=K_EVALUATIONS, label="Recall at K"
            )
            d_filter = gr.CheckboxGroup(
                choices=[(human_format(D), D) for D in DIST_EVALUATIONS],
                value=DIST_EVALUATIONS,
                label="Number of Distractors",
            )
            c_filter = gr.Radio(
                choices=["all", "category", "text"],
                value="all",
                label="Conditioning",
            )

        df_table = gr.Dataframe(type="pandas", interactive=False)
        csv_file = gr.File(interactive=False)
        refresh = gr.Button("Refresh")

        # Actions
        refresh.click(
            load_lrvsf_models,
            inputs=[k_filter, d_filter, c_filter, csv_file],
            outputs=[df_table, csv_file],
        )
        demo.load(
            load_lrvsf_models,
            inputs=[k_filter, d_filter, c_filter, csv_file],
            outputs=[df_table, csv_file],
        )

    demo.launch()