retrieve.smk

# SPDX-FileCopyrightText: Contributors to Open-TYNDP <https://github.com/open-energy-transition/open-tyndp>
# SPDX-FileCopyrightText: Contributors to PyPSA-Eur <https://github.com/pypsa/pypsa-eur>
#
# SPDX-License-Identifier: MIT

import os
import requests
from datetime import datetime
from dateutil.relativedelta import relativedelta
from shutil import move, unpack_archive, rmtree, copy2
from zipfile import ZipFile


# Configure the default storage provider for accessing remote files using http
# and the special storage plugin for accessing Zenodo files
storage:
    provider="http",
    keep_local=True,
    retries=3,


storage cached_http:
    provider="cached-http",


ARCHIVE_SOURCES = {"archive", "tyndp-archive"}


if (EUROSTAT_BALANCES_DATASET := dataset_version("eurostat_balances"))["source"] in [
    "primary",
    *ARCHIVE_SOURCES,
]:

    rule retrieve_eurostat_balances:
        input:
            tsv_gz=storage(EUROSTAT_BALANCES_DATASET["url"]),
        output:
            tsv_gz=f"{EUROSTAT_BALANCES_DATASET['folder']}/estat_nrg_bal_c.tsv.gz",
        message:
            "Retrieving Eurostat balances data"
        run:
            copy2(input["tsv_gz"], output["tsv_gz"])


if (
    EUROSTAT_HOUSEHOLD_BALANCES_DATASET := dataset_version(
        "eurostat_household_balances"
    )
)["source"] in [
    "primary",
    *ARCHIVE_SOURCES,
]:

    rule retrieve_eurostat_household_balances:
        input:
            csv=storage(EUROSTAT_HOUSEHOLD_BALANCES_DATASET["url"]),
        output:
            csv=f"{EUROSTAT_HOUSEHOLD_BALANCES_DATASET['folder']}/nrg_d_hhq.csv",
        message:
            "Retrieving Eurostat household balances data"
        run:
            copy2(input["csv"], output["csv"])


if (SWISS_ENERGY_BALANCES_DATASET := dataset_version("swiss_energy_balances"))[
    "source"
] in [
    "archive",
    "primary",
]:

    rule retrieve_swiss_energy_balances:
        input:
            xlsx=storage(SWISS_ENERGY_BALANCES_DATASET["url"]),
        output:
            xlsx=f"{SWISS_ENERGY_BALANCES_DATASET['folder']}/12361-VWZ_Webtabellen_2024.xlsx",
        message:
            "Retrieving Swiss energy balances data"
        run:
            copy2(input["xlsx"], output["xlsx"])


if (NUTS3_POPULATION_DATASET := dataset_version("nuts3_population"))["source"] in [
    "primary",
    *ARCHIVE_SOURCES,
]:

    rule retrieve_nuts3_population:
        input:
            gz=storage(NUTS3_POPULATION_DATASET["url"]),
        output:
            gz=f"{NUTS3_POPULATION_DATASET['folder']}/nama_10r_3popgdp.tsv.gz",
        retries: 2
        message:
            "Retrieving NUTS3 population data"
        run:
            copy2(input["gz"], output["gz"])


if (CORINE_DATASET := dataset_version("corine"))["source"] in ARCHIVE_SOURCES:

    rule retrieve_corine:
        input:
            zip_file=storage(CORINE_DATASET["url"]),
        output:
            zip_file=f"{CORINE_DATASET['folder']}/corine.zip",
            tif_file=f"{CORINE_DATASET['folder']}/corine.tif",
        message:
            "Retrieving Corine land cover data"
        run:
            output_folder = Path(output["zip_file"]).parent
            unpack_archive(input["zip_file"], output_folder)
            copy2(input["zip_file"], output["zip_file"])
            copy2(
                f"{output_folder}/corine/g250_clc06_V18_5.tif", output["tif_file"]
            )

elif (CORINE_DATASET := dataset_version("corine"))["source"] in ["primary"]:

    rule retrieve_corine:
        output:
            zip=f"{CORINE_DATASET['folder']}/corine.zip",
            tif_file=f"{CORINE_DATASET['folder']}/corine.tif",
        log:
            logs("retrieve_corine_primary.log"),
        retries: 2
        resources:
            mem_mb=1000,
        params:
            apikey=os.environ.get("CORINE_API_TOKEN", ""),
        message:
            "Retrieving Corine land cover data"
        script:
            scripts("retrieve_corine_dataset_primary.py")


if (H2_SALT_CAVERNS_DATASET := dataset_version("h2_salt_caverns"))[
    "source"
] in ARCHIVE_SOURCES:

    rule retrieve_h2_salt_caverns:
        input:
            geojson=storage(H2_SALT_CAVERNS_DATASET["url"]),
        output:
            geojson=f"{H2_SALT_CAVERNS_DATASET['folder']}/h2_salt_caverns_GWh_per_sqkm.geojson",
        retries: 2
        message:
            "Retrieving H2 salt caverns data"
        run:
            copy2(input["geojson"], output["geojson"])


if (GDP_PER_CAPITA_DATASET := dataset_version("gdp_per_capita"))[
    "source"
] in ARCHIVE_SOURCES:

    rule retrieve_gdp_per_capita:
        input:
            gdp=storage(GDP_PER_CAPITA_DATASET["url"]),
        output:
            gdp=f"{GDP_PER_CAPITA_DATASET['folder']}/GDP_per_capita_PPP_1990_2015_v2.nc",
        retries: 2
        message:
            "Retrieving GDP per capita data"
        run:
            copy2(input["gdp"], output["gdp"])


if (POPULATION_COUNT_DATASET := dataset_version("population_count"))["source"] in [
    "primary",
    *ARCHIVE_SOURCES,
]:

    rule retrieve_population_count:
        input:
            tif=storage(POPULATION_COUNT_DATASET["url"]),
        output:
            tif=f"{POPULATION_COUNT_DATASET['folder']}/ppp_2019_1km_Aggregated.tif",
        retries: 2
        message:
            "Retrieving population count data"
        run:
            copy2(input["tif"], output["tif"])

            if POPULATION_COUNT_DATASET["source"] == "primary":
                import xarray as xr
                import rioxarray as rio

                file_path = output["tif"]
                ds = xr.open_dataarray(file_path)
                ds_reqd = ds.sel(x=slice(15.55, 40.41), y=slice(52.49, 41.72))
                ds_reqd.rio.to_raster(file_path)



if (GHG_EMISSIONS_DATASET := dataset_version("ghg_emissions"))["source"] in [
    "primary",
    *ARCHIVE_SOURCES,
]:

    rule retrieve_ghg_emissions:
        input:
            ghg=storage(GHG_EMISSIONS_DATASET["url"]),
        output:
            csv=f"{GHG_EMISSIONS_DATASET['folder']}/UNFCCC_v23.csv",
            zip=(
                f"{GHG_EMISSIONS_DATASET['folder']}/UNFCCC_v23.csv.zip"
                if GHG_EMISSIONS_DATASET["source"] == "primary"
                else []
            ),
            directory=(
                directory(GHG_EMISSIONS_DATASET["folder"])
                if GHG_EMISSIONS_DATASET["source"] == "primary"
                else []
            ),
        retries: 2
        message:
            "Retrieving GHG emissions data"
        run:
            if GHG_EMISSIONS_DATASET["source"] == "primary":
                copy2(input["ghg"], output["zip"])
                unpack_archive(output["zip"], GHG_EMISSIONS_DATASET["folder"])
            else:
                copy2(input["ghg"], output["csv"])



if (GEBCO_DATASET := dataset_version("gebco"))["source"] in [
    "primary",
    *ARCHIVE_SOURCES,
]:

    rule retrieve_gebco:
        input:
            storage(GEBCO_DATASET["url"]),
        output:
            gebco=f"{GEBCO_DATASET['folder']}/GEBCO_2014_2D.nc",
            zip_file=(
                f"{GEBCO_DATASET['folder']}/GEBCO_2014.zip"
                if GEBCO_DATASET["source"] == "primary"
                else []
            ),
        message:
            "Retrieving GEBCO bathymetry data"
        run:
            if GEBCO_DATASET["source"] == "primary":
                import xarray as xr

                copy2(input[0], output["zip_file"])

                output_folder = Path(output["zip_file"]).parent
                unpack_archive(output["zip_file"], output_folder)

                # Limit extent to Europe to reduce file size
                ds = xr.open_dataset(output["gebco"])
                ds = ds.sel(lat=slice(32, 73), lon=slice(-21, 45))
                ds.to_netcdf(output["gebco"])
            else:
                copy2(input[0], output["gebco"])



if (ATTRIBUTED_PORTS_DATASET := dataset_version("attributed_ports"))["source"] in [
    "primary",
    *ARCHIVE_SOURCES,
]:

    rule retrieve_attributed_ports:
        input:
            json=storage(ATTRIBUTED_PORTS_DATASET["url"]),
        output:
            json=f"{ATTRIBUTED_PORTS_DATASET['folder']}/attributed_ports.json",
        retries: 2
        message:
            "Retrieving attributed ports data"
        run:
            copy2(input["json"], output["json"])


if (JRC_IDEES_DATASET := dataset_version("jrc_idees"))["source"] in [
    "primary",
    *ARCHIVE_SOURCES,
]:

    rule retrieve_jrc_idees:
        input:
            zip_file=storage(JRC_IDEES_DATASET["url"]),
        output:
            zip_file=f"{JRC_IDEES_DATASET['folder']}/jrc_idees.zip",
            directory=directory(JRC_IDEES_DATASET["folder"]),
        message:
            "Retrieving JRC IDEES data"
        run:
            copy2(input["zip_file"], output["zip_file"])
            output_folder = Path(output["zip_file"]).parent
            unpack_archive(output["zip_file"], output_folder)


if (EU_NUTS2013_DATASET := dataset_version("eu_nuts2013"))["source"] in [
    "primary",
    *ARCHIVE_SOURCES,
]:

    rule retrieve_eu_nuts_2013:
        input:
            shapes=storage(EU_NUTS2013_DATASET["url"]),
        output:
            zip_file=f"{EU_NUTS2013_DATASET['folder']}/ref-nuts-2013-03m.geojson.zip",
            folder=directory(
                f"{EU_NUTS2013_DATASET['folder']}/ref-nuts-2013-03m.geojson"
            ),
            shapes_level_3=f"{EU_NUTS2013_DATASET['folder']}/ref-nuts-2013-03m.geojson/NUTS_RG_03M_2013_4326_LEVL_3.geojson",
            shapes_level_2=f"{EU_NUTS2013_DATASET['folder']}/ref-nuts-2013-03m.geojson/NUTS_RG_03M_2013_4326_LEVL_2.geojson",
        message:
            "Retrieving EU NUTS 2013 data"
        run:
            copy2(input["shapes"], output["zip_file"])
            unpack_archive(output["zip_file"], Path(output.shapes_level_3).parent)


if (EU_NUTS2021_DATASET := dataset_version("eu_nuts2021"))["source"] in [
    "primary",
    *ARCHIVE_SOURCES,
]:

    rule retrieve_eu_nuts_2021:
        input:
            shapes=storage(EU_NUTS2021_DATASET["url"]),
        output:
            zip_file=f"{EU_NUTS2021_DATASET['folder']}/ref-nuts-2021-01m.geojson.zip",
            folder=directory(
                f"{EU_NUTS2021_DATASET['folder']}/ref-nuts-2021-01m.geojson"
            ),
            shapes_level_3=f"{EU_NUTS2021_DATASET['folder']}/ref-nuts-2021-01m.geojson/NUTS_RG_01M_2021_4326_LEVL_3.geojson",
            shapes_level_2=f"{EU_NUTS2021_DATASET['folder']}/ref-nuts-2021-01m.geojson/NUTS_RG_01M_2021_4326_LEVL_2.geojson",
            shapes_level_1=f"{EU_NUTS2021_DATASET['folder']}/ref-nuts-2021-01m.geojson/NUTS_RG_01M_2021_4326_LEVL_1.geojson",
            shapes_level_0=f"{EU_NUTS2021_DATASET['folder']}/ref-nuts-2021-01m.geojson/NUTS_RG_01M_2021_4326_LEVL_0.geojson",
        message:
            "Retrieving EU NUTS 2021 data"
        run:
            copy2(input["shapes"], output["zip_file"])
            unpack_archive(output["zip_file"], Path(output.shapes_level_3).parent)


if (
    BIDDING_ZONES_ELECTRICITYMAPS_DATASET := dataset_version(
        "bidding_zones_electricitymaps"
    )
)["source"] in ["primary", *ARCHIVE_SOURCES]:

    rule retrieve_bidding_zones_electricitymaps:
        input:
            geojson=storage(BIDDING_ZONES_ELECTRICITYMAPS_DATASET["url"]),
        output:
            geojson=f"{BIDDING_ZONES_ELECTRICITYMAPS_DATASET['folder']}/bidding_zones_electricitymaps.geojson",
        log:
            "logs/retrieve_bidding_zones_electricitymaps.log",
        retries: 2
        resources:
            mem_mb=1000,
        run:
            copy2(input["geojson"], output["geojson"])


if (BIDDING_ZONES_ENTSOEPY_DATASET := dataset_version("bidding_zones_entsoepy"))[
    "source"
] in ["primary", *ARCHIVE_SOURCES]:

    rule retrieve_bidding_zones_entsoepy:
        output:
            geojson=f"{BIDDING_ZONES_ENTSOEPY_DATASET['folder']}/bidding_zones_entsoepy.geojson",
        log:
            "logs/retrieve_bidding_zones_entsoepy.log",
        retries: 2
        resources:
            mem_mb=1000,
        run:
            import entsoe
            import geopandas as gpd
            import logging
            import requests
            from requests.exceptions import HTTPError, Timeout, ConnectionError

            logger = logging.getLogger(__name__)
            logger.setLevel(config["logging"]["level"])
            logger.addHandler(logging.FileHandler(log[0]))

            logger.info("Downloading entsoe-py zones...")
            gdfs: list[gpd.GeoDataFrame] = []
            url = f"{BIDDING_ZONES_ENTSOEPY_DATASET['url']}"
            for area in entsoe.Area:
                name = area.name
                logger.debug(f"Retrieving area {name}")
                try:
                    file_url = f"{url}/{name}.geojson"
                    response = requests.get(file_url, timeout=60)
                    response.raise_for_status()
                    features = response.json().get("features", [])
                    if features:
                        gdf = gpd.GeoDataFrame.from_features(
                            features, crs="EPSG:4326"
                        )
                        gdfs.append(gdf)
                    logger.debug(f"Successfully retrieved area {name}")
                except HTTPError as e:
                    logger.debug(f"Area file not available for {name}: {e}")
                    continue
                except (Timeout, ConnectionError, TimeoutError) as e:
                    raise Exception(f"Network error retrieving {name}: {e}")
            shapes = pd.concat(gdfs, ignore_index=True)  # type: ignore

            logger.info("Downloading entsoe-py zones... Done")

            shapes.to_file(output.geojson)



if (CUTOUT_DATASET := dataset_version("cutout"))["source"] in ARCHIVE_SOURCES:

    rule retrieve_cutout:
        input:
            storage(CUTOUT_DATASET["url"] + "/{cutout}.nc"),
        output:
            CUTOUT_DATASET["folder"] + "/{cutout}.nc",
        log:
            "logs/retrieve_cutout/{cutout}.log",
        retries: 2
        resources:
            mem_mb=5000,
        message:
            "Retrieving cutout data for {wildcards.cutout}"
        run:
            copy2(input[0], output[0])


if (COUNTRY_RUNOFF_DATASET := dataset_version("country_runoff"))[
    "source"
] in ARCHIVE_SOURCES:

    rule retrieve_country_runoff:
        input:
            storage(COUNTRY_RUNOFF_DATASET["url"]),
        output:
            era5_runoff=f"{COUNTRY_RUNOFF_DATASET['folder']}/era5-runoff-per-country.csv",
        message:
            "Retrieving country runoff data"
        run:
            copy2(input[0], output[0])


if (COUNTRY_HDD_DATASET := dataset_version("country_hdd"))["source"] in ARCHIVE_SOURCES:

    rule retrieve_country_hdd:
        input:
            storage(COUNTRY_HDD_DATASET["url"]),
        output:
            era5_runoff=f"{COUNTRY_HDD_DATASET['folder']}/era5-HDD-per-country.csv",
        message:
            "Retrieving country heating degree days data"
        run:
            copy2(input[0], output[0])


if (COSTS_DATASET := dataset_version("costs"))["source"] in [
    "primary",
    *ARCHIVE_SOURCES,
]:

    rule retrieve_cost_data:
        input:
            costs=storage(COSTS_DATASET["url"] + "/costs_{planning_horizons}.csv"),
        output:
            costs=COSTS_DATASET["folder"] + "/costs_{planning_horizons}.csv",
        message:
            "Retrieving cost data for {wildcards.planning_horizons}"
        run:
            copy2(input["costs"], output["costs"])


if (POWERPLANTS_DATASET := dataset_version("powerplants"))["source"] in [
    "primary",
    *ARCHIVE_SOURCES,
]:

    rule retrieve_powerplants:
        input:
            powerplants=storage(POWERPLANTS_DATASET["url"]),
        output:
            powerplants=f"{POWERPLANTS_DATASET['folder']}/powerplants.csv",
        message:
            "Retrieving powerplants data"
        run:
            copy2(input["powerplants"], output["powerplants"])


if (SCIGRID_GAS_DATASET := dataset_version("scigrid_gas"))["source"] in [
    "primary",
    *ARCHIVE_SOURCES,
]:

    rule retrieve_gas_infrastructure_data:
        input:
            zip_file=storage(SCIGRID_GAS_DATASET["url"]),
        output:
            zip_file=f"{SCIGRID_GAS_DATASET['folder']}/IGGIELGN.zip",
            entry=f"{SCIGRID_GAS_DATASET['folder']}/data/IGGIELGN_BorderPoints.geojson",
            storage=f"{SCIGRID_GAS_DATASET['folder']}/data/IGGIELGN_Storages.geojson",
            gas_network=f"{SCIGRID_GAS_DATASET['folder']}/data/IGGIELGN_PipeSegments.geojson",
        message:
            "Retrieving SciGRID gas infrastructure data"
        run:
            copy2(input["zip_file"], output["zip_file"])
            output_folder = Path(output["zip_file"]).parent
            unpack_archive(output["zip_file"], output_folder)


if (OPSD_DEMAND_DATA := dataset_version("opsd_electricity_demand"))["source"] in [
    "build"
]:

    rule retrieve_electricity_demand_opsd:
        output:
            csv=f"{OPSD_DEMAND_DATA['folder']}/electricity_demand_opsd_raw.csv",
        log:
            "logs/retrieve_electricity_demand_opsd.log",
        retries: 2
        resources:
            mem_mb=5000,
        params:
            versions=["2019-06-05", "2020-10-06"],
        message:
            "Retrieving electricity demand data from OPSD from build source"
        script:
            scripts("retrieve_electricity_demand_opsd.py")


if (OPSD_DEMAND_DATA := dataset_version("opsd_electricity_demand"))[
    "source"
] in ARCHIVE_SOURCES:

    rule retrieve_electricity_demand_opsd:
        input:
            csv=storage(OPSD_DEMAND_DATA["url"]),
        output:
            csv=f"{OPSD_DEMAND_DATA['folder']}/electricity_demand_opsd_raw.csv",
        retries: 2
        message:
            "Retrieving electricity demand data from OPSD from archive"
        run:
            copy2(input["csv"], output["csv"])


if (ENTSOE_DEMAND_DATA := dataset_version("entsoe_electricity_demand"))["source"] in [
    "build"
]:

    ENTSOE_COUNTRIES = [
        "AL",
        "AT",
        "BE",
        "BA",
        "BG",
        "CH",
        "CY",
        "CZ",
        "DE",
        "DK",
        "EE",
        "ES",
        "FI",
        "FR",
        "GB",
        "GR",
        "HR",
        "HU",
        "IE",
        "IT",
        "LT",
        "LU",
        "LV",
        "MD",
        "ME",
        "MK",
        "NL",
        "NO",
        "PL",
        "PT",
        "RO",
        "RS",
        "SE",
        "SI",
        "SK",
        "UA",
        "XK",
    ]

    rule retrieve_electricity_demand_entsoe_country:
        output:
            csv=f"{ENTSOE_DEMAND_DATA['folder']}"
            + "/electricity_demand_entsoe_raw_{country}.csv",
        log:
            "logs/retrieve_electricity_demand_entsoe_{country}.log",
        retries: 2
        resources:
            mem_mb=2000,
        params:
            entsoe_token=os.environ.get("ENTSOE_API_TOKEN", ""),
        message:
            "Retrieving electricity demand data from ENTSO-E for {wildcards.country}"
        script:
            scripts("retrieve_electricity_demand_entsoe.py")

    rule retrieve_electricity_demand_entsoe:
        input:
            csvs=expand(
                f"{ENTSOE_DEMAND_DATA['folder']}"
                + "/electricity_demand_entsoe_raw_{country}.csv",
                country=ENTSOE_COUNTRIES,
            ),
        output:
            csv=f"{ENTSOE_DEMAND_DATA['folder']}/electricity_demand_entsoe_raw.csv",
        message:
            "Retrieving electricity demand data from ENTSO-E from build source"
        run:
            import pandas as pd

            loads = [pd.read_csv(csv, index_col=0) for csv in input.csvs]
            df = pd.concat(loads, axis=1, join="outer").sort_index()
            df.to_csv(output.csv)


if (ENTSOE_DEMAND_DATA := dataset_version("entsoe_electricity_demand"))[
    "source"
] in ARCHIVE_SOURCES:

    rule retrieve_electricity_demand_entsoe:
        input:
            csv=storage(ENTSOE_DEMAND_DATA["url"]),
        output:
            csv=f"{ENTSOE_DEMAND_DATA['folder']}/electricity_demand_entsoe_raw.csv",
        retries: 2
        message:
            "Retrieving electricity demand data from ENTSO-E from archive"
        run:
            copy2(input["csv"], output["csv"])


if (NESO_DEMAND_DATA := dataset_version("neso_electricity_demand"))["source"] in [
    "build"
]:

    rule retrieve_electricity_demand_neso:
        output:
            csv=f"{NESO_DEMAND_DATA['folder']}/electricity_demand_neso_raw.csv",
        log:
            "logs/retrieve_electricity_demand_neso.log",
        retries: 2
        resources:
            mem_mb=5000,
        message:
            "Retrieving electricity demand data from NESO from build source"
        script:
            scripts("retrieve_electricity_demand_neso.py")


if (NESO_DEMAND_DATA := dataset_version("neso_electricity_demand"))[
    "source"
] in ARCHIVE_SOURCES:

    rule retrieve_electricity_demand_neso:
        input:
            csv=storage(NESO_DEMAND_DATA["url"]),
        output:
            csv=f"{NESO_DEMAND_DATA['folder']}/electricity_demand_neso_raw.csv",
        retries: 2
        message:
            "Retrieving electricity demand data from NESO from archive"
        run:
            copy2(input["csv"], output["csv"])


if (
    SYNTHETIC_ELECTRICITY_DEMAND_DATASET := dataset_version(
        "synthetic_electricity_demand"
    )
)["source"] in [
    "primary",
    *ARCHIVE_SOURCES,
]:

    rule retrieve_synthetic_electricity_demand:
        input:
            csv=storage(SYNTHETIC_ELECTRICITY_DEMAND_DATASET["url"]),
        output:
            csv=f"{SYNTHETIC_ELECTRICITY_DEMAND_DATASET['folder']}/load_synthetic_raw.csv",
        retries: 2
        message:
            "Retrieving synthetic electricity demand data"
        run:
            copy2(input["csv"], output["csv"])


if (ENERGY_ATLAS_DATASET := dataset_version("jrc_energy_atlas"))["source"] in [
    "primary",
    *ARCHIVE_SOURCES,
]:

    rule retrieve_electricity_demand_energy_atlas:
        output:
            tif=f"{ENERGY_ATLAS_DATASET['folder']}/electricity_tot_demand_2019.tif",
        message:
            "Retrieving JRC Energy Atlas electricity demand data raster"
        run:
            import requests

            url = ENERGY_ATLAS_DATASET["url"]
            response = requests.get(url)
            response.raise_for_status()
            with open(output["tif"], "wb") as f:
                f.write(response.content)



if (
    DESNZ_ELECTRICITY_CONSUMPTION_DATASET := dataset_version(
        "desnz_electricity_consumption"
    )
)["source"] in ["primary", *ARCHIVE_SOURCES]:

    rule retrieve_desnz_electricity_consumption:
        output:
            xlsx=f"{DESNZ_ELECTRICITY_CONSUMPTION_DATASET['folder']}/Subnational_electricity_consumption_statistics_2005-2024.xlsx",
        message:
            "Retrieving DESNZ subnational electricity consumption data"
        run:
            import requests

            url = DESNZ_ELECTRICITY_CONSUMPTION_DATASET["url"]
            response = requests.get(url)
            response.raise_for_status()
            with open(output["xlsx"], "wb") as f:
                f.write(response.content)



if (ONS_LAD_DATASET := dataset_version("ons_lad"))["source"] in [*ARCHIVE_SOURCES]:

    rule retrieve_ons_lad:
        input:
            geojson=storage(ONS_LAD_DATASET["url"]),
        output:
            geojson=f"{ONS_LAD_DATASET['folder']}/Local_Authority_Districts_May_2024_Boundaries__UK_BSC.geojson",
        message:
            "Retrieving UK ONS Local Authority Districts (LAD) Boundaries data"
        run:
            copy2(input["geojson"], output["geojson"])

elif ONS_LAD_DATASET["source"] in ["primary"]:

    rule retrieve_ons_lad:
        output:
            geojson=f"{ONS_LAD_DATASET['folder']}/Local_Authority_Districts_May_2024_Boundaries__UK_BSC.geojson",
        message:
            "Retrieving UK ONS Local Authority Districts (LAD) Boundaries data"
        run:
            import requests

            url = ONS_LAD_DATASET["url"]
            params = {
                "outFields": "*",
                "where": "1=1",
                "f": "geojson",
            }
            response = requests.get(url, params=params)
            with open(output["geojson"], "wb") as f:
                f.write(response.content)



if (SHIP_RASTER_DATASET := dataset_version("ship_raster"))["source"] in [
    "primary",
    *ARCHIVE_SOURCES,
]:

    rule retrieve_ship_raster:
        input:
            zip_file=storage(SHIP_RASTER_DATASET["url"]),
        output:
            zip_file=f"{SHIP_RASTER_DATASET['folder']}/shipdensity_global.zip",
        log:
            "logs/retrieve_ship_raster.log",
        retries: 2
        resources:
            mem_mb=5000,
        message:
            "Retrieving shipping raster data"
        run:
            copy2(input["zip_file"], output["zip_file"])


if (ENSPRESO_BIOMASS_DATASET := dataset_version("enspreso_biomass"))["source"] in [
    "primary",
    *ARCHIVE_SOURCES,
]:

    rule retrieve_enspreso_biomass:
        input:
            xlsx=storage(ENSPRESO_BIOMASS_DATASET["url"]),
        output:
            xlsx=f"{ENSPRESO_BIOMASS_DATASET['folder']}/ENSPRESO_BIOMASS.xlsx",
        retries: 1
        message:
            "Retrieving ENSPRESO biomass data"
        run:
            copy2(input["xlsx"], output["xlsx"])


if (HOTMAPS_INDUSTRIAL_SITES := dataset_version("hotmaps_industrial_sites"))[
    "source"
] in [
    "primary",
    *ARCHIVE_SOURCES,
]:

    rule retrieve_hotmaps_industrial_sites:
        input:
            csv=storage(HOTMAPS_INDUSTRIAL_SITES["url"]),
        output:
            csv=f"{HOTMAPS_INDUSTRIAL_SITES['folder']}/Industrial_Database.csv",
        retries: 1
        message:
            "Retrieving Hotmaps industrial sites"
        run:
            copy2(input["csv"], output["csv"])


if (NITROGEN_STATISTICS_DATASET := dataset_version("nitrogen_statistics"))[
    "source"
] in [
    "primary",
    *ARCHIVE_SOURCES,
]:

    rule retrieve_nitrogen_statistics:
        input:
            xlsx=storage(NITROGEN_STATISTICS_DATASET["url"]),
        output:
            xlsx=f"{NITROGEN_STATISTICS_DATASET['folder']}/nitro-ert.xlsx",
        retries: 1
        message:
            "Retrieving nitrogen statistics data"
        run:
            copy2(input["xlsx"], output["xlsx"])


if (COPERNICUS_LAND_COVER_DATASET := dataset_version("copernicus_land_cover"))[
    "source"
] in ["primary", *ARCHIVE_SOURCES]:

    # Downloading Copernicus Global Land Cover for land cover and land use:
    # Website: https://land.copernicus.eu/global/products/lc
    rule download_copernicus_land_cover:
        input:
            tif=storage(COPERNICUS_LAND_COVER_DATASET["url"]),
        output:
            tif=f"{COPERNICUS_LAND_COVER_DATASET['folder']}/Copernicus_LC100_global_v3.0.1_2019-nrt_Discrete-Classification-map_EPSG-4326.tif",
        message:
            "Retrieving Copernicus land cover data"
        run:
            copy2(input["tif"], output["tif"])


if (LUISA_LAND_COVER_DATASET := dataset_version("luisa_land_cover"))["source"] in [
    "primary",
    *ARCHIVE_SOURCES,
]:

    # Downloading LUISA Base Map for land cover and land use:
    # Website: https://ec.europa.eu/jrc/en/luisa
    rule retrieve_luisa_land_cover:
        input:
            tif=storage(LUISA_LAND_COVER_DATASET["url"]),
        output:
            tif=f"{LUISA_LAND_COVER_DATASET['folder']}/LUISA_basemap_020321_50m.tif",
        message:
            "Retrieving LUISA land cover data"
        run:
            copy2(input["tif"], output["tif"])


if (EEZ_DATASET := dataset_version("eez"))["source"] in ["primary"]:

    rule retrieve_eez:
        output:
            zip_file=f"{EEZ_DATASET['folder']}/World_EEZ_{EEZ_DATASET['version']}_LR.zip",
            gpkg=f"{EEZ_DATASET['folder']}/World_EEZ_{EEZ_DATASET['version']}_LR/eez_{EEZ_DATASET['version'].split('_')[0]}_lowres.gpkg",
        message:
            "Retrieving EEZ data"
        run:
            from uuid import uuid4

            name = str(uuid4())[:8]
            org = str(uuid4())[:8]

            response = requests.post(
                f"{EEZ_DATASET['url']}",
                params={"name": f"World_EEZ_{EEZ_DATASET['version']}_LR.zip"},
                data={
                    "name": name,
                    "organisation": org,
                    "email": f"{name}@{org}.org",
                    "country": "Germany",
                    "user_category": "academia",
                    "purpose_category": "Research",
                    "agree": "1",
                },
            )

            with open(output["zip_file"], "wb") as f:
                f.write(response.content)
            output_folder = Path(output["zip_file"]).parent
            unpack_archive(output["zip_file"], output_folder)


elif (EEZ_DATASET := dataset_version("eez"))["source"] in ARCHIVE_SOURCES:

    rule retrieve_eez:
        input:
            zip_file=storage(
                EEZ_DATASET["url"],
            ),
        output:
            zip_file=f"{EEZ_DATASET['folder']}/World_EEZ_{EEZ_DATASET['version']}_LR.zip",
            gpkg=f"{EEZ_DATASET['folder']}/World_EEZ_{EEZ_DATASET['version']}_LR/eez_{EEZ_DATASET['version'].split('_')[0]}_lowres.gpkg",
        message:
            "Retrieving EEZ data"
        run:
            output_folder = Path(output["zip_file"]).parent
            copy2(input["zip_file"], output["zip_file"])
            unpack_archive(output["zip_file"], output_folder)


if (WB_URB_POP_DATASET := dataset_version("worldbank_urban_population"))["source"] in [
    "primary",
    *ARCHIVE_SOURCES,
]:

    rule retrieve_worldbank_urban_population:
        input:
            zip=storage(WB_URB_POP_DATASET["url"]),
        output:
            zip=f"{WB_URB_POP_DATASET['folder']}/API_SP.URB.TOTL.IN.ZS_DS2_en_csv_v2.zip",
            csv=f"{WB_URB_POP_DATASET['folder']}/API_SP.URB.TOTL.IN.ZS_DS2_en_csv_v2.csv",
        message:
            "Retrieving World Bank urban population data"
        run:
            copy2(input["zip"], output["zip"])
            unpack_archive(output["zip"], WB_URB_POP_DATASET["folder"])

            # Filename contains some added numbers when downloaded,
            # remove them to have a consistent filename across versions
            target_filename = Path(output["csv"])
            origin_filename = next(
                Path(WB_URB_POP_DATASET["folder"]).rglob(
                    target_filename.stem + "*" + target_filename.suffix
                )
            )
            origin_filename.rename(output.csv)


if (CO2STOP_DATASET := dataset_version("co2stop"))["source"] in [
    "primary",
    *ARCHIVE_SOURCES,
]:

    rule retrieve_co2stop:
        input:
            zip_file=storage(CO2STOP_DATASET["url"]),
        output:
            zip_file=f"{CO2STOP_DATASET['folder']}/co2jrc_openformats.zip",
            storage_table=f"{CO2STOP_DATASET['folder']}/CO2JRC_OpenFormats/CO2Stop_DataInterrogationSystem/Hydrocarbon_Storage_Units.csv",
            storage_map=f"{CO2STOP_DATASET['folder']}/CO2JRC_OpenFormats/CO2Stop_Polygons Data/StorageUnits_March13.kml",
            traps_table1=f"{CO2STOP_DATASET['folder']}/CO2JRC_OpenFormats/CO2Stop_DataInterrogationSystem/Hydrocarbon_Traps.csv",
            traps_table2=f"{CO2STOP_DATASET['folder']}/CO2JRC_OpenFormats/CO2Stop_DataInterrogationSystem/Hydrocarbon_Traps_Temp.csv",
            traps_table3=f"{CO2STOP_DATASET['folder']}/CO2JRC_OpenFormats/CO2Stop_DataInterrogationSystem/Hydrocarbon_Traps1.csv",
            traps_map=f"{CO2STOP_DATASET['folder']}/CO2JRC_OpenFormats/CO2Stop_Polygons Data/DaughterUnits_March13.kml",
        message:
            "Retrieving CO2STOP data"
        run:
            output_folder = Path(output["zip_file"]).parent
            output_folder.mkdir(parents=True, exist_ok=True)
            copy2(input["zip_file"], output["zip_file"])
            unpack_archive(output["zip_file"], output_folder)


if (GEM_EUROPE_GAS_TRACKER_DATASET := dataset_version("gem_europe_gas_tracker"))[
    "source"
] in [
    "primary",
    *ARCHIVE_SOURCES,
]:

    rule retrieve_gem_europe_gas_tracker:
        input:
            xlsx=storage(GEM_EUROPE_GAS_TRACKER_DATASET["url"]),
        output:
            xlsx="data/gem/Europe-Gas-Tracker-2024-05.xlsx",
        message:
            "Retrieving GEM Europe Gas Tracker data"
        run:
            copy2(input["xlsx"], output["xlsx"])


if (GEM_GSPT_DATASET := dataset_version("gem_gspt"))["source"] in [
    "primary",
    *ARCHIVE_SOURCES,
]:

    rule retrieve_gem_steel_plant_tracker:
        input:
            xlsx=storage(GEM_GSPT_DATASET["url"]),
        output:
            xlsx=f"{GEM_GSPT_DATASET['folder']}/Global-Steel-Plant-Tracker.xlsx",
        message:
            "Retrieving GEM Global Steel Plant Tracker data"
        run:
            copy2(input["xlsx"], output["xlsx"])


if (GEM_GCCT_DATASET := dataset_version("gem_gcct"))["source"] in [
    "primary",
    *ARCHIVE_SOURCES,
]:

    rule retrieve_gem_cement_concrete_tracker:
        input:
            xlsx=storage(GEM_GCCT_DATASET["url"]),
        output:
            xlsx=f"{GEM_GCCT_DATASET['folder']}/Global-Cement-and-Concrete-Tracker.xlsx",
        run:
            copy2(input["xlsx"], output["xlsx"])


if (BFS_ROAD_VEHICLE_STOCK_DATASET := dataset_version("bfs_road_vehicle_stock"))[
    "source"
] in [
    "primary",
    *ARCHIVE_SOURCES,
]:

    rule retrieve_bfs_road_vehicle_stock:
        input:
            csv=storage(BFS_ROAD_VEHICLE_STOCK_DATASET["url"]),
        output:
            csv=f"{BFS_ROAD_VEHICLE_STOCK_DATASET['folder']}/vehicle_stock.csv",
        message:
            "Retrieving BFS road vehicle stock data"
        run:
            copy2(input["csv"], output["csv"])


if (BFS_GDP_AND_POPULATION_DATASET := dataset_version("bfs_gdp_and_population"))[
    "source"
] in [
    "primary",
    *ARCHIVE_SOURCES,
]:

    rule retrieve_bfs_gdp_and_population:
        input:
            xlsx=storage(BFS_GDP_AND_POPULATION_DATASET["url"]),
        output:
            xlsx=f"{BFS_GDP_AND_POPULATION_DATASET['folder']}/gdp_and_population.xlsx",
        message:
            "Retrieving BFS GDP and population data"
        run:
            copy2(input["xlsx"], output["xlsx"])


def get_wdpa_url(DATASET) -> str:
    """
    Find the right URL for the WDPA / WDPA marine dataset based on the source type.
    """
    if DATASET["source"] in ARCHIVE_SOURCES:
        return DATASET["url"]
    elif DATASET["source"] == "primary":
        # Some logic to find the correct file URL from the WDPA website (primary source)
        # Sometimes files are released delayed or ahead of schedule, check which file is currently available
        def check_file_exists(url):
            response = requests.head(url)
            return response.status_code == 200

        # Basic pattern where WDPA files can be found
        url_pattern = DATASET["url"]

        # 3-letter month + 4 digit year for current/previous/next/pprevious/nnext months to test
        # order reflects priority of testing
        months = [
            datetime.now(),  # current
            (datetime.now() + relativedelta(months=-1)),  # previous month
            (datetime.now() + relativedelta(months=+1)),  # next month
            (datetime.now() + relativedelta(months=-2)),  # two months ago
            (datetime.now() + relativedelta(months=+2)),  # two months ahead
        ]
        months = [m.strftime("%b%Y") for m in months]

        # Test prioritised: current month -> previous -> next
        for bYYYY in months:
            url = url_pattern.format(bYYYY=bYYYY)
            if check_file_exists(url):
                return url

        raise ValueError(
            f"No {DATASET.dataset} files found at {url_pattern} for bY={months}."
        )


if (WDPA_DATASET := dataset_version("wdpa"))["source"] in [
    "primary",
    *ARCHIVE_SOURCES,
]:

    # Downloading protected area database from WDPA
    # extract the main zip and then merge the contained 3 zipped shapefiles
    # Website: https://www.protectedplanet.net/en/thematic-areas/wdpa
    rule retrieve_wdpa:
        input:
            zip_file=storage(get_wdpa_url(WDPA_DATASET)),
        output:
            zip_file=f"{WDPA_DATASET['folder']}/WDPA_shp.zip",
            gpkg=f"{WDPA_DATASET['folder']}/WDPA.gpkg",
        message:
            "Downloading protected area database from WDPA"
        run:
            output_folder = Path(output["zip_file"]).parent
            copy2(input["zip_file"], output["zip_file"])
            unpack_archive(output["zip_file"], output_folder)

            # Extract {bYYYY} from the input file / URL
            bYYYY = re.search(
                r"WDPA_(\w{3}\d{4})_Public_shp.zip",
                input["zip_file"],
            ).group(1)

            for i in range(3):
                # vsizip is special driver for directly working with zipped shapefiles in ogr2ogr
                layer_path = (
                    f"/vsizip/{output_folder}/WDPA_{bYYYY}_Public_shp_{i}.zip"
                )
                print(f"Adding layer {i+1} of 3 to combined output file.")
                shell("ogr2ogr -f gpkg -update -append {output.gpkg} {layer_path}")



if (WDPA_MARINE_DATASET := dataset_version("wdpa_marine"))["source"] in [
    "primary",
    *ARCHIVE_SOURCES,
]:

    rule retrieve_wdpa_marine:
        input:
            zip_file=storage(get_wdpa_url(WDPA_MARINE_DATASET)),
        output:
            zip_file=f"{WDPA_MARINE_DATASET['folder']}/WDPA_WDOECM_marine.zip",
            gpkg=f"{WDPA_MARINE_DATASET['folder']}/WDPA_WDOECM_marine.gpkg",
        # Downloading Marine protected area database from WDPA
        # extract the main zip and then merge the contained 3 zipped shapefiles
        # Website: https://www.protectedplanet.net/en/thematic-areas/marine-protected-areas
        message:
            "Downloading Marine protected area database from WDPA"
        run:
            output_folder = Path(output["zip_file"]).parent
            copy2(input["zip_file"], output["zip_file"])
            unpack_archive(output["zip_file"], output_folder)

            # Extract {bYYYY} from the input file / URL
            bYYYY = re.search(
                r"WDPA_WDOECM_(\w{3}\d{4})_Public_marine_shp.zip",
                input["zip_file"],
            ).group(1)

            for i in range(3):
                # vsizip is special driver for directly working with zipped shapefiles in ogr2ogr
                layer_path = f"/vsizip/{output_folder}/WDPA_WDOECM_{bYYYY}_Public_marine_shp_{i}.zip"
                print(f"Adding layer {i+1} of 3 to combined output file.")
                shell("ogr2ogr -f gpkg -update -append {output.gpkg} {layer_path}")



if (INSTRAT_CO2_PRICES_DATASET := dataset_version("instrat_co2_prices"))["source"] in [
    "primary",
]:

    rule retrieve_co2_prices:
        output:
            csv=f"{INSTRAT_CO2_PRICES_DATASET['folder']}/prices_eu_ets_all.csv",
        log:
            "logs/retrieve_co2_prices.log",
        retries: 2
        resources:
            mem_mb=5000,
        message:
            "Retrieving CO2 emission allowances price in EU ETS system"
        run:
            import pandas as pd

            url = "https://energy-api.instrat.pl/api/prices/co2?all=1"
            headers = {
                "User-Agent": "Mozilla/5.0",
                "Accept": "application/json",
                "Referer": "https://energy.instrat.pl/",
            }

            r = requests.get(url, headers=headers)
            r.raise_for_status()

            df = pd.read_json(r.text)
            df.to_csv(output["csv"], index=False)


if (
    WORLD_BANK_COMMODITY_PRICES_DATASET := dataset_version("worldbank_commodity_prices")
)["source"] in ["primary", *ARCHIVE_SOURCES]:

    rule retrieve_worldbank_commodity_prices:
        input:
            xlsx=storage(WORLD_BANK_COMMODITY_PRICES_DATASET["url"]),
        output:
            xlsx=f"{WORLD_BANK_COMMODITY_PRICES_DATASET['folder']}/CMO-Historical-Data-Monthly.xlsx",
        message:
            "Retrieving monthly commodity price time series (including fossil fuels)"
        run:
            copy2(input["xlsx"], output["xlsx"])


if (TYNDP_DATASET := dataset_version("tyndp"))["source"] in [
    "primary",
    *ARCHIVE_SOURCES,
]:

    if not config["tyndp_scenario"]:

        rule retrieve_tyndp:
            input:
                line_data=storage(TYNDP_DATASET["url"] + "/Line-data.zip"),
                nodes=storage(TYNDP_DATASET["url"] + "/Nodes.zip"),
            output:
                line_data_zip=f"{TYNDP_DATASET['folder']}/Line-data.zip",
                nodes_zip=f"{TYNDP_DATASET['folder']}/Nodes.zip",
                elec_reference_grid=f"{TYNDP_DATASET['folder']}/Line data/ReferenceGrid_Electricity.xlsx",
                nodes=f"{TYNDP_DATASET['folder']}/Nodes/LIST OF NODES.xlsx",
            log:
                "logs/retrieve_tyndp.log",
            message:
                "Retrieving TYNDP network topology data"
            run:
                for key in input.keys():
                    # Keep zip file
                    copy2(input[key], output[f"{key}_zip"])

                    # unzip
                    output_folder = Path(output[f"{key}_zip"]).parent
                    unpack_archive(output[f"{key}_zip"], output_folder)

                    # Remove __MACOSX directory if it exists
                    macosx_dir = output_folder / "__MACOSX"
                    rmtree(macosx_dir, ignore_errors=True)

    else:

        rule retrieve_tyndp:
            input:
                line_data=storage(TYNDP_DATASET["url"] + "/Line-data.zip"),
                nodes=storage(TYNDP_DATASET["url"] + "/Nodes.zip"),
                hydro_inflows=storage(TYNDP_DATASET["url"] + "/Hydro-Inflows.zip"),
                pemmdb=storage(TYNDP_DATASET["url"] + "/PEMMDB2.zip"),
                supply_tool=storage(
                    TYNDP_DATASET["url"] + "/20240518-Supply-Tool.xlsm.zip"
                ),
                benchmark=storage(
                    TYNDP_DATASET["url"] + "/TYNDP-2024-Scenarios-Package-20250128.zip"
                ),
                demand_profiles=storage(TYNDP_DATASET["url"] + "/Demand-Profiles.zip"),
                ev_modelling=storage(TYNDP_DATASET["url"] + "/EV-Modelling-Inputs.zip"),
                hybrid_hp_modelling=storage(
                    TYNDP_DATASET["url"] + "/Hybrid-Heat-Pump-Modelling-Inputs.zip"
                ),
                hydrogen=storage(TYNDP_DATASET["url"] + "/Hydrogen.zip"),
                investment_datasets=storage(
                    TYNDP_DATASET["url"] + "/Investment-Datasets.zip"
                ),
                offshore_hubs=storage(TYNDP_DATASET["url"] + "/Offshore-hubs.zip"),
                market_outputs_NT2030_CY2009=storage(
                    TYNDP_DATASET["url"]
                    + "/MMStandardOutputFile_NT2030_Plexos_CY2009_2.5_v40.xlsx.zip"
                ),
                market_outputs_NT2040_CY2009=storage(
                    TYNDP_DATASET["url"]
                    + "/MMStandardOutputFile_NT2040_Plexos_CY2009_2.5_v40.xlsx.zip"
                ),
                h2_reference_grid_entsoe=storage(
                    TYNDP_DATASET["url"] + "/StartingGrid2030.xlsx.zip"
                ),
            output:
                line_data_zip=f"{TYNDP_DATASET['folder']}/Line-data.zip",
                nodes_zip=f"{TYNDP_DATASET['folder']}/Nodes.zip",
                elec_reference_grid=f"{TYNDP_DATASET['folder']}/Line data/ReferenceGrid_Electricity.xlsx",
                h2_reference_grid_entsos=f"{TYNDP_DATASET['folder']}/Line data/ReferenceGrid_Hydrogen.xlsx",
                h2_reference_grid_entsoe_zip=f"{TYNDP_DATASET['folder']}/StartingGrid2030.xlsx.zip",
                h2_reference_grid_entsoe=f"{TYNDP_DATASET['folder']}/StartingGrid2030.xlsx",
                nodes=f"{TYNDP_DATASET['folder']}/Nodes/LIST OF NODES.xlsx",
                hydro_inflows_zip=f"{TYNDP_DATASET['folder']}/Hydro-Inflows.zip",
                hydro_inflows=directory(f"{TYNDP_DATASET['folder']}/Hydro Inflows"),
                pemmdb_zip=f"{TYNDP_DATASET['folder']}/PEMMDB2.zip",
                pemmdb=directory(f"{TYNDP_DATASET['folder']}/PEMMDB2"),
                supply_tool_zip=f"{TYNDP_DATASET['folder']}/20240518-Supply-Tool.xlsm.zip",
                supply_tool=f"{TYNDP_DATASET['folder']}/20240518-Supply-Tool.xlsm",
                benchmark_zip=f"{TYNDP_DATASET['folder']}/TYNDP-2024-Scenarios-Package-20250128.zip",
                benchmark=f"{TYNDP_DATASET['folder']}/TYNDP-2024-Scenarios-Package/TYNDP_2024-Scenario-Report-Data-Figures_240522.xlsx",
                demand_profiles_zip=f"{TYNDP_DATASET['folder']}/Demand-Profiles.zip",
                demand_profiles=directory(f"{TYNDP_DATASET['folder']}/Demand Profiles"),
                ev_modelling_zip=f"{TYNDP_DATASET['folder']}/EV-Modelling-Inputs.zip",
                ev_modelling=directory(f"{TYNDP_DATASET['folder']}/EV Modelling Inputs"),
                hybrid_hp_modelling_zip=f"{TYNDP_DATASET['folder']}/Hybrid-Heat-Pump-Modelling-Inputs.zip",
                hybrid_hp_modelling=directory(
                    f"{TYNDP_DATASET['folder']}/Hybrid Heat Pump Modelling Inputs"
                ),
                hydrogen_zip=f"{TYNDP_DATASET['folder']}/Hydrogen.zip",
                hydrogen=directory(f"{TYNDP_DATASET['folder']}/Hydrogen"),
                h2_imports=f"{TYNDP_DATASET['folder']}/Hydrogen/H2 IMPORTS GENERATORS PROPERTIES.xlsx",
                h2_storages=f"{TYNDP_DATASET['folder']}/Hydrogen/H2 STORAGES.xlsx",
                smr=f"{TYNDP_DATASET['folder']}/Hydrogen/SMR Figures.xlsx",
                investment_datasets_zip=f"{TYNDP_DATASET['folder']}/Investment-Datasets.zip",
                trajectories=f"{TYNDP_DATASET['folder']}/Investment Datasets/TRAJECTORY.xlsx",
                invest_grid=f"{TYNDP_DATASET['folder']}/Investment Datasets/GRID.xlsx",
                offshore_hubs_zip=f"{TYNDP_DATASET['folder']}/Offshore-hubs.zip",
                offshore_nodes=f"{TYNDP_DATASET['folder']}/Offshore hubs/NODE.xlsx",
                offshore_grid=f"{TYNDP_DATASET['folder']}/Offshore hubs/GRID.xlsx",
                offshore_electrolysers=f"{TYNDP_DATASET['folder']}/Offshore hubs/ELECTROLYSER.xlsx",
                offshore_generators=f"{TYNDP_DATASET['folder']}/Offshore hubs/GENERATOR.xlsx",
                market_outputs_NT2030_CY2009_zip=f"{TYNDP_DATASET['folder']}/MMStandardOutputFile_NT2030_Plexos_CY2009_2.5_v40.xlsx.zip",
                market_outputs_NT2030_CY2009=f"{TYNDP_DATASET['folder']}/MMStandardOutputFile_NT2030_Plexos_CY2009_2.5_v40.xlsx",
                market_outputs_NT2040_CY2009_zip=f"{TYNDP_DATASET['folder']}/MMStandardOutputFile_NT2040_Plexos_CY2009_2.5_v40.xlsx.zip",
                market_outputs_NT2040_CY2009=f"{TYNDP_DATASET['folder']}/MMStandardOutputFile_NT2040_Plexos_CY2009_2.5_v40.xlsx",
            log:
                "logs/retrieve_tyndp.log",
            message:
                "Retrieving TYNDP 2024 data package"
            run:
                for key in input.keys():
                    # Keep zip file
                    copy2(input[key], output[f"{key}_zip"])

                    # unzip
                    output_folder = Path(output[f"{key}_zip"]).parent
                    unpack_archive(output[f"{key}_zip"], output_folder)

                    # Remove __MACOSX directory if it exists
                    macosx_dir = output_folder / "__MACOSX"
                    rmtree(macosx_dir, ignore_errors=True)



def get_osm_archive_files(version):
    return [
        "buses.csv",
        "converters.csv",
        "lines.csv",
        "links.csv",
        "transformers.csv",
        # Newer versions include the additional map.html file for visualisation
        *(["map.html"] if float(version) >= 0.6 else []),
    ]


def get_osm_network_incumbent(
    version: str = "latest",
    source: str = "archive",
) -> pd.Series:
    fp = workflow.source_path("../data/versions.csv")
    data_versions = load_data_versions(fp)
    name = "osm"

    dataset = data_versions.loc[
        (data_versions["dataset"] == name)
        & (data_versions["source"] == source)
        & (data_versions["supported"])  # Limit to supported versions only
        & (data_versions["version"] == version if "latest" != version else True)
        & (data_versions["latest"] if "latest" == version else True)
    ]

    if dataset.empty:
        raise ValueError(
            f"OSM network for version '{version}' not found in data/versions.csv."
        )

    # Return single-row DataFrame as a Series
    dataset = dataset.squeeze()

    # Generate output folder path in the `data` directory
    dataset["folder"] = Path(
        "data", name, dataset["source"], dataset["version"]
    ).as_posix()

    return dataset


def input_base_network_incumbent(w):
    version = config_provider("osm_network_release", "compare_to", "version")(w)
    source = config_provider("osm_network_release", "compare_to", "source")(w)
    osm_dataset = get_osm_network_incumbent(version, source)
    osm_path = osm_dataset["folder"]
    components = {"buses", "lines", "links", "converters", "transformers"}
    inputs = {c: f"{osm_path}/{c}.csv" for c in components}
    return inputs


if OSM_DATASET["source"] in ARCHIVE_SOURCES:
    OSM_ARCHIVE_FILES = get_osm_archive_files(OSM_DATASET["version"])

    rule retrieve_osm_archive:
        input:
            **{
                file: storage(f"{OSM_DATASET['url']}/{file}")
                for file in OSM_ARCHIVE_FILES
            },
        output:
            **{file: f"{OSM_DATASET['folder']}/{file}" for file in OSM_ARCHIVE_FILES},
        log:
            "logs/retrieve_osm_archive.log",
        threads: 1
        resources:
            mem_mb=500,
        message:
            "Retrieving OSM archive data"
        run:
            for key in input.keys():
                copy2(input[key], output[key])



# Only create incumbent rule if it points to a different folder
OSM_DATASET_INCUMBENT = get_osm_network_incumbent(
    version=config.get("osm_network_release", {})
    .get("compare_to", {})
    .get("version", "latest"),
    source=config.get("osm_network_release", {})
    .get("compare_to", {})
    .get("source", "archive"),
)

if OSM_DATASET_INCUMBENT["source"] in ARCHIVE_SOURCES and OSM_DATASET_INCUMBENT[
    "folder"
] != OSM_DATASET.get("folder"):

    OSM_ARCHIVE_FILES_INCUMBENT = get_osm_archive_files(
        OSM_DATASET_INCUMBENT["version"]
    )

    rule retrieve_osm_archive_incumbent:
        input:
            **{
                file: storage(f"{OSM_DATASET_INCUMBENT['url']}/{file}")
                for file in OSM_ARCHIVE_FILES_INCUMBENT
            },
        output:
            **{
                file: f"{OSM_DATASET_INCUMBENT['folder']}/{file}"
                for file in OSM_ARCHIVE_FILES_INCUMBENT
            },
        log:
            "logs/retrieve_osm_archive_incumbent.log",
        threads: 1
        resources:
            mem_mb=500,
        message:
            "Retrieving OSM archive incumbent data"
        run:
            for key in input.keys():
                copy2(input[key], output[key])



if OSM_DATASET["source"] == "build":
    OSM_RAW_JSON = [
        "cables_way.json",
        "lines_way.json",
        "routes_relation.json",
        "substations_way.json",
        "substations_relation.json",
    ]

    rule retrieve_osm_data_raw:
        output:
            **{
                file.replace(
                    ".json", ""
                ): f"{OSM_DATASET['folder']}/{{country}}/{file}"
                for file in OSM_RAW_JSON
            },
        log:
            "logs/retrieve_osm_data_{country}.log",
        threads: 1
        params:
            overpass_api=config_provider("overpass_api"),
        message:
            "Retrieving OSM electricity grid raw data for {wildcards.country}"
        script:
            scripts("retrieve_osm_data.py")

    rule retrieve_osm_data_raw_all:
        input:
            expand(
                f"{OSM_DATASET['folder']}/{{country}}/{{file}}",
                country=config_provider("countries"),
                file=OSM_RAW_JSON,
            ),


if (NATURA_DATASET := dataset_version("natura"))["source"] in ARCHIVE_SOURCES:

    rule retrieve_natura:
        input:
            storage(NATURA_DATASET["url"]),
        output:
            f"{NATURA_DATASET['folder']}/natura.tiff",
        log:
            "logs/retrieve_natura.log",
        message:
            "Retrieving Natura 2000 raster data"
        run:
            copy2(input[0], output[0])

elif NATURA_DATASET["source"] == "build":

    rule build_natura_raster:
        input:
            online=storage(NATURA_DATASET["url"]),
            cutout=lambda w: input_cutout(w),
        output:
            zip=f"{NATURA_DATASET['folder']}/raw/natura.zip",
            raw=directory(f"{NATURA_DATASET['folder']}/raw"),
            raster=f"{NATURA_DATASET['folder']}/natura.tiff",
        log:
            "logs/build_natura.log",
        resources:
            mem_mb=5000,
        message:
            "Building Natura 2000 raster data"
        script:
            scripts("build_natura.py")


if (OSM_BOUNDARIES_DATASET := dataset_version("osm_boundaries"))["source"] in [
    "primary"
]:

    rule retrieve_osm_boundaries:
        output:
            json=f"{OSM_BOUNDARIES_DATASET['folder']}/{country}_adm1.json",
        log:
            "logs/retrieve_osm_boundaries_{country}_adm1.log",
        threads: 1
        message:
            "Retrieving OSM admin boundaries for {wildcards.country}"
        script:
            scripts("retrieve_osm_boundaries.py")

elif (OSM_BOUNDARIES_DATASET := dataset_version("osm_boundaries"))[
    "source"
] in ARCHIVE_SOURCES:

    rule retrieve_osm_boundaries:
        input:
            storage(
                f"{OSM_BOUNDARIES_DATASET['url']}",
            ),
        output:
            json1=f"{OSM_BOUNDARIES_DATASET['folder']}/XK_adm1.json",
            json2=f"{OSM_BOUNDARIES_DATASET['folder']}/UA_adm1.json",
            json3=f"{OSM_BOUNDARIES_DATASET['folder']}/MD_adm1.json",
            json4=f"{OSM_BOUNDARIES_DATASET['folder']}/BA_adm1.json",
            zip_file=f"{OSM_BOUNDARIES_DATASET['folder']}/osm_boundaries.zip",
        message:
            "Retrieving OSM admin boundaries data"
        run:
            output_folder = Path(output["zip_file"]).parent
            copy2(input[0], output["zip_file"])
            unpack_archive(output["zip_file"], output_folder)


if (
    GEOTHERMAL_HEAT_UTILISATION_POTENTIALS_DATASET := dataset_version(
        "geothermal_heat_utilisation_potentials"
    )
)["source"] in ["primary", *ARCHIVE_SOURCES]:

    rule retrieve_geothermal_heat_utilisation_potentials:
        input:
            isi_heat_potentials=storage(
                GEOTHERMAL_HEAT_UTILISATION_POTENTIALS_DATASET["url"]
            ),
        output:
            isi_heat_potentials=f"{GEOTHERMAL_HEAT_UTILISATION_POTENTIALS_DATASET['folder']}/isi_heat_utilisation_potentials.xlsx",
        log:
            "logs/retrieve_geothermal_heat_utilisation_potentials.log",
        retries: 2
        threads: 1
        message:
            "Retrieving geothermal heat utilisation potentials"
        run:
            copy2(input["isi_heat_potentials"], output["isi_heat_potentials"])


if (LAU_REGIONS_DATASET := dataset_version("lau_regions"))["source"] in [
    "primary",
    *ARCHIVE_SOURCES,
]:

    rule retrieve_lau_regions:
        input:
            lau_regions=storage(LAU_REGIONS_DATASET["url"]),
        output:
            zip=f"{LAU_REGIONS_DATASET['folder']}/lau_regions.zip",
        log:
            "logs/retrieve_lau_regions.log",
        retries: 2
        threads: 1
        message:
            "Retrieving Local Administrative Units and Administation Unit regions"
        run:
            copy2(input["lau_regions"], output["zip"])

    rule retrieve_seawater_temperature:
        output:
            seawater_temperature="data/seawater_temperature_{year}.nc",
        log:
            "logs/retrieve_seawater_temperature_{year}.log",
        resources:
            mem_mb=10000,
        params:
            default_cutout=config_provider("atlite", "default_cutout"),
            test_data_url=dataset_version("seawater_temperature")["url"],
        message:
            "Retrieving seawater temperature data for {wildcards.year}"
        script:
            scripts("retrieve_seawater_temperature.py")

    rule retrieve_hera_data_test_cutout:
        input:
            hera_data_url=storage(
                f"https://zenodo.org/records/15828866/files/hera_be_2013-03-01_to_2013-03-08.zip"
            ),
        output:
            river_discharge=f"data/hera_be_2013-03-01_to_2013-03-08/river_discharge_be_2013-03-01_to_2013-03-08.nc",
            ambient_temperature=f"data/hera_be_2013-03-01_to_2013-03-08/ambient_temp_be_2013-03-01_to_2013-03-08.nc",
        log:
            "logs/retrieve_hera_data_test_cutout.log",
        retries: 2
        resources:
            mem_mb=10000,
        params:
            folder="data",
        message:
            "Retrieving HERA test cutout data"
        run:
            unpack_archive(input[0], params.folder)

    rule retrieve_hera_data:
        input:
            river_discharge=storage(
                "https://jeodpp.jrc.ec.europa.eu/ftp/jrc-opendata/CEMS-EFAS/HERA/VER1-0/Data/NetCDF/river_discharge/dis.HERA{year}.nc"
            ),
            ambient_temperature=storage(
                "https://jeodpp.jrc.ec.europa.eu/ftp/jrc-opendata/CEMS-EFAS/HERA/VER1-0/Data/NetCDF/climate_inputs/ta6/ta6_{year}.nc"
            ),
        output:
            river_discharge="data/hera_{year}/river_discharge_{year}.nc",
            ambient_temperature="data/hera_{year}/ambient_temp_{year}.nc",
        log:
            "logs/retrieve_hera_data_{year}.log",
        retries: 2
        resources:
            mem_mb=10000,
        params:
            snapshot_year="{year}",
        message:
            "Retrieving HERA data for {wildcards.year}"
        run:
            move(input.river_discharge, output.river_discharge)
            move(input.ambient_temperature, output.ambient_temperature)


if (JRC_ARDECO_DATASET := dataset_version("jrc_ardeco"))["source"] in [
    "primary",
]:

    rule retrieve_jrc_ardeco:
        input:
            ardeco_gdp=storage(
                f"{JRC_ARDECO_DATASET['url']}/SUVGDP?versions=2021&unit=EUR&format=csv-table"
            ),
            ardeco_pop=storage(
                f"{JRC_ARDECO_DATASET['url']}/SNPTD?versions=2021&unit=EUR&format=csv-table"
            ),
        output:
            ardeco_gdp=f"{JRC_ARDECO_DATASET['folder']}/ARDECO-SUVGDP.2021.table.csv",
            ardeco_pop=f"{JRC_ARDECO_DATASET['folder']}/ARDECO-SNPTD.2021.table.csv",
        message:
            "Retrieving JRC ARDECO data"
        run:
            for key in input.keys():
                copy2(input[key], output[key])


elif (JRC_ARDECO_DATASET := dataset_version("jrc_ardeco"))["source"] in ARCHIVE_SOURCES:

    rule retrieve_jrc_ardeco:
        input:
            ardeco_gdp=storage(
                f"{JRC_ARDECO_DATASET['url']}/ARDECO-SUVGDP.2021.table.csv"
            ),
            ardeco_pop=storage(
                f"{JRC_ARDECO_DATASET['url']}/ARDECO-SNPTD.2021.table.csv"
            ),
        output:
            ardeco_gdp=f"{JRC_ARDECO_DATASET['folder']}/ARDECO-SUVGDP.2021.table.csv",
            ardeco_pop=f"{JRC_ARDECO_DATASET['folder']}/ARDECO-SNPTD.2021.table.csv",
        message:
            "Retrieving JRC ARDECO data"
        run:
            for key in input.keys():
                copy2(input[key], output[key])



if (AQUIFER_DATA_DATASET := dataset_version("aquifer_data"))["source"] in [
    "primary",
    *ARCHIVE_SOURCES,
]:

    rule retrieve_aquifer_data_bgr:
        input:
            zip_file=storage(AQUIFER_DATA_DATASET["url"]),
        output:
            zip_file=f"{AQUIFER_DATA_DATASET['folder']}/ihme1500_aquif_ec4060_v12_poly.zip",
            aquifer_shapes=expand(
                f"{AQUIFER_DATA_DATASET['folder']}/IHME1500_v12/shp/ihme1500_aquif_ec4060_v12_poly.{{ext}}",
                ext=[
                    "shp",
                    "shx",
                    "dbf",
                    "cpg",
                    "prj",
                    "sbn",
                    "sbx",
                ],
            ),
        message:
            "Retrieving BGR aquifer data"
        run:
            copy2(input["zip_file"], output["zip_file"])
            unpack_archive(
                output["zip_file"],
                AQUIFER_DATA_DATASET["folder"],
            )


if (DH_AREAS_DATASET := dataset_version("dh_areas"))["source"] in [
    "primary",
    *ARCHIVE_SOURCES,
]:

    rule retrieve_dh_areas:
        input:
            dh_areas=storage(DH_AREAS_DATASET["url"]),
        output:
            dh_areas=f"{DH_AREAS_DATASET['folder']}/dh_areas.gpkg",
        log:
            "logs/retrieve_dh_areas.log",
        message:
            "Retrieving District Heating areas"
        run:
            copy2(input["dh_areas"], output["dh_areas"])


if (MOBILITY_PROFILES_DATASET := dataset_version("mobility_profiles"))[
    "source"
] in ARCHIVE_SOURCES:

    rule retrieve_mobility_profiles:
        input:
            kfz=storage(MOBILITY_PROFILES_DATASET["url"] + "/kfz.csv"),
            pkw=storage(MOBILITY_PROFILES_DATASET["url"] + "/pkw.csv"),
        output:
            kfz=f"{MOBILITY_PROFILES_DATASET['folder']}/kfz.csv",
            pkw=f"{MOBILITY_PROFILES_DATASET['folder']}/pkw.csv",
        log:
            "logs/retrieve_mobility_profiles.log",
        benchmark:
            "benchmarks/performances/retrieve_mobility_profiles"
        threads: 1
        resources:
            mem_mb=1000,
        message:
            "Retrieving mobility profiles data"
        run:
            copy2(input["kfz"], output["kfz"])
            copy2(input["pkw"], output["pkw"])