# data_loaders/load_datasets.py

from datasets import load_dataset

def load_all_datasets():
    datasets = {}
    try:
        # Linear Algebra
        datasets["linear_algebra_composed"] = load_dataset("marcov/math_dataset_algebra__linear_2d_composed_promptsource")
        datasets["linear_algebra_standard"] = load_dataset("marcov/math_dataset_algebra__linear_2d_promptsource")

        # Probability NLI variants
        datasets["probability_nli_1hop"] = load_dataset("sileod/probability_words_nli", "reasoning_1hop")
        datasets["probability_nli_2hop"] = load_dataset("sileod/probability_words_nli", "reasoning_2hop")
        datasets["probability_nli_usnli"] = load_dataset("sileod/probability_words_nli", "usnli")

        # Qwen Math Probability
        datasets["probability_qwen"] = load_dataset("yxsllgz-uts-org/Math_Consistency-Probability-Qwen2.5-Math-1.5B-Instruct-style1")

        # Supply Chain QA
        datasets["supply_chain_qa"] = load_dataset("sebdg/supply_chain_q_a")

        # Game Theory
        datasets["game_theory"] = load_dataset("ebowwa/game-theory-io")

        # Time Series - Stock Prices
        datasets["time_series_stocks"] = load_dataset("chuyin0321/timeseries-daily-stocks")

    except Exception as e:
        print(f"❌ Error loading one or more datasets: {e}")

    print("📊 Loaded datasets:")
    for key in datasets.keys():
        print(f" - {key}")

    return datasets