from app import add_new_eval, LeaderboardData
from pathlib import Path
import gradio as gr
import os

# Create a mock profile for testing
class MockProfile:
    def __init__(self, username):
        self.username = username

mock_profile = MockProfile("mshuaibi_test")

evals = {
    # "IE_EA": "unoptimized_ie_ea_results.json",
    # "Ligand pocket": "pdb_pocket_results.json",
    "Ligand strain": "ligand_strain_results.json",
    # "Conformers": "geom_conformers_results.json",
    # "Protonation": "protonation_energies_results.json",
    # "Distance scaling": "distance_scaling_results.json",
    # "Spin gap": "unoptimized_spin_gap_results.json",
    # "Validation": "val_predictions.npz",
    # "Test": "test_predictions.npz"
}

models = {
    # "esen-s-c-4M": {
    #     "name": "eSEN-sm-cons.",
    #     "dataset_size": "OMol-4M",
    #     "results_dir": "/large_experiments/opencatalyst/foundation_models/data/omol/leaderboard/predictions/043025_esen_sm_conserving_4M",
    #     "paper_link": "https://arxiv.org/pdf/2505.08762",
    # },
    # "esen-s-c-All": {
    #     "name": "eSEN-sm-cons.",
    #     "dataset_size": "OMol-All",
    #     "results_dir": "/large_experiments/opencatalyst/foundation_models/data/omol/leaderboard/predictions/043025_esen_sm_conserving_all",
    #     "paper_link": "https://arxiv.org/pdf/2505.08762",
    # },
    # "esen-m-d-4M": {
    #     "name": "eSEN-md-d.",
    #     "dataset_size": "OMol-4M",
    #     "results_dir": "/large_experiments/opencatalyst/foundation_models/data/omol/leaderboard/predictions/043025_esen_md_direct_4M_finetune",
    #     "paper_link": "https://arxiv.org/pdf/2505.08762",
    # },
    # "esen-m-d-All": {
    #     "name": "eSEN-md-d.",
    #     "dataset_size": "OMol-All",
    #     "results_dir": "/large_experiments/opencatalyst/foundation_models/data/omol/leaderboard/predictions/043025_esen_md_direct_all_finetune",
    #     "paper_link": "https://arxiv.org/pdf/2505.08762",
    # },
    # "goc-4M": {
    #     "name": "GemNet-OC",
    #     "dataset_size": "OMol-4M",
    #     "results_dir": "/large_experiments/opencatalyst/foundation_models/data/omol/leaderboard/predictions/043025_gemnet_oc_4M",
    #     "paper_link": "https://arxiv.org/pdf/2505.08762",
    # },
    # "goc-All": {
    #     "name": "GemNet-OC",
    #     "dataset_size": "OMol-All",
    #     "results_dir": "/large_experiments/opencatalyst/foundation_models/data/omol/leaderboard/predictions/050325_gemnet_oc_all",
    #     "paper_link": "https://arxiv.org/pdf/2505.08762",
    # },
    # "uma-s-1p1": {
    #     "name": "UMA-S-1p1",
    #     "dataset_size": "UMA-459M",
    #     "results_dir": "/large_experiments/opencatalyst/foundation_models/data/omol/leaderboard/predictions/uma_sm_1p1",
    #     "paper_link": "https://arxiv.org/pdf/2506.23971",
    # },
    # "uma-m-1p1": {
    #     "name": "UMA-M-1p1",
    #     "dataset_size": "UMA-459M",
    #     "results_dir": "/large_experiments/opencatalyst/foundation_models/data/omol/leaderboard/predictions/uma_md_1p1",
    #     "paper_link": "https://arxiv.org/pdf/2506.23971",
    # },
    "mace": {
        "name": "mace-omol-L-0",
        "dataset_size": "OMol-All",
        "results_dir": "/large_experiments/opencatalyst/foundation_models/data/omol/leaderboard/predictions/mace",
        "paper_link": "https://github.com/ACEsuit/mace/releases/tag/v0.3.14",
        "org": "MACE-Cambridge"
    },
}

for model, model_info in models.items():
    model_name = model_info["name"]
    dataset_size = model_info["dataset_size"]
    results_dir = model_info["results_dir"]
    paper_link = model_info["paper_link"]
    org = model_info.get("org", "Meta")

    for _eval, eval_path in evals.items():
        generator = add_new_eval(
            path_to_file=os.path.join(results_dir, eval_path),
            eval_type=_eval,
            organization=org,
            model=model_name,
            url=paper_link,
            mail="mshuaibi@meta.com",
            training_set=dataset_size,
            additional_info="",
            profile=mock_profile,
        )
        for i in generator:
            print(i)