src/task_groups.py

"""Task groupings for filtering the leaderboard by subsets."""

# All tasks in the benchmark (100 tasks)
ALL_TASKS = [
    "ETT_15T", "ETT_1D", "ETT_1H", "ETT_1W",
    "LOOP_SEATTLE_1D", "LOOP_SEATTLE_1H", "LOOP_SEATTLE_5T",
    "M_DENSE_1D", "M_DENSE_1H",
    "SZ_TAXI_15T", "SZ_TAXI_1H",
    "australian_tourism",
    "bizitobs_l2c_1H", "bizitobs_l2c_5T",
    "boomlet_1062", "boomlet_1209", "boomlet_1225", "boomlet_1230", "boomlet_1282",
    "boomlet_1487", "boomlet_1631", "boomlet_1676", "boomlet_1855", "boomlet_1975",
    "boomlet_2187", "boomlet_285", "boomlet_619", "boomlet_772", "boomlet_963",
    "ecdc_ili",
    "entsoe_15T", "entsoe_1H", "entsoe_30T",
    "epf_be", "epf_de", "epf_fr", "epf_np", "epf_pjm",
    "ercot_1D", "ercot_1H", "ercot_1M", "ercot_1W",
    "favorita_stores_1D", "favorita_stores_1M", "favorita_stores_1W",
    "favorita_transactions_1D", "favorita_transactions_1M", "favorita_transactions_1W",
    "fred_md_2025/cee", "fred_md_2025/macro",
    "fred_qd_2025/cee", "fred_qd_2025/macro",
    "gvar",
    "hermes",
    "hierarchical_sales_1D", "hierarchical_sales_1W",
    "hospital", "hospital_admissions_1D", "hospital_admissions_1W",
    "jena_weather_10T", "jena_weather_1D", "jena_weather_1H",
    "kdd_cup_2022_10T", "kdd_cup_2022_1D", "kdd_cup_2022_30T",
    "m5_1D", "m5_1M", "m5_1W",
    "proenfo_gfc12", "proenfo_gfc14", "proenfo_gfc17",
    "redset_15T", "redset_1H", "redset_5T",
    "restaurant",
    "rohlik_orders_1D", "rohlik_orders_1W", "rohlik_sales_1D", "rohlik_sales_1W",
    "rossmann_1D", "rossmann_1W",
    "solar_1D", "solar_1W", "solar_with_weather_15T", "solar_with_weather_1H",
    "uci_air_quality_1D", "uci_air_quality_1H",
    "uk_covid_nation_1D/cumulative", "uk_covid_nation_1D/new",
    "uk_covid_nation_1W/cumulative", "uk_covid_nation_1W/new",
    "uk_covid_utla_1D/new", "uk_covid_utla_1W/cumulative",
    "us_consumption_1M", "us_consumption_1Q", "us_consumption_1Y",
    "walmart",
    "world_co2_emissions", "world_life_expectancy", "world_tourism",
]

# Mini benchmark - representative subset (20 tasks)
MINI_TASKS = [
    "jena_weather_1H",
    "M_DENSE_1D",
    "bizitobs_l2c_5T",
    "rohlik_orders_1D",
    "boomlet_1282",
    "rossmann_1D",
    "rossmann_1W",
    "boomlet_1676",
    "solar_with_weather_1H",
    "boomlet_619",
    "uci_air_quality_1H",
    "uk_covid_nation_1D/cumulative",
    "us_consumption_1Y",
    "epf_np",
    "world_co2_emissions",
    "ETT_15T",
    "ETT_1H",
    "proenfo_gfc14",
    "hospital_admissions_1D",
    "hospital_admissions_1W",
]

# Frequency-based groupings
FREQUENCY_GROUPS = {
    "Sub-hourly": [
        # T (1 minute)
        "boomlet_1225", "boomlet_1282", "boomlet_285", "boomlet_619", "boomlet_772", "boomlet_963",
        # 5T (5 minutes)
        "LOOP_SEATTLE_5T", "bizitobs_l2c_5T", "redset_5T",
        "boomlet_1062", "boomlet_1209", "boomlet_1230", "boomlet_1487",
        # 10T (10 minutes)
        "jena_weather_10T", "kdd_cup_2022_10T",
        # 15T (15 minutes)
        "ETT_15T", "SZ_TAXI_15T", "entsoe_15T", "redset_15T", "solar_with_weather_15T",
        # 30T (30 minutes)
        "entsoe_30T", "kdd_cup_2022_30T", "boomlet_1631", "boomlet_1676",
    ],
    "Hourly": [
        "ETT_1H", "LOOP_SEATTLE_1H", "M_DENSE_1H", "SZ_TAXI_1H",
        "bizitobs_l2c_1H", "entsoe_1H", "ercot_1H",
        "epf_be", "epf_de", "epf_fr", "epf_np", "epf_pjm",
        "jena_weather_1H",
        "proenfo_gfc12", "proenfo_gfc14", "proenfo_gfc17",
        "redset_1H", "solar_with_weather_1H", "uci_air_quality_1H",
        "boomlet_1855", "boomlet_1975", "boomlet_2187",
    ],
    "Daily": [
        "ETT_1D", "LOOP_SEATTLE_1D", "M_DENSE_1D",
        "ercot_1D", "kdd_cup_2022_1D", "solar_1D",
        "favorita_stores_1D", "favorita_transactions_1D",
        "hierarchical_sales_1D", "m5_1D",
        "restaurant",
        "rohlik_orders_1D", "rohlik_sales_1D", "rossmann_1D",
        "jena_weather_1D", "uci_air_quality_1D",
        "hospital_admissions_1D",
        "uk_covid_nation_1D/cumulative", "uk_covid_nation_1D/new", "uk_covid_utla_1D/new",
    ],
    "Weekly": [
        "ETT_1W", "ercot_1W", "solar_1W",
        "favorita_stores_1W", "favorita_transactions_1W",
        "hierarchical_sales_1W", "m5_1W",
        "hermes", "walmart",
        "rohlik_orders_1W", "rohlik_sales_1W", "rossmann_1W",
        "ecdc_ili",
        "hospital_admissions_1W",
        "uk_covid_nation_1W/cumulative", "uk_covid_nation_1W/new", "uk_covid_utla_1W/cumulative",
    ],
    "Monthly+": [
        # Monthly
        "ercot_1M",
        "favorita_stores_1M", "favorita_transactions_1M", "m5_1M",
        "fred_md_2025/cee", "fred_md_2025/macro",
        "hospital",
        "us_consumption_1M",
        # Quarterly
        "australian_tourism", "gvar",
        "fred_qd_2025/cee", "fred_qd_2025/macro",
        "us_consumption_1Q",
        # Yearly
        "us_consumption_1Y",
        "world_co2_emissions", "world_life_expectancy", "world_tourism",
    ],
}

# Domain-based groupings
DOMAIN_GROUPS = {
    "Energy": [
        "ETT_15T", "ETT_1D", "ETT_1H", "ETT_1W",
        "entsoe_15T", "entsoe_1H", "entsoe_30T",
        "epf_be", "epf_de", "epf_fr", "epf_np", "epf_pjm",
        "ercot_1D", "ercot_1H", "ercot_1M", "ercot_1W",
        "kdd_cup_2022_10T", "kdd_cup_2022_1D", "kdd_cup_2022_30T",
        "proenfo_gfc12", "proenfo_gfc14", "proenfo_gfc17",
        "solar_1D", "solar_1W", "solar_with_weather_15T", "solar_with_weather_1H",
    ],
    "Retail": [
        "favorita_stores_1D", "favorita_stores_1M", "favorita_stores_1W",
        "favorita_transactions_1D", "favorita_transactions_1M", "favorita_transactions_1W",
        "hermes",
        "hierarchical_sales_1D", "hierarchical_sales_1W",
        "m5_1D", "m5_1M", "m5_1W",
        "restaurant",
        "rohlik_orders_1D", "rohlik_orders_1W", "rohlik_sales_1D", "rohlik_sales_1W",
        "rossmann_1D", "rossmann_1W",
        "walmart",
    ],
    "Nature": [
        "jena_weather_10T", "jena_weather_1D", "jena_weather_1H",
        "uci_air_quality_1D", "uci_air_quality_1H",
    ],
    "Cloud": [
        "bizitobs_l2c_1H", "bizitobs_l2c_5T",
        "boomlet_1062", "boomlet_1209", "boomlet_1225", "boomlet_1230", "boomlet_1282",
        "boomlet_1487", "boomlet_1631", "boomlet_1676", "boomlet_1855", "boomlet_1975",
        "boomlet_2187", "boomlet_285", "boomlet_619", "boomlet_772", "boomlet_963",
        "redset_15T", "redset_1H", "redset_5T",
    ],
    "Health": [
        "ecdc_ili",
        "hospital", "hospital_admissions_1D", "hospital_admissions_1W",
        "uk_covid_nation_1D/cumulative", "uk_covid_nation_1D/new",
        "uk_covid_nation_1W/cumulative", "uk_covid_nation_1W/new",
        "uk_covid_utla_1D/new", "uk_covid_utla_1W/cumulative",
    ],
    "Econ": [
        "australian_tourism",
        "fred_md_2025/cee", "fred_md_2025/macro",
        "fred_qd_2025/cee", "fred_qd_2025/macro",
        "gvar",
        "us_consumption_1M", "us_consumption_1Q", "us_consumption_1Y",
        "world_co2_emissions", "world_life_expectancy", "world_tourism",
    ],
    "Mobility": [
        "LOOP_SEATTLE_1D", "LOOP_SEATTLE_1H", "LOOP_SEATTLE_5T",
        "M_DENSE_1D", "M_DENSE_1H",
        "SZ_TAXI_15T", "SZ_TAXI_1H",
    ],
}


def get_task_group(group_type: str, group_value: str | None = None) -> list[str]:
    """Get the list of tasks for a given group type and value.

    Args:
        group_type: One of "full", "mini", "frequency", "domain"
        group_value: Required for "frequency" and "domain" types

    Returns:
        List of task names belonging to the group
    """
    if group_type == "full":
        return ALL_TASKS
    elif group_type == "mini":
        return MINI_TASKS
    elif group_type == "frequency":
        if group_value is None:
            raise ValueError("group_value required for frequency grouping")
        return FREQUENCY_GROUPS[group_value]
    elif group_type == "domain":
        if group_value is None:
            raise ValueError("group_value required for domain grouping")
        return DOMAIN_GROUPS[group_value]
    else:
        raise ValueError(f"Unknown group_type: {group_type}")