cba.smk

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

import logging
import os
import shutil
from pathlib import Path
from zipfile import ZipFile

import pandas as pd

from scripts.cba._helpers import filter_projects_by_specs
from scripts._helpers import fill_wildcards
from shutil import unpack_archive, copy2

logger = logging.getLogger(__name__)


wildcard_constraints:
    cba_project=r"(s|t)\d+",
    run="(?!None)[-a-zA-Z0-9]+",  # Disallow None as a run wildcard


# Retrieve
##########

if (CBA_PROJECTS_DATASET := dataset_version("tyndp_cba_projects"))[
    "source"
] in ARCHIVE_SOURCES:

    rule retrieve_tyndp_cba_projects:
        input:
            zip_file=storage(CBA_PROJECTS_DATASET["url"]),
        output:
            dir=directory(CBA_PROJECTS_DATASET["folder"]),
        log:
            "logs/retrieve_tyndp_cba_projects.log",
        params:
            source="CBA project explorer",
        run:
            copy2(input["zip_file"], output["dir"] + ".zip")
            unpack_archive(output["dir"] + ".zip", output["dir"])
            os.remove(output["dir"] + ".zip")


if (CBA_NON_CO2_DATASET := dataset_version("tyndp_cba_non_co2_emissions"))[
    "source"
] in ARCHIVE_SOURCES:

    rule retrieve_tyndp_cba_non_co2_emissions:
        input:
            file=storage(CBA_NON_CO2_DATASET["url"]),
        output:
            file=f"{CBA_NON_CO2_DATASET['folder']}/a.3_non-co2-emissions.csv",
        log:
            "logs/retrieve_tyndp_cba_non_co2_emissions.log",
        run:
            copy2(input["file"], output["file"])


if (CBA_GUIDELINES_DATASET := dataset_version("cba_guidelines_reference_projects"))[
    "source"
] in ARCHIVE_SOURCES:

    rule retrieve_cba_guidelines_reference_projects:
        input:
            file=storage(CBA_GUIDELINES_DATASET["url"]),
        output:
            file=f"{CBA_GUIDELINES_DATASET['folder']}/table_B1_CBA_Implementations_Guidelines_TYNDP2024.csv",
        log:
            "logs/retrieve_cba_guidelines_reference_projects.log",
        run:
            copy2(input["file"], output["file"])


def _effective_horizon(h, warn_fn=None, msg=None):
    if h not in [2030, 2040]:
        if warn_fn:
            warn_fn(msg or "Using 2040 for unsupported planning horizon %s.", h)
        return 2040
    return h


def get_run_name(w):
    """
    Return a single run name from Snakemake wildcards or config.

    There are instances when the run name given in the config is a string
    or a list. This helper is used to ensure we always get a single string
    (or empty string if no run is given) to work with in the rules.
    """
    run = w.get("run", config_provider("run", "name")(w))
    if isinstance(run, list):
        return run[0] if run else ""
    return run or ""


def presolved_sb_network_path(w, horizon=None):
    sb_version = config_provider(
        "cba", "cba_scenario_input", "sb_version", default="latest"
    )(w)
    target_horizon = horizon if horizon is not None else w.planning_horizons
    return RESULTS + f"networks/presolved-{sb_version}/base_s_all___{target_horizon}.nc"


if config.get("cba", {}).get("cba_scenario_input", {}).get("use_presolved", False):
    sb_version = (
        config.get("cba", {}).get("cba_scenario_input", {}).get("sb_version", "latest")
    )
    if (
        SB_SOLVED_NETWORKS_DATASET := dataset_version(
            "open_tyndp_prelim", version=sb_version
        )
    )["source"] in ARCHIVE_SOURCES:

        rule retrieve_presolved_sb_networks:
            input:
                zip_file=storage(SB_SOLVED_NETWORKS_DATASET["url"]),
            output:
                network=RESULTS
                + f"networks/presolved-{config['cba']['cba_scenario_input']['sb_version']}/base_s_all___{{planning_horizons}}.nc",
            log:
                logs("cba/retrieve_presolved_sb_networks_{planning_horizons}.log"),
            run:
                target_suffix = (
                    f"networks/base_s_all___{wildcards.planning_horizons}.nc"
                )
                with ZipFile(input["zip_file"], "r") as zf:
                    matches = [
                        m for m in zf.namelist() if m.endswith(target_suffix)
                    ]
                    if not matches:
                        raise ValueError(
                            f"Could not find '{target_suffix}' in {input['zip_file']}."
                        )
                    member = matches[0]
                    out_path = Path(output["network"])
                    out_path.parent.mkdir(parents=True, exist_ok=True)
                    with zf.open(member) as src, out_path.open("wb") as dst:
                        shutil.copyfileobj(src, dst)



# Build MSV
############


# read in transmission and storage projects from excel sheets
#
def input_clustered_network(w):
    scenario = config_provider("scenario")(w)
    (clusters,) = scenario["clusters"]
    return fill_wildcards(rules.cluster_network.output.network, clusters=clusters)


checkpoint clean_projects:
    input:
        dir=rules.retrieve_tyndp_cba_projects.output.dir,
        buses=rules.retrieve_tyndp.output.nodes,
        guidelines=rules.retrieve_cba_guidelines_reference_projects.output.file,
    output:
        # TODO: The toot_projects and pint_projects outputs are likely only
        # transmission projects (no storage). In order to confirm, we should check
        # if Table B.1 from the guidelines (table_B1_CBA_Implementations_Guidelines_TYNDP2024.csv)
        # contains only transmission or also storage projects.
        transmission_projects=resources("cba/transmission_projects.csv"),
        storage_projects=resources("cba/storage_projects.csv"),
        methods=resources("cba/cba_project_methods.csv"),
    log:
        logs("cba/clean_projects.log"),
    benchmark:
        benchmarks("performances/cba/clean_projects")
    script:
        scripts("cba/clean_projects.py")


rule clean_tyndp_indicators:
    input:
        dir=rules.retrieve_tyndp_cba_projects.output.dir,
    output:
        indicators=resources("cba/tyndp_indicators.csv"),
        readme=resources("cba/tyndp_indicators_name_unit.csv"),
    log:
        logs("cba/clean_tyndp_indicators.log"),
    benchmark:
        benchmarks("performances/cba/clean_tyndp_indicators")
    script:
        scripts("cba/clean_tyndp_indicators.py")


def input_sb_network(w, run=None):
    scenario = config_provider("scenario")(w)
    (clusters,) = scenario["clusters"]
    (opts,) = scenario["opts"]
    (sector_opts,) = scenario["sector_opts"]

    if config_provider("cba", "cba_scenario_input", "use_presolved", default=False)(w):
        scenario_name = config_provider("tyndp_scenario")(w)
        if scenario_name != "NT":
            raise ValueError(
                "Pre-solved SB networks are only currently available for the NT scenario."
            )
        # Check that options match the pre-solved network naming convention
        if clusters != "all" or opts != "" or sector_opts != "":
            raise ValueError(
                "Pre-solved SB runs require scenario.clusters=['all'], "
                "scenario.opts=[''], and scenario.sector_opts=[''] to match "
                "the Zenodo network naming (base_s_all___{planning_horizons}.nc)."
            )
        horizon = _effective_horizon(
            int(w.planning_horizons),
            warn_fn=logger.warning,
            msg=(
                "Pre-solved SB networks are only available for 2030 and 2040. "
                "Falling back to 2040 for CBA planning horizon %s."
            ),
        )
        return presolved_sb_network_path(w, horizon)

    expanded_wildcards = {
        "clusters": clusters,
        "opts": opts,
        "sector_opts": sector_opts,
    }
    if run is not None:
        expanded_wildcards["run"] = run

    match config_provider("foresight")(w):
        case "perfect":
            expanded_wildcards["planning_horizons"] = "all"
        case "myopic":
            expanded_wildcards["planning_horizons"] = _effective_horizon(
                int(w.planning_horizons),
                warn_fn=logger.warning,
                msg=(
                    "CBA planning horizon %s is not supported for SB inputs. "
                    "Using 2040 inputs instead."
                ),
            )
            # converts the same value to a string so fill_wildcards() can safely call .replace().
            expanded_wildcards["planning_horizons"] = str(
                expanded_wildcards["planning_horizons"]
            )
        case _:
            raise ValueError('config["foresight"] must be one of "perfect" or "myopic"')

    return fill_wildcards(
        RESULTS
        + "networks/base_s_{clusters}_{opts}_{sector_opts}_{planning_horizons}.nc",
        **expanded_wildcards,
    )


# Simplify scenario building network for CBA
# Fixes capacities, adds hurdle costs, extends primary fuel sources, disables volume limits
rule simplify_sb_network:
    input:
        network=input_sb_network,
        sb_network=lambda w: input_sb_network(
            w, run=config_provider("cba", "sb_scenario")(w)
        ),
    output:
        network=resources("cba/networks/simple_{planning_horizons}.nc"),
    log:
        logs("cba/simplify_sb_network_{planning_horizons}.log"),
    benchmark:
        benchmarks("performances/cba/simplify_sb_network_{planning_horizons}")
    params:
        tyndp_conventional_carriers=config_provider(
            "electricity", "tyndp_conventional_carriers"
        ),
        hurdle_costs=config_provider("cba", "hurdle_costs"),
    script:
        scripts("cba/simplify_sb_network.py")


# build reference corrections between SB investments and CBA guidelines
def get_elec_project_build_years(w):
    return config_provider("tyndp_investment_candidates", "elec_projects")(w).get(
        int(w.planning_horizons)
    )


rule fix_reference_sb_to_cba:
    input:
        invest_grid=rules.retrieve_tyndp.output.invest_grid,
        guidelines=rules.retrieve_cba_guidelines_reference_projects.output.file,
        transmission_projects=rules.clean_projects.output.transmission_projects,
        buses=rules.build_tyndp_network.output.substations_geojson,
    output:
        corrections=resources("cba/reference_sb_to_cba_{planning_horizons}.csv"),
    log:
        logs("cba/fix_reference_sb_to_cba_{planning_horizons}.log"),
    benchmark:
        benchmarks("performances/cba/fix_reference_sb_to_cba_{planning_horizons}")
    params:
        build_years=get_elec_project_build_years,
    script:
        scripts("cba/fix_reference_sb_to_cba.py")


# Build reference network with all TOOT projects included
# Ensures MSV extraction and rolling horizon use the same baseline
rule prepare_reference:
    input:
        network=rules.simplify_sb_network.output.network,
        transmission_projects=rules.clean_projects.output.transmission_projects,
        storage_projects=rules.clean_projects.output.storage_projects,
        corrections=rules.fix_reference_sb_to_cba.output.corrections,
        costs=resources("costs_{planning_horizons}_processed.csv"),
    output:
        network=resources("cba/networks/reference_{planning_horizons}.nc"),
    log:
        logs("cba/prepare_reference_{planning_horizons}.log"),
    benchmark:
        benchmarks("performances/cba/prepare_reference_{planning_horizons}")
    params:
        hurdle_costs=config_provider("cba", "hurdle_costs"),
        patch_sb_with_annexe=config_provider(
            "tyndp_investment_candidates", "patch_sb_with_annexe"
        ),
    script:
        scripts("cba/prepare_reference.py")


# Generate snapshot weightings for MSV extraction temporal aggregation
rule build_msv_snapshot_weightings:
    input:
        network=rules.prepare_reference.output.network,
    output:
        snapshot_weightings=resources(
            "cba/msv_snapshot_weightings_{planning_horizons}.csv"
        ),
    log:
        logs("cba/build_msv_snapshot_weightings_{planning_horizons}.log"),
    benchmark:
        benchmarks("performances/cba/build_msv_snapshot_weightings_{planning_horizons}")
    params:
        msv_resolution=config_provider("cba", "msv_extraction", "resolution"),
        drop_leap_day=config_provider("enable", "drop_leap_day"),
    script:
        scripts("cba/build_msv_snapshot_weightings.py")


def input_msv_snapshot_weightings(w):
    """Return snapshot weightings file only if MSV resolution requires it."""
    resolution = config_provider("cba", "msv_extraction", "resolution")(w)
    if resolution and "h" in str(resolution).lower():
        return rules.build_msv_snapshot_weightings.output.snapshot_weightings
    return []


# Extract marginal storage values via perfect foresight solve (full year with cyclicity enabled)
rule solve_cba_msv_extraction:
    input:
        network=rules.prepare_reference.output.network,
        snapshot_weightings=input_msv_snapshot_weightings,
    output:
        network=resources("cba/networks/msv_{planning_horizons}.nc"),
    log:
        solver=RESULTS + "logs/cba/msv/{planning_horizons}_solver.log",
        memory=RESULTS + "logs/cba/msv/{planning_horizons}_memory.log",
        python=RESULTS + "logs/cba/msv/{planning_horizons}_python.log",
    benchmark:
        RESULTS + "benchmarks/performances/cba/msv/{planning_horizons}"
    threads: solver_threads
    params:
        solving=config_provider("solving"),
        cba_solving=config_provider("cba", "msv_extraction", "solving"),
        msv_resolution=config_provider("cba", "msv_extraction", "resolution"),
        cyclic_carriers=config_provider("cba", "storage", "cyclic_carriers"),
    script:
        scripts("cba/solve_cba_msv_extraction.py")


# Build rolling horizon
#######################


# Prepare network for rolling horizon: disable seasonal cyclicity, apply marginal storage value
rule prepare_rolling_horizon:
    input:
        network=rules.prepare_reference.output.network,
        network_msv=rules.solve_cba_msv_extraction.output.network,
    output:
        network=resources("cba/networks/rl_{planning_horizons}.nc"),
    log:
        logs("cba/prepare_rolling_horizon_{planning_horizons}.log"),
    benchmark:
        benchmarks("performances/cba/prepare_rolling_horizon_{planning_horizons}")
    params:
        cyclic_carriers=config_provider("cba", "storage", "cyclic_carriers"),
        soc_boundary_carriers=config_provider("cba", "storage", "soc_boundary_carriers"),
        msv_resample_method=config_provider("cba", "msv_extraction", "resample_method"),
    script:
        scripts("cba/prepare_rolling_horizon.py")


# add or remove the cba project based on assigned method
rule prepare_project:
    input:
        network=rules.prepare_rolling_horizon.output.network,
        network_msv=rules.solve_cba_msv_extraction.output.network,
        transmission_projects=rules.clean_projects.output.transmission_projects,
        storage_projects=rules.clean_projects.output.storage_projects,
        methods=rules.clean_projects.output.methods,
        costs=resources("costs_{planning_horizons}_processed.csv"),
    output:
        network=temp(
            resources("cba/networks/project_{cba_project}_{planning_horizons}.nc")
        ),
    log:
        logs("cba/prepare_project_{cba_project}_{planning_horizons}.log"),
    benchmark:
        benchmarks("performances/cba/prepare_project_{cba_project}_{planning_horizons}")
    params:
        hurdle_costs=config_provider("cba", "hurdle_costs"),
        cyclic_carriers=config_provider("cba", "storage", "cyclic_carriers"),
        soc_boundary_carriers=config_provider("cba", "storage", "soc_boundary_carriers"),
    script:
        scripts("cba/prepare_project.py")


# Solve reference network with rolling horizon (MSV already applied)
rule solve_cba_reference_network:
    input:
        network=rules.prepare_rolling_horizon.output.network,
    output:
        network=RESULTS + "cba/networks/reference_{planning_horizons}.nc",
    log:
        solver=RESULTS + "logs/cba/reference/reference_{planning_horizons}_solver.log",
        memory=RESULTS + "logs/cba/reference/reference_{planning_horizons}_memory.log",
        python=RESULTS + "logs/cba/reference/reference_{planning_horizons}_python.log",
    benchmark:
        RESULTS + "benchmarks/performances/cba/reference/reference_{planning_horizons}"
    threads: 1
    params:
        solving=config_provider("solving"),
        cba_solving=config_provider("cba", "solving"),
        foresight=config_provider("foresight"),
        time_resolution=config_provider("clustering", "temporal", "resolution_sector"),
        custom_extra_functionality=None,
    script:
        scripts("cba/solve_cba_network.py")


# Solve TOOT/PINT project network with rolling horizon
rule solve_cba_network:
    input:
        network=resources("cba/networks/project_{cba_project}_{planning_horizons}.nc"),
    output:
        network=RESULTS + "cba/networks/project_{cba_project}_{planning_horizons}.nc",
    log:
        solver=RESULTS
        + "logs/cba/projects/project_{cba_project}_{planning_horizons}_solver.log",
        memory=RESULTS
        + "logs/cba/projects/project_{cba_project}_{planning_horizons}_memory.log",
        python=RESULTS
        + "logs/cba/projects/project_{cba_project}_{planning_horizons}_python.log",
    benchmark:
        RESULTS
        + "benchmarks/performances/cba/projects/project_{cba_project}_{planning_horizons}"
    threads: 1
    params:
        solving=config_provider("solving"),
        cba_solving=config_provider("cba", "solving"),
        foresight=config_provider("foresight"),
        time_resolution=config_provider("clustering", "temporal", "resolution_sector"),
        custom_extra_functionality=None,
    script:
        scripts("cba/solve_cba_network.py")


# Postprocess
##############


# Compute CBA indicators comparing reference and project networks
rule make_indicators:
    input:
        reference=RESULTS + "cba/networks/reference_{planning_horizons}.nc",
        project=RESULTS + "cba/networks/project_{cba_project}_{planning_horizons}.nc",
        non_co2_emissions=rules.retrieve_tyndp_cba_non_co2_emissions.output.file,
        benchmark=rules.clean_tyndp_indicators.output.indicators,
        methods=rules.clean_projects.output.methods,
    output:
        indicators=RESULTS
        + "cba/results/{planning_horizons}/project_{cba_project}_{planning_horizons}.csv",
    log:
        logs("cba/make_indicators_{cba_project}_{planning_horizons}.log"),
    benchmark:
        benchmarks("performances/cba/make_indicators_{cba_project}_{planning_horizons}")
    script:
        scripts("cba/make_indicators.py")


def input_indicators(w):
    """
    List all indicators csv

    Works for collection scenarios and regular scenarios.
    """
    run = get_run_name(w)
    projects = pd.read_csv(checkpoints.clean_projects.get(run=run).output.methods)
    horizon = _effective_horizon(
        int(w.planning_horizons),
        warn_fn=logger.warning,
        msg=(
            "CBA methods are only available for 2030 or 2040. "
            "Using 2040 for planning horizon %s."
        ),
    )
    if "planning_horizon" in projects.columns:
        projects = projects.loc[projects["planning_horizon"] == horizon]
    cba_projects = [f"t{pid}" for pid in projects["project_id"].unique()]

    # Collection scenarios look for results within nested source runs,
    # regular scenarios look within their own run.
    runs = cba_source_runs(w)
    # NOTE: project_specs filtering happens on the collection scenario (and does not descend)
    project_specs = config_provider("cba", "projects")(w)

    return expand(
        rules.make_indicators.output.indicators,
        cba_project=filter_projects_by_specs(cba_projects, project_specs),
        planning_horizons=[w.planning_horizons],
        run=runs,
    )


# Combine indicators for all projects into overview CSV
rule combine_indicators:
    input:
        indicators=input_indicators,
    output:
        indicators=RESULTS
        + "cba/results/{planning_horizons}/indicators_{planning_horizons}.csv",
    log:
        logs("cba/combine_indicators_{planning_horizons}.log"),
    benchmark:
        benchmarks("performances/cba/combine_indicators_{planning_horizons}")
    script:
        scripts("cba/combine_indicators.py")


rule plot_indicators:
    input:
        indicators=rules.combine_indicators.output.indicators,
        transmission_projects=rules.clean_projects.output.transmission_projects,
    output:
        plot_dir=directory(RESULTS + "cba/graphs/{planning_horizons}/summary"),
    log:
        logs("cba/plot_indicators_{planning_horizons}.log"),
    benchmark:
        benchmarks("performances/cba/plot_indicators_{planning_horizons}")
    params:
        plotting=config_provider("plotting"),
    script:
        scripts("cba/plot_indicators.py")


# Benchmarking
###############


rule plot_cba_benchmark:
    input:
        indicators=rules.make_indicators.output.indicators,
    output:
        plot_file=RESULTS
        + "cba/validation/{planning_horizons}/project_{cba_project}_{planning_horizons}.png",
    log:
        logs("cba/plot_cba_benchmark_{cba_project}_{planning_horizons}.log"),
    benchmark:
        benchmarks(
            "performances/cba/plot_cba_benchmark_{cba_project}_{planning_horizons}"
        )
    script:
        scripts("cba/plot_benchmark_indicators.py")


# rule plot_all_cba_benchmark:
#     input:
#         indicators=rules.combine_indicators.output.indicators,
#     output:
#         plot_dir=directory(RESULTS + "cba/validation/{planning_horizons}"),
#     script:
#         scripts("cba/plot_benchmark_indicators.py")


rule plot_weather_benchmark:
    input:
        # indicators=rules.combine_indicators.output.indicators,
        indicators=rules.make_indicators.output.indicators,
    output:
        plot_file=RESULTS
        + "cba/graphs/{planning_horizons}/ensemble_{cba_project}_{planning_horizons}.png",
    log:
        logs("cba/plot_weather_benchmark_{cba_project}_{planning_horizons}.log"),
    benchmark:
        benchmarks(
            "performances/cba/plot_weather_benchmark_{cba_project}_{planning_horizons}"
        )
    script:
        scripts("cba/plot_benchmark_indicators.py")


rule average_indicators_per_project_and_planning_horizon:
    input:
        indicators=lambda w: expand(
            rules.make_indicators.output.indicators,
            cba_project=[w.cba_project],
            planning_horizons=[w.planning_horizons],
            run=cba_source_runs(w),
        ),
    output:
        indicators=RESULTS
        + "cba/results/all/ensemble_indicators_{cba_project}_{planning_horizons}.csv",
    log:
        logs("cba/average_indicators_{cba_project}_{planning_horizons}.log"),
    benchmark:
        benchmarks(
            "performances/cba/average_indicators_{cba_project}_{planning_horizons}"
        )
    script:
        scripts("cba/average_indicators.py")


rule summarize_indicators_per_project:
    input:
        indicators=lambda w: expand(
            rules.average_indicators_per_project_and_planning_horizon.output.indicators,
            planning_horizons=config["cba"]["planning_horizons"],
            cba_project=[w.cba_project],
            run=[w.run],
        ),
    output:
        plot_file=RESULTS + "cba/graphs/all/ensemble_{cba_project}_all_horizons.png",
    log:
        logs("cba/summarize_indicators_{cba_project}.log"),
    benchmark:
        benchmarks("performances/cba/summarize_indicators_{cba_project}")
    script:
        scripts("cba/summarize_indicators.py")


def summary_benchmark_indicators(w):
    """
    Returns Indicator CSVs as inputs for the per-horizon summary benchmark plot.
    If collection scenarios, returns the weighted-average ensemble indicators CSV as inputs for plotting.
    """
    run = get_run_name(w)
    if run in cba_collection_scenarios(w):
        return expand(
            rules.average_indicators_per_project_and_planning_horizon.output.indicators,
            planning_horizons=[w.planning_horizons],
            cba_project=cba_projects(w),
            run=[run],
        )
    return expand(
        rules.combine_indicators.output.indicators,
        planning_horizons=[w.planning_horizons],
        run=[run],
    )


rule plot_summary_projects_benchmark:
    input:
        indicators=summary_benchmark_indicators,
    output:
        plot_file=RESULTS
        + "cba/validation/{planning_horizons}/summary_benchmark_{planning_horizons}.png",
    log:
        logs("cba/plot_summary_projects_benchmark_{planning_horizons}.log"),
    benchmark:
        benchmarks(
            "performances/cba/plot_summary_projects_benchmark_{planning_horizons}"
        )
    script:
        scripts("cba/plot_benchmark_indicators.py")


rule summarize_all_indicators:
    input:
        indicators=lambda w: expand(
            rules.plot_weather_benchmark.input.indicators,
            planning_horizons=config["cba"]["planning_horizons"],
            cba_project=cba_projects(w),
            run=cba_source_runs(w),
        ),
    output:
        plot_file=RESULTS + "cba/graphs/all/ensemble_all.png",
    log:
        logs("cba/summarize_all_indicators.log"),
    benchmark:
        benchmarks("performances/cba/summarize_all_indicators")
    script:
        scripts("cba/summarize_all.py")


def cba_target_runs(w):
    """
    Return runs explicitly requested through config["run"]["name"].
    """
    names = config["run"]["name"]
    if isinstance(names, str):
        return [names]
    return names


def cba_collection_scenarios(w):
    """
    Return actual cba collection scenarios, i.e. runs with cba.scenarios.
    """
    scenarios = []
    for name in cba_target_runs(w):
        try:
            scn = scenario_config(name)
        except KeyError:
            continue
        if scn.get("cba", {}).get("scenarios"):
            scenarios.append(name)
    return scenarios


def cba_source_runs(w):
    """
    Return the runs that provide CBA project indicator CSVs.

    Collection scenarios read from their nested cba.scenarios;
    regular runs read from their own run name.
    """
    run = get_run_name(w)
    if not run:
        return []
    try:
        scn = scenario_config(run)
    except KeyError:
        return [run]
    runs = scn.get("cba", {}).get("scenarios")
    return runs if runs else [run]


def cba_scenarios(w):
    """
    Return cba: scenarios of a cba collection scenario
    """
    run = get_run_name(w)
    try:
        scn = scenario_config(run)
    except KeyError:
        return [run] if run else []
    return scn.get("cba", {}).get("scenarios", [run] if run else [])


def cba_projects_run(w):
    """
    Return the run from which project methods should be read.
    """
    runs = cba_source_runs(w)
    return runs[0] if runs else ""


def validate_cba_project_methods_consistency(w):
    """
    Validate that nested runs of a collection scenario share the same project ids
    and CBA methods across planning horizons.
    """
    runs = cba_source_runs(w)
    if len(runs) <= 1:
        return

    reference_run = runs[0]
    reference = pd.read_csv(
        checkpoints.clean_projects.get(run=reference_run).output.methods
    )
    ref_cols = ["project_id", "planning_horizon", "method"]
    reference = reference[ref_cols].sort_values(ref_cols).reset_index(drop=True)

    for run in runs[1:]:
        current = pd.read_csv(checkpoints.clean_projects.get(run=run).output.methods)
        current = current[ref_cols].sort_values(ref_cols).reset_index(drop=True)
        if not reference.equals(current):
            raise ValueError(
                "CBA project methods differ across nested collection runs. "
                f"Reference run '{reference_run}' does not match '{run}'."
            )


def cba_projects(w):
    """
    List all indicators csv
    """
    validate_cba_project_methods_consistency(w)

    run = cba_projects_run(w)
    projects = pd.read_csv(checkpoints.clean_projects.get(run=run).output.methods)
    cba_projects = [f"t{pid}" for pid in projects["project_id"].unique()]
    project_specs = config_provider("cba", "projects")(w)
    cba_project = filter_projects_by_specs(cba_projects, project_specs)

    return expand(cba_project)


def collect_cba_scenario_inputs(w):
    inputs = []
    inputs.extend(
        expand(
            rules.plot_indicators.output.plot_dir,
            planning_horizons=config_provider("cba", "planning_horizons")(w),
            run=cba_scenarios(w),
        )
    )
    inputs.extend(
        expand(
            rules.plot_cba_benchmark.output.plot_file,
            planning_horizons=config_provider("cba", "planning_horizons")(w),
            cba_project=cba_projects(w),
            run=cba_scenarios(w),
        )
    )
    inputs.extend(
        expand(
            rules.plot_summary_projects_benchmark.output.plot_file,
            planning_horizons=config_provider("cba", "planning_horizons")(w),
            run=cba_scenarios(w),
        )
    )

    run = get_run_name(w)
    if run in cba_collection_scenarios(w):
        inputs.extend(
            expand(
                rules.plot_weather_benchmark.output.plot_file,
                planning_horizons=config_provider("cba", "planning_horizons")(w),
                cba_project=cba_projects(w),
                run=cba_scenarios(w),
            )
        )

    return inputs


# Collect
##########


# collect files to be stored in the scenario directory, e.g., NT-cy1995
rule collect_cba_scenario:
    input:
        collect_cba_scenario_inputs,
    output:
        touch(RESULTS + "cba/all_scenarios.txt"),
    log:
        logs("cba/collect_cba_scenario.log"),
    benchmark:
        benchmarks("performances/cba/collect_cba_scenario")


def cba_ensemble_inputs(w):
    runs = cba_collection_scenarios(w)
    if not runs:
        return []

    inputs = []
    inputs.extend(
        expand(
            rules.average_indicators_per_project_and_planning_horizon.output.indicators,
            planning_horizons=config["cba"]["planning_horizons"],
            cba_project=cba_projects(w),
            run=runs,
        )
    )
    inputs.extend(
        expand(
            rules.summarize_indicators_per_project.output.plot_file,
            cba_project=cba_projects(w),
            run=runs,
        )
    )
    inputs.extend(
        expand(
            rules.summarize_all_indicators.output.plot_file,
            planning_horizons=config["cba"]["planning_horizons"],
            cba_project=cba_projects(w),
            run=runs,
        )
    )
    inputs.extend(
        expand(
            rules.plot_summary_projects_benchmark.output.plot_file,
            planning_horizons=config["cba"]["planning_horizons"],
            run=runs,
        )
    )
    return inputs


# collect files to be stored in the scenario collection directory, e.g., NT-cyears
rule cba:
    input:
        cba_ensemble_inputs,
        # lambda w: expand(
        #     rules.plot_all_cba_benchmark.output.plot_dir,
        #     planning_horizons=config["cba"]["planning_horizons"],
        #     run=cba_collection_scenarios(w),
        # ),
        # collect files to be stored in the scenario directory, e.g., NT-cy1995
        lambda w: expand(
            rules.collect_cba_scenario.output[0],
            run=cba_target_runs(w),
        ),


# collect rules
rule prepare_references:
    input:
        lambda w: expand(
            resources("cba/networks/reference_{planning_horizons}.nc"),
            **config["scenario"],
            run=config["run"]["name"],
        ),