# tests/data/resource_groups_base/resource_group_stats.py

"""
Create a summary file with the resource group capacity by lcoe bin for every region
"""

import pandas as pd
import numpy as np
from pathlib import Path
from powergenome.util import regions_to_keep, map_agg_region_names, load_settings

settings_path = None

if settings_path:
    settings = load_settings(settings_path)
    keep_regions, region_agg_map = regions_to_keep(settings)


LCOE_BIN_EDGES = [0, 10, 20, 30, 40, 50, 60, 70, 80, 90, 100, 125, 150, 175, 200, 5000]

meta_files = list(Path.cwd().glob("*metadata.csv"))

name_len = len(meta_files[0].stem.split("_"))
if "offshorewind" in meta_files[0].stem:
    resource_position = name_len - 7
else:
    resource_position = name_len - 5


def resource_summary(df: pd.DataFrame) -> pd.DataFrame:
    # Loop though metros and only include base (highest number level) clusters, which
    # include all capacity.
    region_list = []
    for region, _df in df.groupby(["ipm_region", "metro_id"]):
        max_level = _df["level"].max()
        region_list.append(_df.query("level==@max_level"))

    refined_df = pd.concat(region_list, ignore_index=True)
    refined_df["lcoe_bin"] = pd.cut(refined_df["lcoe"], bins=LCOE_BIN_EDGES)
    stats = refined_df.groupby(["ipm_region", "lcoe_bin"])["mw"].sum().unstack()
    stats.columns = stats.columns.astype(str)  # Convert from categorical index to str
    stats["total_mw"] = stats.sum(axis=1)
    return stats.reset_index()


df_list = []
for m_f in meta_files:
    resource = m_f.stem.split("_")[resource_position]
    if resource == "existing":
        continue
    if resource == "offshorewind":
        anchor = m_f.stem.split("_")[resource_position + 4]
        pref = m_f.stem.split("_")[resource_position + 5]
        resource = f"{resource}_{anchor}_{pref}"
    df = pd.read_csv(m_f)

    stats = resource_summary(df)
    cols = list(stats.columns)
    stats["resource"] = resource

    # Make sure capacicity numbers match
    assert np.allclose(stats["total_mw"].sum(), df.query("level == 1")["mw"].sum())

    df_list.append(stats[["resource"] + cols])

results_df = pd.concat(df_list, ignore_index=True).fillna(0)

if settings_path:
    results_df = map_agg_region_names(
        results_df,
        region_agg_map=region_agg_map,
        original_col_name="ipm_region",
        new_col_name="model_region",
    )

    results_df = results_df.groupby(["model_region", "resource"], as_index=False).sum()


results_df.to_csv("resource_summary_stats.csv", index=False)


# tests/data/create_data_from_pudl.py

import sqlite3
import sqlalchemy as sa

import pandas as pd
from powergenome.params import DATA_PATHS
from powergenome.util import init_pudl_connection

GENS860_COLS = [
    "report_date",
    "plant_id_eia",
    "generator_id",
    # "associated_combined_heat_power",
    # "balancing_authority_code_eia",
    # "bypass_heat_recovery",
    "capacity_mw",
    # "county",
    "current_planned_operating_date",
    "energy_source_code_1",
    # "ferc_cogen_status",
    # "iso_rto_code",
    # "latitude",
    # "longitude",
    "minimum_load_mw",
    # "operating_date",
    "operational_status_code",
    # "original_planned_operating_date",
    # "state",
    "summer_capacity_mw",
    "technology_description",
    # "unit_id_pudl",
    "winter_capacity_mw",
    "fuel_type_code_pudl",
    # "zip_code",
    "planned_retirement_date",
    "time_cold_shutdown_full_load_code",
    "switch_oil_gas",
    "planned_new_capacity_mw",
    "energy_source_code_2",
    "region",
]
GEN_FUEL_COLS = [
    "report_date",
    "plant_id_eia",
    "energy_source_code",
    "fuel_consumed_for_electricity_mmbtu",
    "fuel_consumed_for_electricity_units",
    "fuel_consumed_mmbtu",
    "fuel_consumed_units",
    "fuel_mmbtu_per_unit",
    "net_generation_mwh",
    "prime_mover_code",
    "fuel_type_code_pudl",
]
ENTITY_COLS = ["plant_id_eia", "generator_id", "prime_mover_code", "operating_date"]


def create_testing_db():
    pudl_engine, pudl_out, pg_engine = init_pudl_connection(
        start_year=2018, end_year=2020
    )
    pudl_test_conn = sqlite3.connect(DATA_PATHS["test_data"] / "pudl_test_data.db")

    plant_region = pd.read_sql_table("plant_region_map_epaipm", pg_engine)
    # gens_860 = pudl_out.gens_eia860()
    s = "SELECT * from generators_eia860 where strftime('%Y',report_date)='2020'"
    gens_860 = pd.read_sql_query(s, pudl_engine, parse_dates=["report_date"])
    # gens_860 = gens_860.loc[gens_860.report_date.dt.year == 2020, :]
    gens_860 = pd.merge(gens_860, plant_region, on="plant_id_eia", how="inner")
    gens_860 = gens_860.loc[:, GENS860_COLS]
    gens_860 = gens_860.groupby(
        ["region", "technology_description"], as_index=False
    ).head(10)
    gens_860 = gens_860.drop(columns="region")
    eia_plant_ids = gens_860["plant_id_eia"].unique()

    gen_entity = pd.read_sql_table("generators_entity_eia", pudl_engine)
    gen_entity = gen_entity.loc[
        gen_entity["plant_id_eia"].isin(eia_plant_ids), ENTITY_COLS
    ]

    bga = pudl_out.bga_eia860()
    bga = bga.loc[
        (bga.report_date.dt.year == 2020) & (bga.plant_id_eia.isin(eia_plant_ids)), :
    ]

    s = "SELECT * from generation_fuel_eia923 where strftime('%Y',report_date)='2020'"
    gen_fuel = pd.read_sql_query(s, pudl_engine, parse_dates=["report_date"])
    gen_fuel = gen_fuel.loc[
        gen_fuel.plant_id_eia.isin(eia_plant_ids),
        GEN_FUEL_COLS,
    ]
    s = "SELECT * from generation_fuel_nuclear_eia923 where strftime('%Y',report_date)='2020'"
    gen_fuel_nuc = pd.read_sql_query(s, pudl_engine, parse_dates=["report_date"])
    gen_fuel_nuc = gen_fuel_nuc.loc[
        gen_fuel_nuc.plant_id_eia.isin(eia_plant_ids),
        GEN_FUEL_COLS,
    ]
    # gen_fuel = pd.concat([gen_fuel, gen_fuel_nuc], ignore_index=True)

    s = "SELECT * from generation_eia923 where strftime('%Y',report_date)='2020'"
    gen_923 = pd.read_sql_query(s, pudl_engine, parse_dates=["report_date"])
    gen_923 = gen_923.loc[
        gen_923.plant_id_eia.isin(eia_plant_ids),
        :,
    ]

    s = "SELECT * from boiler_fuel_eia923 where strftime('%Y',report_date)='2020'"
    boiler_fuel = pd.read_sql_query(s, pudl_engine, parse_dates=["report_date"])
    boiler_fuel = boiler_fuel.loc[
        boiler_fuel.plant_id_eia.isin(eia_plant_ids),
        :,
    ]

    plant_entity = pd.read_sql_table("plants_entity_eia", pudl_engine)
    plant_entity = plant_entity.loc[plant_entity["plant_id_eia"].isin(eia_plant_ids), :]
    plants_eia_860 = pd.read_sql_table("plants_eia860", pudl_engine)
    plants_eia_860 = plants_eia_860.loc[
        plants_eia_860["plant_id_eia"].isin(eia_plant_ids), :
    ]

    plants_eia = pd.read_sql_table("plants_eia", pudl_engine)
    plants_eia = plants_eia.loc[plants_eia["plant_id_eia"].isin(eia_plant_ids), :]

    utilities_eia = pd.read_sql_table("utilities_eia", pudl_engine)
    utilities_entity = pd.read_sql_table("utilities_entity_eia", pudl_engine)
    utilities_860 = pd.read_sql_table("utilities_eia860", pudl_engine)

    plant_region.to_sql(
        "plant_region_map_epaipm", pudl_test_conn, index=False, if_exists="replace"
    )
    gens_860.to_sql(
        "generators_eia860", pudl_test_conn, index=False, if_exists="replace"
    )
    gen_entity.to_sql(
        "generators_entity_eia", pudl_test_conn, index=False, if_exists="replace"
    )
    bga.to_sql(
        "boiler_generator_assn_eia860", pudl_test_conn, index=False, if_exists="replace"
    )
    gen_fuel.to_sql(
        "generation_fuel_eia923", pudl_test_conn, index=False, if_exists="replace"
    )
    gen_fuel_nuc.to_sql(
        "generation_fuel_nuclear_eia923",
        pudl_test_conn,
        index=False,
        if_exists="replace",
    )
    gen_923.to_sql(
        "generation_eia923", pudl_test_conn, index=False, if_exists="replace"
    )
    boiler_fuel.to_sql(
        "boiler_fuel_eia923", pudl_test_conn, index=False, if_exists="replace"
    )
    plant_entity.to_sql(
        "plants_entity_eia", pudl_test_conn, index=False, if_exists="replace"
    )
    plants_eia_860.to_sql(
        "plants_eia860", pudl_test_conn, index=False, if_exists="replace"
    )
    plants_eia.to_sql("plants_eia", pudl_test_conn, index=False, if_exists="replace")
    utilities_eia.to_sql(
        "utilities_eia", pudl_test_conn, index=False, if_exists="replace"
    )
    utilities_entity.to_sql(
        "utilities_entity_eia", pudl_test_conn, index=False, if_exists="replace"
    )
    utilities_860.to_sql(
        "utilities_eia860", pudl_test_conn, index=False, if_exists="replace"
    )


if __name__ == "__main__":
    create_testing_db()


# tests/data/create_slim_pg_misc.py

"""
Create a slimmed-down version of the pg_misc database for testing.
"""

import sqlite3

import pandas as pd
from powergenome.params import DATA_PATHS
from powergenome.util import init_pudl_connection


def create_testing_db():
    pudl_engine, pudl_out, pg_engine = init_pudl_connection(
        start_year=2018, end_year=2020
    )

    pg_test_conn = sqlite3.connect(DATA_PATHS["test_data"] / "pg_misc_tables.sqlite3")

    # Read tables. Filter data where necessary
    s = """
        SELECT time_index, region_id_epaipm, load_mw, year
        FROM load_curves_ferc
        WHERE region_id_epaipm in ('WECC_AZ', 'WEC_CALN', 'WEC_BANC', 'WEC_LADW')
    """
    load_curves_ferc = pd.read_sql_query(s, pg_engine)

    s = """
        SELECT time_index, region, sector, subsector, load_mw, year
        FROM load_curves_nrel_efs
        WHERE region in ('WECC_AZ', 'WEC_CALN', 'WEC_BANC', 'WEC_LADW')
    """
    load_curves_efs = pd.read_sql_query(s, pg_engine)

    s = """
        SELECT *
        FROM offshore_spur_costs_nrelatb
        WHERE atb_year = 2022
        AND basis_year >= 2020
    """
    offshore_spur = pd.read_sql_query(s, pg_engine)

    s = """
        SELECT *
        FROM technology_costs_nrelatb
        WHERE atb_year = 2022
        AND parameter in ('capex_mw', 'capex_mwh', 'variable_o_m_mwh', 'fixed_o_m_mw', 'fixed_o_m_mwh', 'wacc_real')
    """
    tech_costs = pd.read_sql_query(s, pg_engine)

    s = """
        SELECT *
        FROM technology_heat_rates_nrelatb
        WHERE atb_year = 2022
    """
    tech_hr = pd.read_sql_query(s, pg_engine)

    plant_region_map = pd.read_sql_table("plant_region_map_epaipm", pg_engine)
    regions = pd.read_sql_table("regions_entity_epaipm", pg_engine)
    transmission = pd.read_sql_table("transmission_single_epaipm", pg_engine).drop(
        columns=["id", "tariff_mills_kwh"]
    )

    # Write tables to test db
    load_curves_ferc.to_sql(
        "load_curves_ferc", pg_test_conn, if_exists="replace", index=False
    )
    load_curves_efs.to_sql(
        "load_curves_nrel_efs", pg_test_conn, if_exists="replace", index=False
    )
    offshore_spur.to_sql(
        "offshore_spur_costs_nrelatb", pg_test_conn, if_exists="replace", index=False
    )
    tech_costs.to_sql(
        "technology_costs_nrelatb", pg_test_conn, if_exists="replace", index=False
    )
    tech_hr.to_sql(
        "technology_heat_rates_nrelatb", pg_test_conn, if_exists="replace", index=False
    )
    plant_region_map.to_sql(
        "plant_region_map_epaipm", pg_test_conn, if_exists="replace", index=False
    )
    regions.to_sql(
        "regions_entity_epaipm", pg_test_conn, if_exists="replace", index=False
    )
    transmission.to_sql(
        "transmission_single_epaipm", pg_test_conn, if_exists="replace", index=False
    )


if __name__ == "__main__":
    create_testing_db()


# bin/resource_group_stats.py

"""
Create a summary file with the resource group capacity by lcoe bin for every region
"""
import argparse
import sys
from typing import List
import pandas as pd
import numpy as np
from pathlib import Path
from powergenome.util import regions_to_keep, map_agg_region_names, load_settings
from powergenome.params import SETTINGS


def parse_command_line(argv):
    """
    Parse command line arguments. See the -h option.

    :param argv: arguments on the command line must include caller file name.
    """
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "-sf",
        "--settings_file",
        dest="settings_file",
        type=str,
        default=None,
        help="Specify a YAML settings file to aggregate capacity for model regions.",
    )
    parser.add_argument(
        "--sub_region",
        dest="sub_region",
        type=str,
        default=None,
        help="Name of the sub-region column (division below ipm region).",
    )
    arguments = parser.parse_args(argv[1:])
    return arguments


LCOE_BIN_EDGES = [0, 10, 20, 30, 40, 50, 60, 70, 80, 90, 100, 125, 150, 175, 200, 5000]


def resource_summary(df: pd.DataFrame, sub_region: str = None) -> pd.DataFrame:
    # Loop though metros and only include base (highest number level) clusters, which
    # include all capacity.
    by = ["ipm_region", "metro_id"]
    if sub_region:
        by.insert(-1, sub_region)
    region_list = []
    for region, _df in df.groupby(by):
        max_level = _df["level"].max()
        region_list.append(_df.query("level==@max_level"))

    region_by = ["ipm_region", "lcoe_bin"]
    if sub_region:
        region_by.insert(-1, sub_region)
    refined_df = pd.concat(region_list, ignore_index=True)
    refined_df["lcoe_bin"] = pd.cut(refined_df["lcoe"], bins=LCOE_BIN_EDGES)
    stats = refined_df.groupby(region_by)["mw"].sum().unstack()
    stats.columns = stats.columns.astype(str)  # Convert from categorical index to str
    stats["total_mw"] = stats.sum(axis=1)
    stats = stats.query("total_mw > 0")
    return stats.reset_index()


def find_resource_idx(s: str) -> int:
    name_list = s.split("_")
    for resource in ["offshorewind", "landbasedwind", "utilitypv"]:
        if resource in name_list:
            resource_idx = name_list.index(resource)
    return resource_idx


def main():
    args = parse_command_line(sys.argv)
    meta_files = list(Path(SETTINGS["RESOURCE_GROUPS"]).glob("*metadata.csv"))

    # This might fail with different naming structures of resource group files.
    # TODO: #139 Generalize finding resource type in metadata files.
    # name_len = len(meta_files[0].stem.split("_"))
    # if "offshorewind" in meta_files[0].stem:
    #     resource_position = name_len - 13
    # else:
    #     resource_position = name_len - 11
    resource_position = find_resource_idx(meta_files[0].stem)

    if args.settings_file:
        settings = load_settings(args.settings_file)
        keep_regions, region_agg_map = regions_to_keep(
            settings["model_regions"], settings.get("region_aggregations")
        )
    else:
        settings = {}

    df_list = []
    for m_f in meta_files:
        resource = m_f.stem.split("_")[resource_position]
        # from IPython import embed

        # embed()
        if resource == "existing":
            continue
        if resource == "offshorewind":
            anchor = None
            i = 0
            while not anchor and i <= len(m_f.stem.split("_")):
                i += 1
                if m_f.stem.split("_")[resource_position + i] in ["fixed", "floating"]:
                    anchor = m_f.stem.split("_")[resource_position + i]
            # anchor = m_f.stem.split("_")[resource_position + 4]
            pref = m_f.stem.split("_")[resource_position + i + 1]
            # pref = m_f.stem.split("_")[resource_position + 5]
            resource = f"{resource}_{anchor}_{pref}"
        df = pd.read_csv(m_f)

        stats = resource_summary(df, args.sub_region)
        # from IPython import embed

        # embed()
        cols = list(stats.columns)
        stats["resource"] = resource

        # Make sure capacicity numbers match
        assert np.allclose(stats["total_mw"].sum(), df.query("level == 1")["mw"].sum())

        df_list.append(stats[["resource"] + cols])

    results_df = pd.concat(df_list, ignore_index=True)
    results_df = results_df.fillna(0)

    if args.settings_file:
        results_df = map_agg_region_names(
            results_df,
            region_agg_map=region_agg_map,
            original_col_name="ipm_region",
            new_col_name="model_region",
        )
        results_df = results_df.query("ipm_region in @keep_regions")

        by = ["model_region", "resource"]
        if args.sub_region:
            by.insert(-1, args.sub_region)
        results_df = results_df.groupby(by, as_index=False).sum()
        summary_fn = (
            Path(args.settings_file).parent
            / f"{Path(args.settings_file).stem}_resource_summary_stats.csv"
        )
        print(
            "Summary stats for each model region are being saved to the project folder with your settings file."
        )
    else:
        print(
            "Summary stats for all regions are being saved to the RESOURCE_GROUPS folder listed in your .env file."
        )
        summary_fn = Path(SETTINGS["RESOURCE_GROUPS"]) / "resource_summary_stats.csv"

    results_df.to_csv(summary_fn, index=False, float_format="%.1f")


if __name__ == "__main__":
    main()


# tests/fuel_test.py

"Test functions related to fuel assignment, price, and emissions"

import logging
import os
import numpy as np
import pytest

import powergenome
from powergenome.eia_opendata import fetch_fuel_prices, modify_fuel_prices
from powergenome.fuels import add_user_fuel_prices, add_carbon_tax, fuel_cost_table
from powergenome.generators import GeneratorClusters
from powergenome.params import DATA_PATHS
from powergenome.util import init_pudl_connection

logger = logging.getLogger(powergenome.__name__)
logger.setLevel(logging.INFO)
handler = logging.StreamHandler()
formatter = logging.Formatter(
    # More extensive test-like formatter...
    "%(asctime)s [%(levelname)8s] %(name)s:%(lineno)s %(message)s",
    # This is the datetime format string.
    "%Y-%m-%d %H:%M:%S",
)
handler.setFormatter(formatter)
logger.addHandler(handler)

if os.name == "nt":
    # if user is using a windows system
    sql_prefix = "sqlite:///"
else:
    sql_prefix = "sqlite:////"
pudl_engine, pudl_out, pg_engine = init_pudl_connection(
    start_year=2018,
    end_year=2020,
    pudl_db=sql_prefix + str(DATA_PATHS["test_data"] / "pudl_test_data.db"),
    pg_db=sql_prefix + str(DATA_PATHS["test_data"] / "pg_misc_tables.sqlite3"),
)


@pytest.fixture()
def fuel_settings():
    settings = {
        "modified_atb_new_gen": {
            "ZCF_CombinedCycle1": {
                "new_technology": "ZCF",
                "new_tech_detail": "CCAvgCF",
                "new_cost_case": "Advanced",
                "atb_technology": "NaturalGas",
                "atb_tech_detail": "CCAvgCF",
                "atb_cost_case": "Advanced",
                "size_mw": 500,
            },
            "ZCF_CombinedCycle2": {
                "new_technology": "ZeroCarbon",
                "new_tech_detail": "CCAvgCF",
                "new_cost_case": "Advanced",
                "atb_technology": "NaturalGas",
                "atb_tech_detail": "CCAvgCF",
                "atb_cost_case": "Advanced",
                "size_mw": 500,
            },
            "biopower_ccs": {
                "new_technology": "Biopower",
                "new_tech_detail": "DedicatedCCS",
                "new_cost_case": "Moderate",
                "atb_technology": "Biopower",
                "atb_tech_detail": "Dedicated",
                "atb_cost_case": "Moderate",
                "size_mw": 500,
            },
        },
        "eia_atb_tech_map": {
            "Biomass": ["Biopower_Dedicated"],
            "Zero Carbon": ["ZCF"],
        },
        "atb_new_gen": [
            ["Biopower", "Dedicated", "Moderate", 100],
            ["NaturalGas", "CTAvgCF", "Moderate", 100],
            ["NaturalGas", "CCCCSAvgCF", "Moderate", 100],
        ],
        "atb_data_year": 2020,
        "atb_cap_recovery_years": 20,
        "fuel_eia_aeo_year": 2021,
        "aeo_fuel_usd_year": 2021,
        "eia_series_scenario_names": {"reference": "REF2021"},
        "aeo_fuel_region_map": {"south_atlantic": ["S_VACA", "PJM_Dom"]},
        "eia_series_region_names": {"south_atlantic": "SOATL"},
        "eia_series_fuel_names": {"naturalgas": "NG"},
        "aeo_fuel_scenarios": {"naturalgas": "reference"},
        "user_fuel_price": {
            "zerocarbonfuel1": 14,
            "zerocarbonfuel2": 10,
            "biomass": {"S_VACA": 10, "PJM_Dom": 5},
        },
        "user_fuel_usd_year": {
            "zerocarbonfuel1": 2020,
            "zerocarbonfuel2": 2020,
            "biomass": 2019,
        },
        "tech_fuel_map": {
            "ZeroCarbon_CCAvgCF_Advanced": "zerocarbonfuel1",
            "Zero Carbon": "zerocarbonfuel2",
            "Biomass": "biomass",
            "NaturalGas": "naturalgas",
        },
        "ccs_fuel_map": {
            "biopower_dedicatedccs": "biomass_ccs",
            "naturalgas_ccccs": "naturalgas_ccs90",
        },
        "ccs_capture_rate": {"biomass_ccs": 0.9, "naturalgas_ccs90": 0.9},
        "model_regions": ["S_VACA", "PJM_Dom"],
        "model_year": 2040,
        "model_first_planning_year": 2035,
        "cost_multiplier_region_map": {
            "SRCA": ["S_VACA"],
            "PJMD": ["PJM_Dom"],
        },
        "cost_multiplier_technology_map": {
            "Biomass": ["Biopower_Dedicated"],
            "CC - multi shaft": ["ZeroCarbon_CCAvgCF", "ZCF_CCAvgCF"],
        },
        "fuel_emission_factors": {
            "biomass": 0.1,  # Dummy value for biomass
            "naturalgas": 0.05306,
        },
        "ccs_disposal_cost": 25,
    }
    return settings


def test_fuel_labels_and_prices(fuel_settings):
    df_base = add_user_fuel_prices(fuel_settings)

    for fuel in [
        "S_VACA_biomass",
        "PJM_Dom_biomass",
        "zerocarbonfuel1",
        "zerocarbonfuel2",
    ]:
        assert fuel in df_base["full_fuel_name"].unique()

    fuel_settings["target_usd_year"] = 2020
    df_inflate = add_user_fuel_prices(fuel_settings)
    assert (
        df_base.loc[df_base["fuel"] == "biomass", "price"].mean()
        < df_inflate.loc[df_inflate["fuel"] == "biomass", "price"].mean()
    )
    assert np.allclose(
        df_base.loc[df_base["fuel"] == "zerocarbonfuel1", "price"].mean(),
        df_inflate.loc[df_inflate["fuel"] == "zerocarbonfuel1", "price"].mean(),
    )

    gc = GeneratorClusters(
        pudl_engine, pudl_out, pg_engine, fuel_settings, current_gens=False
    )
    gens = gc.create_new_generators()
    assert gens["Fuel"].isna().any() == False
    assert gens["Fuel"].str.contains("ccs", case=False).any() == True

    fuel_table = fuel_cost_table(gc.fuel_prices, gens, fuel_settings)
    for r in ["PJM_Dom", "S_VACA"]:
        assert (
            fuel_table.loc[0, f"{r}_biomass"]
            == fuel_settings["fuel_emission_factors"]["biomass"]
        )
        assert np.allclose(
            fuel_table.loc[0, f"{r}_biomass_ccs"],
            fuel_settings["fuel_emission_factors"]["biomass"]
            * (1 - fuel_settings["ccs_capture_rate"]["biomass_ccs"]),
        )
    assert (fuel_table.loc[1:, :] == 0).any().any() == False


def test_fetch_fuel_price_errors(fuel_settings, caplog):
    aeo_year = fuel_settings.pop("fuel_eia_aeo_year")
    with pytest.raises(KeyError):
        fetch_fuel_prices(fuel_settings)

    fuel_settings["eia_aeo_year"] = {"bad": "value"}
    with pytest.raises(TypeError):
        fetch_fuel_prices(fuel_settings)

    fuel_settings["eia_aeo_year"] = aeo_year

    eia_series_scenario_names = fuel_settings.pop("eia_series_scenario_names")
    with pytest.raises(KeyError):
        fetch_fuel_prices(fuel_settings)

    fuel_settings["eia_series_scenario_names"] = eia_series_scenario_names
    caplog.set_level(logging.WARNING)
    fetch_fuel_prices(fuel_settings)
    # assert "Unable to inflate fuel prices" in caplog.text

    eia_series_region_names = fuel_settings.pop("eia_series_region_names")
    eia_series_fuel_names = fuel_settings.pop("eia_series_fuel_names")

    fetch_fuel_prices(fuel_settings)
    assert "EIA fuel region names were not found" in caplog.text
    assert "EIA fuel names were not found" in caplog.text


def test_regional_fuel_price_mod(fuel_settings):
    fuel_prices = fetch_fuel_prices(fuel_settings)
    mod_fuel_prices = modify_fuel_prices(
        fuel_prices,
        fuel_settings["aeo_fuel_region_map"],
        fuel_settings.get("regional_fuel_adjustments"),
    )
    assert np.allclose(fuel_prices["price"].values, mod_fuel_prices["price"].values)

    fuel_settings["regional_fuel_adjustments"] = {
        "S_VACA": ["mul", 2],
        "PJM_Dom": {"naturalgas": ["add", 1]},
    }

    mod_fuel_prices = modify_fuel_prices(
        fuel_prices,
        fuel_settings["aeo_fuel_region_map"],
        fuel_settings.get("regional_fuel_adjustments"),
    )
    assert np.isclose(
        fuel_prices.query("region == 'south_atlantic'")["price"].mean(),
        mod_fuel_prices.query("region == 'S_VACA'")["price"].mean() / 2,
    )

    assert np.isclose(
        fuel_prices.query("region == 'south_atlantic' and fuel == 'naturalgas'")[
            "price"
        ].mean(),
        mod_fuel_prices.query("region == 'PJM_Dom' and fuel == 'naturalgas'")[
            "price"
        ].mean()
        - 1,
    )

    fuel_settings["regional_fuel_adjustments"] = {
        "S_VACA": ["mul", 2],
        "PJM_Dom": {"coal": ["add", 1]},
    }
    with pytest.raises(KeyError):
        modify_fuel_prices(
            fuel_prices,
            fuel_settings["aeo_fuel_region_map"],
            fuel_settings.get("regional_fuel_adjustments"),
        )

    fuel_settings["regional_fuel_adjustments"] = {
        "S_VACA": ["mul", 2],
        "PJM_Dom": {"coal": 1},
    }
    with pytest.raises(KeyError):
        modify_fuel_prices(
            fuel_prices,
            fuel_settings["aeo_fuel_region_map"],
            fuel_settings.get("regional_fuel_adjustments"),
        )
    fuel_settings["regional_fuel_adjustments"] = {
        "S_VACA": 1,
        "PJM_Dom": {"coal": 1},
    }
    with pytest.raises(TypeError):
        modify_fuel_prices(
            fuel_prices,
            fuel_settings["aeo_fuel_region_map"],
            fuel_settings.get("regional_fuel_adjustments"),
        )

    fuel_settings["regional_fuel_adjustments"] = {
        "S_VACA": ["mul", 2],
        "PJM_Dom": {"naturalgas": ["add", 1]},
    }
    with pytest.raises(KeyError):
        modify_fuel_prices(
            fuel_prices,
            None,
            fuel_settings.get("regional_fuel_adjustments"),
        )

    fuel_settings["regional_fuel_adjustments"] = {
        "S_VACA": ["mul", 2],
        "PJM_Dom": {"naturalga": ["add", 1]},
    }
    with pytest.raises(KeyError):
        modify_fuel_prices(
            fuel_prices,
            fuel_settings["aeo_fuel_region_map"],
            fuel_settings.get("regional_fuel_adjustments"),
        )

    fuel_settings["regional_fuel_adjustments"] = {
        "S_VACA": ["mul", 2],
        "PJM_Dom": {"naturalgas": ["ad", 1]},
    }
    with pytest.raises(KeyError):
        modify_fuel_prices(
            fuel_prices,
            fuel_settings["aeo_fuel_region_map"],
            fuel_settings.get("regional_fuel_adjustments"),
        )


def test_regional_mod_fuel_labels(fuel_settings):
    fuel_settings["regional_fuel_adjustments"] = {
        "S_VACA": ["mul", 2],
        "PJM_Dom": {"naturalgas": ["add", 1]},
    }
    fuel_settings["target_usd_year"] = 2020
    gc = GeneratorClusters(
        pudl_engine, pudl_out, pg_engine, fuel_settings, current_gens=False
    )
    gens = gc.create_new_generators()
    assert "S_VACA_reference_naturalgas" in gens.Fuel.values
    assert "PJM_Dom_reference_naturalgas" in gens.Fuel.values


# tests/generation_test.py

"""Test functions in generation.py"""
import logging
import sqlite3
import os
from pathlib import Path
from powergenome.GenX import (
    RESOURCE_TAGS,
    add_cap_res_network,
    add_misc_gen_values,
    check_resource_tags,
    create_policy_req,
    create_regional_cap_res,
    hydro_energy_to_power,
    max_cap_req,
    min_cap_req,
    network_line_loss,
    network_max_reinforcement,
    network_reinforcement_cost,
    reduce_time_domain,
    round_col_values,
    set_int_cols,
)
from powergenome.eia_opendata import add_user_fuel_prices
from powergenome.external_data import make_generator_variability

from powergenome.fuels import fuel_cost_table
from powergenome.nrelatb import db_col_values

CWD = Path.cwd()
# os.environ["RESOURCE_GROUPS"] = str(CWD / "data" / "resource_groups_base")
# os.environ["PUDL_DB"] = "sqlite:////" + str(
#     CWD / "tests" / "data" / "pudl_test_data.db"
# )
# os.environ["PG_DB"] = "sqlite:////" + str(
#     CWD / "tests" / "data" / "pg_misc_tables.sqlite3"
# )

import numpy as np
import pandas as pd
import sqlalchemy
import powergenome
import pytest
from powergenome.generators import (
    fill_missing_tech_descriptions,
    gentype_region_capacity_factor,
    group_technologies,
    label_retirement_year,
    label_small_hydro,
    load_demand_response_efs_profile,
    remove_leading_zero,
    unit_generator_heat_rates,
    load_860m,
    GeneratorClusters,
    energy_storage_mwh,
)
from powergenome.load_profiles import (
    add_load_growth,
    make_final_load_curves,
    make_load_curves,
)
from powergenome.params import DATA_PATHS, SETTINGS  # , SETTINGS
from powergenome.transmission import (
    agg_transmission_constraints,
    transmission_line_distance,
)
from powergenome.util import (
    build_scenario_settings,
    find_region_col,
    init_pudl_connection,
    check_settings,
    load_settings,
    map_agg_region_names,
    regions_to_keep,
    remove_fuel_scenario_name,
    reverse_dict_of_lists,
    write_results_file,
)

logger = logging.getLogger(powergenome.__name__)
logger.setLevel(logging.INFO)
handler = logging.StreamHandler()
formatter = logging.Formatter(
    # More extensive test-like formatter...
    "%(asctime)s [%(levelname)8s] %(name)s:%(lineno)s %(message)s",
    # This is the datetime format string.
    "%Y-%m-%d %H:%M:%S",
)
handler.setFormatter(formatter)
logger.addHandler(handler)

# pudl_engine = sqlite3.connect(DATA_PATHS["test_data"] / "pudl_test_data.db")
# pg_engine = sqlalchemy.create_engine(
#     "sqlite:////" + str(DATA_PATHS["test_data"] / "pg_misc_tables.sqlite3")
# )
if os.name == "nt":
    # if user is using a windows system
    sql_prefix = "sqlite:///"
else:
    sql_prefix = "sqlite:////"
pudl_engine, pudl_out, pg_engine = init_pudl_connection(
    start_year=2018,
    end_year=2020,
    pudl_db=sql_prefix + str(DATA_PATHS["test_data"] / "pudl_test_data.db"),
    pg_db=sql_prefix + str(DATA_PATHS["test_data"] / "pg_misc_tables.sqlite3"),
)


@pytest.fixture(scope="module")
def generation_fuel_eia923_data():
    gen_fuel = pd.read_sql_query(
        "SELECT * FROM generation_fuel_eia923",
        pudl_engine,
        parse_dates=["report_date"],
    )
    return gen_fuel


@pytest.fixture(scope="module")
def generators_eia860_data():
    sql = """
        SELECT *
        FROM generators_eia860
        WHERE operational_status_code = 'OP'
    """
    gens_860 = pd.read_sql_query(
        sql, pudl_engine, parse_dates=["report_date", "planned_retirement_date"]
    )
    return gens_860


@pytest.fixture(scope="module")
def generators_entity_eia_data():
    gen_entity = pd.read_sql_query(
        "SELECT * FROM generators_entity_eia",
        pudl_engine,
        parse_dates=["operating_date"],
    )
    return gen_entity


@pytest.fixture(scope="module")
def plant_region_map_ipm_data():
    plant_region_map = pd.read_sql_query(
        "SELECT * FROM plant_region_map_epaipm", pudl_engine
    )
    return plant_region_map


@pytest.fixture(scope="module")
def test_settings():
    settings = load_settings(DATA_PATHS["test_data"] / "test_settings.yml")
    settings["RESOURCE_GROUPS"] = DATA_PATHS["test_data"] / "resource_groups_base"
    return settings


@pytest.fixture(scope="module")
def CA_AZ_settings():
    settings = load_settings(
        DATA_PATHS["powergenome"].parent / "example_systems" / "CA_AZ" / "settings"
    )
    settings["input_folder"] = Path(
        DATA_PATHS["powergenome"].parent
        / "example_systems"
        / "CA_AZ"
        / settings["input_folder"]
    )
    settings["RESOURCE_GROUPS"] = DATA_PATHS["test_data"] / "resource_groups_base"
    scenario_definitions = pd.read_csv(
        settings["input_folder"] / settings["scenario_definitions_fn"]
    )
    scenario_settings = build_scenario_settings(settings, scenario_definitions)

    return scenario_settings[2030]["p1"]


class MockPudlOut:
    """
    The methods in this class read pre-calculated tables from a sqlite db and return
    the expected values from pudl_out methods.
    """

    def hr_by_unit():
        "Heat rate by unit over multiple years"
        hr_by_unit = pd.read_sql_query(
            "SELECT * FROM hr_by_unit", pudl_engine, parse_dates=["report_date"]
        )
        return hr_by_unit

    def bga():
        "Boiler generator associations with unit_id_pudl values"
        bga = pd.read_sql_query(
            "SELECT * FROM boiler_generator_assn_eia860", pudl_engine
        )
        return bga


def test_load_single_settings():
    settings = load_settings(
        DATA_PATHS["powergenome"].parent
        / "example_systems"
        / "CA_AZ"
        / "combined_settings.yml"
    )


def test_load_multiple_settings():
    settings = load_settings(
        DATA_PATHS["powergenome"].parent / "example_systems" / "CA_AZ" / "settings"
    )


def test_group_technologies(generators_eia860_data, test_settings):
    df = generators_eia860_data.loc[
        generators_eia860_data.report_date.dt.year == 2020, :
    ]
    # df = df.query("report_date.dt.year==2017")
    df = df.drop_duplicates(subset=["plant_id_eia", "generator_id"])

    grouped_by_tech = group_technologies(
        df,
        test_settings.get("group_technologies"),
        test_settings.get("tech_groups", {}) or {},
        test_settings.get("regional_no_grouping", {}) or {},
    )
    techs = grouped_by_tech["technology_description"].unique()
    capacities = grouped_by_tech.groupby("technology_description")[
        test_settings["capacity_col"]
    ].sum()
    # expected_hydro_cap = 48.1
    hydro_cap = capacities["Conventional Hydroelectric"]
    expected_peaker_cap = 354.8
    # peaker_cap = capacities["Peaker"]

    assert len(df) == len(grouped_by_tech)
    assert df["capacity_mw"].sum() == grouped_by_tech["capacity_mw"].sum()
    assert "Peaker" in techs
    # assert np.allclose(hydro_cap, expected_hydro_cap)
    # assert np.allclose(peaker_cap, expected_peaker_cap)


def test_fill_missing_tech_descriptions(generators_eia860_data):
    filled = fill_missing_tech_descriptions(generators_eia860_data)

    assert len(generators_eia860_data) == len(
        filled.dropna(subset=["technology_description"])
    )


def test_label_small_hyro(
    generators_eia860_data, test_settings, plant_region_map_ipm_data
):
    region_agg_map = reverse_dict_of_lists(test_settings.get("region_aggregations", {}))
    model_region_map_df = map_agg_region_names(
        df=plant_region_map_ipm_data,
        region_agg_map=region_agg_map,
        original_col_name="region",
        new_col_name="model_region",
    )
    df = pd.merge(
        generators_eia860_data, model_region_map_df, on="plant_id_eia", how="left"
    )
    logger.info(df[["plant_id_eia", "technology_description", "model_region"]].head())

    # df["model_region"] = df["region"].map(reverse_dict_of_lists)

    df = label_small_hydro(df, test_settings, by=["plant_id_eia", "report_date"])
    print(df.query("plant_id_eia==34"))
    logger.info(df["technology_description"].unique())

    assert "Small Hydroelectric" in df["technology_description"].unique()
    assert np.allclose(
        df.loc[df.technology_description == "Small Hydroelectric", "capacity_mw"].sum(),
        140.5,
    )


# def test_label_retirement_year(
#     generators_eia860_data, generators_entity_eia_data, test_settings
# ):
#     gens = pd.merge(
#         generators_eia860_data,
#         generators_entity_eia_data,
#         on=["plant_id_eia", "generator_id"],
#         how="left",
#     )
#     df = label_retirement_year(gens, test_settings)
#     print(df)

#     assert df.loc[df["retirement_year"].isnull(), :].empty is True


def test_unit_generator_heat_rates(data_years=[2020]):
    hr_df = unit_generator_heat_rates(MockPudlOut, data_years)

    assert hr_df.empty is False
    assert "heat_rate_mmbtu_mwh" in hr_df.columns
    assert np.allclose(
        hr_df.query("plant_id_eia==117 & unit_id_pudl == 2")[
            "heat_rate_mmbtu_mwh"
        ].values,
        [7.635626],
    )


def test_load_860m(test_settings):
    eia_860m = load_860m(test_settings)
    test_settings["eia_860m_fn"] = None
    eia_860m = load_860m(test_settings)


def test_agg_transmission_constraints(test_settings):
    agg_transmission_constraints(pg_engine, test_settings)


def test_demand_curve(test_settings):
    make_load_curves(pg_engine, test_settings)


def test_alt_table_load_sources(CA_AZ_settings):
    # Test with a single non-default load table
    CA_AZ_settings["load_source_table_name"] = {
        # "EFS": "load_curves_nrel_efs",
        "FERC": "load_curves_ferc",
    }
    CA_AZ_settings["regional_load_source"] = "FERC"
    CA_AZ_settings["EFS_DATA"] = DATA_PATHS["test_data"] / "efs"
    make_final_load_curves(pg_engine, CA_AZ_settings)


def test_combined_load_sources(CA_AZ_settings):
    # Test with a combination of user and database load sources
    CA_AZ_settings["regional_load_fn"] = "test_regional_load_profiles.csv"
    CA_AZ_settings["load_source_table_name"] = {"EFS": "load_curves_nrel_efs"}
    CA_AZ_settings["regional_load_source"] = {
        "USER": ["CA_N", "CA_S"],
        "EFS": ["WECC_AZ"],
    }
    CA_AZ_settings["load_source_table_name"] = {
        "EFS": "load_curves_nrel_efs",
        "FERC": "load_curves_ferc",
    }
    CA_AZ_settings["EFS_DATA"] = DATA_PATHS["test_data"] / "efs"
    make_final_load_curves(pg_engine, CA_AZ_settings)


def test_check_settings(test_settings):
    check_settings(test_settings, pg_engine)


# def test_gentype_region_capacity_factor(plant_region_map_ipm_data, test_settings):
#     cf_techs = test_settings["capacity_factor_techs"]

#     plant_region_map_ipm_data = plant_region_map_ipm_data.rename(
#         columns={"region": "model_region"}
#     )
#     df = gentype_region_capacity_factor(
#         pudl_engine, plant_region_map_ipm_data, test_settings
#     )
#     print(df.technology.unique())
#     assert "Biomass" in df.technology.unique()
#     # CF can sometime be greater than 1, but shouldn't be significantly higher.
#     assert df.loc[df["technology"].isin(cf_techs), "capacity_factor"].max() < 2


def test_gen_integration(CA_AZ_settings, tmp_path):
    CA_AZ_settings["atb_modifiers"] = {
        "ngccccs": {
            "technology": "NaturalGas",
            "tech_detail": "CCCCSAvgCF",
            "Heat_Rate_MMBTU_per_MWh": 7.159,
        }
    }
    CA_AZ_settings["modified_atb_new_gen"]["NGCCS100"]["heat_rate"] = 7.5
    CA_AZ_settings["EFS_DATA"] = DATA_PATHS["test_data"] / "efs"
    gc = GeneratorClusters(
        pudl_engine, pudl_out, pg_engine, CA_AZ_settings, supplement_with_860m=True
    )
    all_gens = gc.create_all_generators()
    all_gens = add_misc_gen_values(all_gens, CA_AZ_settings)

    assert (
        all_gens.loc[all_gens["Resource"] == "CA_N_biomass_1", "Eff_Up"].values[0] == 1
    )
    assert np.allclose(
        all_gens.query("technology.str.contains('NaturalGas_CCCCS', case=False)")[
            "Heat_Rate_MMBTU_per_MWh"
        ].mean(),
        7.159,
    )
    assert np.allclose(
        all_gens.query("technology.str.contains('CCS100', case=False)")[
            "Heat_Rate_MMBTU_per_MWh"
        ].mean(),
        7.5,
    )
    # Capacity in existing clusters should be > 0
    assert all(gc.results["Existing_Cap_MW"] > 0)
    batteries = gc.results.query("technology == 'Batteries'")

    # Battery energy capacity should be larger than battery power capacity
    assert all(batteries["Existing_Cap_MWh"] > batteries["Existing_Cap_MW"])

    gen_variability = make_generator_variability(all_gens)
    assert (gen_variability >= 0).all().all()

    fuels = fuel_cost_table(
        fuel_costs=gc.fuel_prices,
        generators=gc.all_resources,
        settings=gc.settings,
    )
    fuels.index.name = "Time_Index"
    write_results_file(
        df=remove_fuel_scenario_name(fuels, gc.settings)
        .pipe(set_int_cols)
        .pipe(round_col_values),
        folder=tmp_path,
        file_name="Fuels_data.csv",
        include_index=True,
    )
    load = make_final_load_curves(pg_engine=pg_engine, settings=gc.settings)
    (
        reduced_resource_profile,
        reduced_load_profile,
        time_series_mapping,
        representative_point,
    ) = reduce_time_domain(gen_variability, load, gc.settings)
    if gc.settings["reduce_time_domain"]:
        assert len(representative_point) == gc.settings["time_domain_periods"]
        assert (
            time_series_mapping["Rep_Period"].nunique()
            == gc.settings["time_domain_periods"]
        )
        assert representative_point.isna().any().all() == False
        assert time_series_mapping.isna().any().all() == False
    assert len(reduced_load_profile) == len(reduced_resource_profile)

    gc.settings["distributed_gen_method"]["CA_N"] = "fraction_load"
    gc.settings["distributed_gen_values"][2030]["CA_N"] = 0.1
    gc.settings["regional_load_fn"] = "test_regional_load_profiles.csv"
    gc.settings["regional_load_includes_demand_response"] = False
    make_final_load_curves(pg_engine=pg_engine, settings=gc.settings)

    model_regions_gdf = gc.model_regions_gdf
    transmission = (
        agg_transmission_constraints(pg_engine=pg_engine, settings=gc.settings)
        .pipe(
            transmission_line_distance,
            ipm_shapefile=model_regions_gdf,
            settings=gc.settings,
            units="mile",
        )
        .pipe(network_line_loss, settings=gc.settings)
        .pipe(network_max_reinforcement, settings=gc.settings)
        .pipe(network_reinforcement_cost, settings=gc.settings)
        .pipe(set_int_cols)
        .pipe(round_col_values)
        .pipe(add_cap_res_network, settings=gc.settings)
    )

    if gc.settings.get("emission_policies_fn"):
        energy_share_req = create_policy_req(gc.settings, col_str_match="ESR")
        co2_cap = create_policy_req(gc.settings, col_str_match="CO_2")
    min_cap = min_cap_req(gc.settings)

    cap_res = create_regional_cap_res(gc.settings)


def test_existing_gen_profiles():
    ipm_regions = pd.read_sql_table("regions_entity_epaipm", pg_engine)
    regions = [r for r in ipm_regions.region_id_epaipm.to_list() if "CN_" not in r]

    s = """
    SELECT DISTINCT technology_description
    FROM generators_eia860
    """
    technologies = (
        pd.read_sql_query(s, pudl_engine).dropna()["technology_description"].to_list()
    )
    technologies.remove("Natural Gas with Compressed Air Storage")
    settings = dict(
        RESOURCE_GROUPS=DATA_PATHS["test_data"] / "resource_groups_base",
        target_usd_year=2019,
        model_year=2030,
        model_first_planning_year=2022,
        model_regions=regions,
        data_years=[2020],
        capacity_col="capacity_mw",
        num_clusters={tech: 2 for tech in technologies},
        retirement_ages={tech: 200 for tech in technologies},
        atb_data_year=2022,
        atb_existing_year=2020,
        fuel_eia_aeo_year=2022,
        load_eia_aeo_year=2020,
        aeo_fuel_usd_year=2021,
        eia_series_region_names={
            "mountain": "MTN",
        },
        eia_series_fuel_names={
            "naturalgas": "NG",
        },
        eia_series_scenario_names={
            "reference": "REF2022",
        },
        aeo_fuel_scenarios={
            "coal": "reference",
            "naturalgas": "reference",
            "distillate": "reference",
            "uranium": "reference",
        },
        eia_atb_tech_map={
            "Battery": "Battery_*",
            "Batteries": "Battery_*",
            "Biomass": "Biopower_Dedicated",
            "Solar Thermal without Energy Storage": "CSP_Class1",
            "Conventional Steam Coal": "Coal_newAvgCF",
            "Coal Integrated Gasification Combined Cycle": "NaturalGas_CCAvgCF",
            "Natural Gas Fired Combined Cycle": "NaturalGas_CCAvgCF",  # [NaturalGas_CCAvgCF, NETL_NGCC]
            "Natural Gas Fired Combustion Turbine": "NaturalGas_CTAvgCF",
            "Peaker": "NaturalGas_CTAvgCF",
            "Natural Gas Internal Combustion Engine": "NaturalGas_CTAvgCF",
            "Landfill Gas": "NaturalGas_CTAvgCF",
            "Petroleum Liquids": "NaturalGas_CTAvgCF",
            "Municipal Solid Waste": "Biopower_Dedicated",
            "Other Waste Biomass": "Biopower_Dedicated",
            "Wood/Wood Waste Biomass": "Biopower_Dedicated",
            "Solar Photovoltaic": "UtilityPV_Class1",
            "Geothermal": "Geothermal_HydroFlash",  # assume installed capacity is dominated by flash
            "Conventional Hydroelectric": "Hydropower_NSD4",  # Large variability based on choice
            "Hydroelectric Pumped Storage": "Hydropower_NSD4",  # Large variability based on choice
            "Small Hydroelectric": "Hydropower_NSD3",  # Large variability based on choice
            "Onshore Wind Turbine": "LandbasedWind_Class4",  # All onshore wind is the same
            "Offshore Wind Turbine": "OffShoreWind_Class10",  # Mid-range of floating offshore wind
            "Nuclear": "Nuclear_Nuclear",
            "Natural Gas Steam Turbine": "Coal_newAvgCF",  # No gas steam turbines in ATB, using coal instead
            "Solar Thermal with Energy Storage": "CSP_Class1",
            "Solar Thermal without Energy Storage": "CSP_Class1",
            "Other Gases": "NaturalGas_CTAvgCF",
            "Other Natural Gas": "NaturalGas_CTAvgCF",
            "Petroleum Coke": "Coal_newAvgCF",
            "All Other": "NaturalGas_CTAvgCF",
            "Flywheels": "Battery_*",
            "Natural Gas with Compressed Air Storage": "NaturalGas_CTAvgCF",
        },
        startup_vom_costs_mw={
            "coal_small_sub": 2.81,
            "coal_large_sub": 2.69,
            "coal_supercritical": 2.98,
            "gas_cc": 1.03,
            "gas_large_ct": 0.77,
            "gas_aero_ct": 0.70,
            "gas_steam": 1.03,
            "nuclear": 5.4,
        },
        startup_vom_costs_usd_year=2011,
        startup_costs_type="startup_costs_per_cold_start_mw",
        startup_costs_per_cold_start_usd_year=2011,
        startup_costs_per_cold_start_mw={
            "coal_small_sub": 147,
            "coal_large_sub": 105,
            "coal_supercritical": 104,
            "gas_cc": 79,
            "gas_large_ct": 103,
            "gas_aero_ct": 32,
            "gas_steam": 75,
            "nuclear": 210,
        },
        existing_startup_costs_tech_map={
            "Conventional Steam Coal": "coal_large_sub",
            "Natural Gas Fired Combined Cycle": "gas_cc",
            "Natural Gas Fired Combustion Turbine": "gas_large_ct",
            "Natural Gas Steam Turbine": "gas_steam",
            "Nuclear": "nuclear",
        },
    )
    gc = GeneratorClusters(
        pudl_engine, pudl_out, pg_engine, settings, supplement_with_860m=False
    )
    existing_gen = gc.create_region_technology_clusters()
    gen_variability = make_generator_variability(existing_gen)
    assert (gen_variability >= 0).all().all()


def test_cap_req():
    settings = {
        "model_tag_names": ["MinCapTag_1", "MinCapTag_2", "MaxCapTag_1", "MaxCapTag_2"],
        "MinCapReq": {
            "MinCapTag_1": {"description": "Landbasedwind", "min_mw": 8000},
            "MinCapTag_2": {"description": "CA_S_solar", "min_mw": 10000},
        },
        "MaxCapReq": {
            "MaxCapTag_1": {"description": "Landbasedwind", "max_mw": 8000},
            "MaxCapTag_2": {"description": "CA_S_solar", "max_mw": 10000},
        },
        "generator_columns": [],
    }

    max_cap = max_cap_req(settings)
    min_cap = min_cap_req(settings)

    assert set(settings["generator_columns"]) == set(settings["model_tag_names"])
    assert min_cap.isna().any().all() == False
    assert max_cap.isna().any().all() == False


def test_check_resource_tags():
    # Check something that should fail
    cols = ["region", "technology"] + RESOURCE_TAGS
    data = [pd.Series(["a", "b"] + [1] * len(RESOURCE_TAGS), index=cols)]
    df = pd.DataFrame(data)

    with pytest.raises(Exception):
        check_resource_tags(df)

    # Check something that should pass
    cols = ["region", "technology"] + RESOURCE_TAGS
    data = [pd.Series(["a", "b", 1] + [0] * (len(RESOURCE_TAGS) - 1), index=cols)]
    df = pd.DataFrame(data)

    check_resource_tags(df)

    # Check something that should pass
    cols = ["region", "technology"] + RESOURCE_TAGS
    data = [pd.Series(["a", "b", 2] + [0] * (len(RESOURCE_TAGS) - 1), index=cols)]
    df = pd.DataFrame(data)

    check_resource_tags(df)


def test_storage_duration(caplog):
    settings = {
        "energy_storage_duration": {
            "hydroelectric pumped": 15.5,
            "batteries": {"A": 2, "B": 1, "D": 1},
        }
    }

    data = {
        "region": ["A", "A", "A", "B", "B", "B", "C"],
        "technology": ["Hydroelectric Pumped Storage", "Batteries", "Nuclear"] * 2
        + ["Batteries"],
        "Existing_Cap_MW": [1] * 7,
    }
    df = pd.DataFrame(data)

    caplog.set_level(logging.WARNING)
    df_mwh = energy_storage_mwh(
        df,
        settings["energy_storage_duration"],
        "technology",
        "Existing_Cap_MW",
        "Existing_Cap_MWh",
    )
    assert "The regions ['C'] are missing from technology batteries" in caplog.text
    assert "technology 'batteries' has the region 'D'" in caplog.text

    assert df.equals(df_mwh[df.columns])
    mwh = pd.Series([15.5, 2, 0, 15.5, 1, 0, 0])
    assert mwh.equals(df_mwh["Existing_Cap_MWh"])


def test_hydro_energy_to_power():
    settings = {
        "hydro_factor": 2,
        "regional_hydro_factor": {
            "A": 4,
            "B": 1,
        },
    }

    data = {
        "region": ["A", "A", "A", "B", "B", "B", "C"],
        "technology": ["Hydro", "NG", "Nuclear"] * 2 + ["Hydro"],
        "HYDRO": [1, 0, 0] * 2 + [1],
        "profile": [[0.5], np.nan, np.nan] * 2 + [[0.6]],
    }
    df = pd.DataFrame(data)

    df_hydro_ratio = hydro_energy_to_power(
        df, settings["hydro_factor"], settings["regional_hydro_factor"]
    )
    assert df.equals(df_hydro_ratio[df.columns])
    hydro_ratio = pd.Series([2, 0, 0, 1, 0, 0, 1.2])
    assert hydro_ratio.equals(df_hydro_ratio["Hydro_Energy_to_Power_Ratio"])


def test_efs_flex_demand():
    flex_demand = load_demand_response_efs_profile(
        resource="trans_light_duty",
        electrification_stock_fn="EFS_STOCK_AGG.parquet",
        model_year=2034,
        electrification_scenario="REFERENCE ELECTRIFICATION - MODERATE TECHNOLOGY ADVANCEMENT",
        model_regions=["WECC_AZ", "WEC_BANC", "WEC_CALN"],
        region_aggregations={"CA_N": ["WEC_BANC", "WEC_CALN"]},
        path_in=DATA_PATHS["test_data"] / "efs",
    )

    assert set(flex_demand.columns) == set(["WECC_AZ", "CA_N"])
    assert flex_demand.isnull().any().any() == False
    assert flex_demand.min().min() >= 0

    with pytest.raises(ValueError):
        load_demand_response_efs_profile(
            resource="trans_light_duty",
            electrification_stock_fn="EFS_STOCK_AGG.parquet",
            model_year=2034,
            electrification_scenario="REFERENCE ELECTRIFICATION - MODERATE TECHNOLOGY",
            model_regions=["WECC_AZ", "WEC_BANC", "WEC_CALN"],
            region_aggregations={"CA_N": ["WEC_BANC", "WEC_CALN"]},
            path_in=DATA_PATHS["test_data"] / "efs",
        )


def test_flex_resources(CA_AZ_settings):
    CA_AZ_settings["model_year"] = 2035
    CA_AZ_settings["model_first_planning_year"] = 2030
    CA_AZ_settings["electrification_stock_fn"] = "EFS_STOCK_AGG.parquet"
    CA_AZ_settings[
        "electrification_scenario"
    ] = "REFERENCE ELECTRIFICATION - MODERATE TECHNOLOGY ADVANCEMENT"
    CA_AZ_settings["flexible_demand_resources"] = {
        2035: {
            "trans_light_duty": {
                "fraction_shiftable": 0.8,
                "parameter_values": {"Max_DSM_delay": 5, "DR": 2},
            },
            "res_water_heat": {
                "fraction_shiftable": 0.1,
                "parameter_values": {"Max_DSM_delay": 2, "DR": 2},
            },
        }
    }
    CA_AZ_settings["atb_new_gen"] = []
    CA_AZ_settings["renewables_clusters"] = None
    CA_AZ_settings["additional_technologies_fn"] = None
    CA_AZ_settings["modified_atb_new_gen"] = None
    CA_AZ_settings["atb_modifiers"] = None
    CA_AZ_settings["efs_path"] = DATA_PATHS["test_data"] / "efs"

    gc = GeneratorClusters(
        pudl_engine, pudl_out, pg_engine, CA_AZ_settings, current_gens=False
    )
    flex_resources = gc.create_new_generators()

    assert len(flex_resources) == 6
    non_null_cols = ["profile", "Max_DSM_delay", "DR"]
    assert flex_resources[non_null_cols].isnull().any().any() == False


def test_usr_tx(tmp_path):
    settings = {
        "input_folder": tmp_path,
        "user_transmission_constraints_fn": "usr_tx.csv",
        "model_regions": ["A", "B", "C"],
        "tx_value_col": "nonfirm_ttc_mw",
    }

    usr_tx = pd.DataFrame(
        data={
            "region_from": ["A", "C", "C"],
            "region_to": ["B", "A", "B"],
            "nonfirm_ttc_mw": [100, 200, 300],
        }
    )
    usr_tx.to_csv(tmp_path / "usr_tx.csv", index=False)

    tx_constraints = agg_transmission_constraints(pg_engine, settings=settings)

    assert tx_constraints["Line_Max_Flow_MW"].to_list() == [100, 200, 300]

    usr_tx = pd.DataFrame(
        data={
            "region_from": ["A", "C", "C"],
            "region_to": ["B", "A", "B"],
            "firm_ttc_mw": [100, 200, 300],
        }
    )
    usr_tx.to_csv(tmp_path / "usr_tx.csv", index=False)

    settings["tx_value_col"] = "firm_ttc_mw"
    tx_constraints = agg_transmission_constraints(pg_engine, settings=settings)

    assert tx_constraints["Line_Max_Flow_MW"].to_list() == [100, 200, 300]

    usr_tx = pd.DataFrame(
        data={
            "region_from": ["A", "C", "C"],
            "region_to": ["B", "A", "B"],
            "nonfirm_ttc_mw": [100, 200, 300],
        }
    )
    usr_tx = pd.concat(
        [
            usr_tx,
            usr_tx.rename(
                columns={"region_from": "region_to", "region_to": "region_from"}
            ),
        ]
    )
    usr_tx.to_csv(tmp_path / "usr_tx.csv", index=False)

    settings["tx_value_col"] = "nonfirm_ttc_mw"
    tx_constraints = agg_transmission_constraints(pg_engine, settings=settings)

    assert tx_constraints["Line_Max_Flow_MW"].to_list() == [100, 200, 300]

    usr_tx = pd.DataFrame(
        data={
            "region_from": ["A", "C", "C", "A"],
            "region_to": ["B", "A", "B", "B"],
            "nonfirm_ttc_mw": [100, 200, 300, 50],
        }
    )
    usr_tx.to_csv(tmp_path / "usr_tx.csv", index=False)

    with pytest.raises(KeyError):
        agg_transmission_constraints(pg_engine, settings=settings)


def test_load_growth(CA_AZ_settings):
    data = {
        "year": [2019] * 3,
        "region": ["WECC_AZ"] * 3,
        "time_index": [1] * 3,
        "sector": ["commercial", "residential", "industrial"],
        "load_mw": [1, 1, 1],
    }

    load_curves = pd.DataFrame(data)

    load_curves = add_load_growth(load_curves, CA_AZ_settings)

    assert all(load_curves["load_mw"] > 1)

    data = {
        "year": [2017] * 3,
        "region": ["WECC_AZ"] * 3,
        "time_index": [1] * 3,
        "sector": ["commercial", "residential", "industrial"],
        "load_mw": [1, 1, 1],
    }

    load_curves_2 = pd.DataFrame(data)

    load_curves_2 = add_load_growth(load_curves_2, CA_AZ_settings)

    assert all(load_curves["load_mw"] < load_curves_2["load_mw"])

    data = {
        "year": [2019] * 3,
        "region": ["WECC_AZ"] * 3,
        "time_index": [1] * 3,
        "load_mw": [1, 1, 1],
    }

    load_curves_3 = pd.DataFrame(data)

    load_curves_3 = add_load_growth(load_curves_3, CA_AZ_settings)

    assert all(load_curves_3["load_mw"] > 1)
    assert load_curves_3.iloc[0, -1] == load_curves_3.iloc[1, -1]

    data = {
        "year": [2017] * 3,
        "region": ["WECC_AZ"] * 3,
        "time_index": [1] * 3,
        "load_mw": [1, 1, 1],
    }

    load_curves_4 = pd.DataFrame(data)

    load_curves_4 = add_load_growth(load_curves_4, CA_AZ_settings)

    assert load_curves_4.iloc[0, -1] == load_curves_4.iloc[1, -1]
    assert all(load_curves_3["load_mw"] < load_curves_4["load_mw"])


def test_db_col_values():
    values = db_col_values(pg_engine, "technology_costs_nrelatb", ["technology"])
    assert "NaturalGas" in values


def test_find_region_col():
    df = pd.DataFrame(columns=["A", "Region", "C"])
    region_col = find_region_col(df.columns)
    assert region_col == "Region"

    df = pd.DataFrame(columns=["A", "model_region", "C"])
    region_col = find_region_col(df.columns)
    assert region_col == "model_region"

    with pytest.raises(ValueError):
        df = pd.DataFrame(columns=["A", "model_region", "region"])
        region_col = find_region_col(df.columns)

    with pytest.raises(ValueError):
        df = pd.DataFrame(columns=["A", "B", "C"])
        region_col = find_region_col(df.columns)


def test_remove_leading_zeros():
    s = "01"
    assert remove_leading_zero(s) == "1"

    i = 1
    assert remove_leading_zero(i) == 1

    s = "GEN1"
    assert remove_leading_zero(s) == "GEN1"


# tests/load_test.py

"""
Test functions related to load profiles
"""

import pandas as pd
from powergenome.load_profiles import grow_historical_load


def test_grow_historical_load():
    base_load = {"region": ["A", "A", "B", "B"], "load_mw": [1, 1, 1, 1]}
    base_load_df = pd.DataFrame(base_load)
    hist_region_map = {
        "A": "MWRCE",
        "B": "TRE",
    }
    future_region_map = {
        "A": "MCW",
        "B": "TRE",
    }
    keep_regions = ["A", "B"]
    start_year = 2012
    load_aeo_year = 2022
    end_year, load_2019 = grow_historical_load(
        base_load_df.copy(),
        start_year,
        load_aeo_year,
        keep_regions,
        hist_region_map,
        future_region_map,
    )

    assert end_year > start_year
    assert not base_load_df.equals(load_2019)

    end_year, load_2021 = grow_historical_load(
        load_2019.copy(),
        end_year,
        load_aeo_year,
        keep_regions,
        hist_region_map,
        future_region_map,
    )

    assert end_year == 2021
    assert not load_2021.equals(load_2019)

    aeo_sector_map = {
        "commercial": "COMM",
        "industrial": "IDAL",
        "residential": "RESD",
        "transportation": "TRN",
    }
    base_load = {
        "region": ["A", "A", "B", "B"],
        "load_mw": [1, 1, 1, 1],
        "sector": ["commercial"] * 4,
    }
    base_load_df = pd.DataFrame(base_load)
    hist_region_map = {
        "A": "MWRCE",
        "B": "TRE",
    }
    future_region_map = {
        "A": "MCW",
        "B": "TRE",
    }
    keep_regions = ["A", "B"]
    start_year = 2013
    end_year, load_2019_sector = grow_historical_load(
        base_load_df.copy(),
        start_year,
        load_aeo_year,
        keep_regions,
        hist_region_map,
        future_region_map,
        aeo_sector_map,
    )

    assert not load_2019_sector.equals(load_2019)

    assert end_year > start_year
    assert not base_load_df.equals(load_2019_sector)

    end_year, load_2021_sector = grow_historical_load(
        load_2019_sector.copy(),
        end_year,
        load_aeo_year,
        keep_regions,
        hist_region_map,
        future_region_map,
        aeo_sector_map,
    )

    assert end_year == 2021
    assert not load_2021_sector.equals(load_2019_sector)


# tests/price_test.py

"Test functions from price_adjustment.py"

import pandas as pd
from powergenome.eia_opendata import read_eia_api
from powergenome.params import SETTINGS
from powergenome.price_adjustment import inflation_price_adjustment, load_cpi_data


def test_load_cpi(tmp_path):
    # Get data the first time
    cpi = load_cpi_data(data_path=tmp_path / "cpi_data.csv")

    # Get data saved to file
    cpi = load_cpi_data(data_path=tmp_path / "cpi_data.csv")

    assert all(cpi["value"].notna())
    assert all(cpi["period"] == 12)


def test_inflate_price(tmp_path):
    p = 10
    p2 = inflation_price_adjustment(p, 2000, 2010, data_path=tmp_path / "cpi_data.csv")
    assert p2 > p

    # remove later years from saved CPI file
    cpi = pd.read_csv(tmp_path / "cpi_data.csv")
    first_year = cpi["year"].min()
    last_year = cpi["year"].max()
    cpi.loc[5:10, :].to_csv(tmp_path / "cpi_data.csv", float_format="%g")

    p = pd.Series([1, 10])
    p2 = inflation_price_adjustment(p, first_year, last_year)
    assert all(p2 > p)


def test_eia_api():
    series_id = "AEO.2022.REF2022.PRCE_REAL_ELEP_NA_NG_NA_NEENGL_Y13DLRPMMBTU.A"
    df = read_eia_api(series_id, SETTINGS["EIA_API_KEY"], columns=["year", "price"])
    assert not df.empty


# tests/transmission_test.py

"Test functions for interregional transmission lines"

from collections import namedtuple
import os
from pathlib import Path
import numpy as np
import pandas as pd
import pytest

from powergenome.util import (
    build_scenario_settings,
    load_settings,
    init_pudl_connection,
)
from powergenome.params import DATA_PATHS
from powergenome.external_data import load_user_tx_costs, insert_user_tx_costs
from powergenome.transmission import agg_transmission_constraints

if os.name == "nt":
    # if user is using a windows system
    sql_prefix = "sqlite:///"
else:
    sql_prefix = "sqlite:////"
pudl_engine, pudl_out, pg_engine = init_pudl_connection(
    pudl_db=sql_prefix + str(DATA_PATHS["test_data"] / "pudl_test_data.db"),
    pg_db=sql_prefix + str(DATA_PATHS["test_data"] / "pg_misc_tables.sqlite3"),
)


@pytest.fixture(scope="module")
def CA_AZ_settings():
    settings = load_settings(
        DATA_PATHS["powergenome"].parent / "example_systems" / "CA_AZ" / "settings"
    )
    settings["input_folder"] = Path(
        DATA_PATHS["powergenome"].parent
        / "example_systems"
        / "CA_AZ"
        / settings["input_folder"]
    )
    settings["RESOURCE_GROUPS"] = DATA_PATHS["test_data"] / "resource_groups_base"
    scenario_definitions = pd.read_csv(
        settings["input_folder"] / settings["scenario_definitions_fn"]
    )
    scenario_settings = build_scenario_settings(settings, scenario_definitions)

    return scenario_settings[2030]["p1"]


def test_load_user_tx(tmp_path):
    cols = [
        "start_region",
        "dest_region",
        "total_interconnect_annuity_mw",
        "total_interconnect_cost_mw",
        "total_line_loss_frac",
        "dollar_year",
    ]
    tx_line = namedtuple("tx_line", cols)
    lines = [
        tx_line("CA_S", "WECC_AZ", 1000, 100000, 0.07, 2018),
        tx_line("CA_N", "CA_N", 2000, 200000, 0.06, 2018),
    ]
    user_tx = pd.DataFrame(lines)
    user_tx.to_csv(tmp_path / "tx_lines.csv", index=False)

    model_regions = ["CA_N", "CA_S", "WECC_AZ"]
    target_usd_year = 2020

    user_tx_costs = load_user_tx_costs(
        tmp_path / "tx_lines.csv", model_regions, target_usd_year
    )

    assert all(
        user_tx_costs["total_interconnect_annuity_mw"]
        > user_tx["total_interconnect_annuity_mw"]
    )

    user_tx_costs = load_user_tx_costs(
        tmp_path / "tx_lines.csv", model_regions, target_usd_year=None
    )

    assert np.allclose(
        user_tx_costs["total_interconnect_annuity_mw"],
        user_tx["total_interconnect_annuity_mw"],
    )


def test_insert_user_tx_costs(tmp_path, CA_AZ_settings):
    cols = [
        "start_region",
        "dest_region",
        "total_interconnect_annuity_mw",
        "total_interconnect_cost_mw",
        "total_line_loss_frac",
        "dollar_year",
    ]
    tx_line = namedtuple("tx_line", cols)
    lines = [
        tx_line("CA_S", "WECC_AZ", 1000, 100000, 0.07, 2018),
        tx_line("CA_N", "CA_N", 2000, 200000, 0.06, 2018),
    ]
    user_tx = pd.DataFrame(lines)
    user_tx.to_csv(tmp_path / "tx_lines.csv", index=False)
    model_regions = ["CA_N", "CA_S", "WECC_AZ"]
    target_usd_year = 2020

    user_tx = load_user_tx_costs(
        tmp_path / "tx_lines.csv", model_regions, target_usd_year
    )

    tx_constraints = agg_transmission_constraints(
        pg_engine,
        CA_AZ_settings,
    )

    combined_tx = insert_user_tx_costs(tx_constraints, user_tx)

    assert len(combined_tx) == 2

    req_cols = [
        "Network_Lines",
        "z1",
        "z2",
        "z3",
        "Line_Reinforcement_Cost_per_MWyr",
        "Line_Reinforcement_Cost_per_MW",
        "Line_Loss_Percentage",
    ]
    for col in req_cols:
        assert col in combined_tx.columns
        assert combined_tx[col].notnull().all()


# setup.py

from setuptools import find_packages, setup

setup(
    name="powergenome",
    packages=find_packages(),
    version="0.6.0",
    description="Extract PUDL data for use in power system models",
    author="Greg Schivley",
    entry_points={
        "console_scripts": [
            "run_powergenome_multiple = powergenome.run_powergenome_multiple_outputs_cli:main",
        ]
    },
)


# powergenome/GenX.py

"Functions specific to GenX outputs"

from itertools import product
import logging
from pathlib import Path
from typing import Dict, List
import pandas as pd

from powergenome.external_data import (
    load_policy_scenarios,
    load_demand_segments,
    load_user_genx_settings,
    make_generator_variability,
)
from powergenome.load_profiles import make_distributed_gen_profiles
from powergenome.time_reduction import kmeans_time_clustering
from powergenome.util import find_region_col, load_settings, snake_case_col
from powergenome.nrelatb import investment_cost_calculator

logger = logging.getLogger(__name__)

INT_COLS = [
    "Inv_Cost_per_MWyr",
    "Fixed_OM_Cost_per_MWyr",
    "Inv_Cost_per_MWhyr",
    "Fixed_OM_Cost_per_MWhyr",
    "Line_Reinforcement_Cost_per_MW_yr",
    "Up_Time",
    "Down_Time",
]

COL_ROUND_VALUES = {
    "Var_OM_Cost_per_MWh": 2,
    "Var_OM_Cost_per_MWh_in": 2,
    "Start_Cost_per_MW": 0,
    "Cost_per_MMBtu": 2,
    "CO2_content_tons_per_MMBtu": 5,
    "Cap_Size": 2,
    "Heat_Rate_MMBTU_per_MWh": 2,
    "distance_mile": 4,
    "Line_Max_Reinforcement_MW": 0,
    "distance_miles": 1,
    "distance_km": 1,
    "Existing_Cap_MW": 1,
    "Existing_Cap_MWh": 1,
    "Max_Cap_MW": 1,
    "Max_Cap_MWh": 1,
    "Min_Cap_MW": 1,
    "Min_Cap_MWh": 1,
    "Line_Loss_Percentage": 4,
}

# RESOURCE_TAGS = ["THERM", "VRE", "MUST_RUN", "STOR", "FLEX", "HYDRO", "LDS"]
RESOURCE_TAGS = ["THERM", "VRE", "MUST_RUN", "STOR", "FLEX", "HYDRO"]


def create_policy_req(settings: dict, col_str_match: str) -> pd.DataFrame:
    model_year = settings["model_year"]
    case_id = settings["case_id"]

    policies = load_policy_scenarios(settings)
    policy_cols = [c for c in policies.columns if col_str_match in c]
    if len(policy_cols) == 0:
        return None

    year_case_policy = policies.loc[(case_id, model_year), ["region"] + policy_cols]
    if isinstance(year_case_policy, pd.DataFrame):
        year_case_policy = year_case_policy.dropna(subset=policy_cols)
    elif isinstance(year_case_policy, pd.Series):
        year_case_policy = year_case_policy.dropna()
    else:
        raise TypeError(
            "Somehow the object 'year_case_policy' is not a dataframe of a series."
            "Please submit this as an issue in the PowerGenome repository."
            f"{year_case_policy}"
        )
    # Bug where multiple regions for a case will return this as a df, even if the policy
    # for this case applies to all regions (code below expects a Series)
    ycp_shape = year_case_policy.shape
    if ycp_shape[0] == 1 and len(ycp_shape) > 1:
        year_case_policy = year_case_policy.squeeze()  # convert to series

    zones = settings["model_regions"]
    zone_num_map = {
        zone: f"z{number + 1}" for zone, number in zip(zones, range(len(zones)))
    }

    zone_cols = ["Region_description", "Network_zones"] + policy_cols
    zone_df = pd.DataFrame(columns=zone_cols, dtype=float)
    zone_df["Region_description"] = zones
    zone_df["Network_zones"] = zone_df["Region_description"].map(zone_num_map)
    # If there is only one region, assume that the policy is applied across all regions.
    if isinstance(year_case_policy, pd.Series):
        logger.info(
            "Only one zone was found in the emissions policy file."
            " The same emission policies are being applied to all zones."
        )
        for col, value in year_case_policy[policy_cols].iteritems():
            if "CO_2_Max_Mtons" in col:
                zone_df.loc[:, col] = 0
                if value > 0:
                    zone_df.loc[0, col] = value
            else:
                zone_df.loc[:, col] = value
    else:
        for region, col in product(
            year_case_policy["region"].unique(), year_case_policy[policy_cols].columns
        ):
            zone_df.loc[
                zone_df["Region_description"] == region, col
            ] = year_case_policy.loc[year_case_policy.region == region, col].values[0]

    # zone_df = zone_df.drop(columns="region")

    return zone_df


def create_regional_cap_res(settings: dict) -> pd.DataFrame:
    """Create a dataframe of regional capacity reserve constraints from settings params

    Parameters
    ----------
    settings : dict
        PowerGenome settings dictionary with parameters 'model_regions' and
        'regional_capacity_reserves'

    Returns
    -------
    pd.DataFrame
        A dataframe with a 'Network_zones' column and one 'CapRes_*' for each capacity
        reserve constraint

    Raises
    ------
    KeyError
        A region listed in a capacity reserve constraint is not a valid model region
    """

    if not settings.get("regional_capacity_reserves"):
        return None
    else:
        zones = settings["model_regions"]
        zone_num_map = {
            zone: f"z{number + 1}" for zone, number in zip(zones, range(len(zones)))
        }
        cap_res_cols = list(settings["regional_capacity_reserves"])
        cap_res_df = pd.DataFrame(index=zones, columns=["Network_zones"] + cap_res_cols)
        cap_res_df["Network_zones"] = cap_res_df.index.map(zone_num_map)

        for col, cap_res_dict in settings["regional_capacity_reserves"].items():
            for region, val in cap_res_dict.items():
                if region not in zones:
                    raise KeyError(
                        f"The region {region} in 'regional_capacity_reserves', {col} "
                        "is not a valid model region."
                    )
                cap_res_df.loc[region, col] = val

        cap_res_df = cap_res_df.fillna(0)

        return cap_res_df


def label_cap_res_lines(path_names: List[str], dest_regions: List[str]) -> List[int]:
    """Label if each transmission line is part of a capacity reserve constraint and
    if line flow is into or out of the constraint region.

    Parameters
    ----------
    path_names : List[str]
        String names of transmission lines, with format <region>_to_<region>
    dest_regions : List[str]
        Names of model regions, corresponding to region names used in 'path_names'

    Returns
    -------
    List[int]
        Same length as 'path_names'. Values of 1 mean the line connects a region within
        a capacity reserve constraint to a region outside the constraint. Values of -1
        mean it connects a region outside a capacity reserve constraint to a region
        within the constraint. Values of 0 mean it connects two regions that are both
        within or outside the constraint.
    """
    cap_res_list = []
    for name in path_names:
        s_r = name.split("_to_")[0]
        e_r = name.split("_to_")[-1]
        if (s_r in dest_regions) and not (e_r in dest_regions):
            cap_res_list.append(1)
        elif (e_r in dest_regions) and not (s_r in dest_regions):
            cap_res_list.append(-1)
        else:
            cap_res_list.append(0)

    assert len(cap_res_list) == len(path_names)

    return cap_res_list


def add_cap_res_network(tx_df: pd.DataFrame, settings: dict) -> pd.DataFrame:
    """Add capacity reserve colums to the transmission dataframe (Network.csv)

    Parameters
    ----------
    tx_df : pd.DataFrame
        Transmission dataframe with rows for each transmission line, a column
        'transmission_path_name', and columns 'z1', 'z2', etc for each model region. The
        'z*' columns have values of -1 (line is outbound from region), 0 (not connected),
        or 1 (inbound to region).
    settings : dict
        PowerGenome settings, with parameters 'model_regions', 'regional_capacity_reserves',
        and 'cap_res_network_derate_default'.


    Returns
    -------
    pd.DataFrame
        A copy of the input dataframe with additional columns 'CapRes_*', 'DerateCapRes_*',
        and 'CapRes_Excl_*' for each capacity reserve constraint
    """
    if (
        "Transmission Path Name" in tx_df.columns
        and "transmission_path_name" not in tx_df.columns
    ):
        tx_df["transmission_path_name"] = tx_df["Transmission Path Name"]
    original_cols = tx_df.columns.to_list()

    zones = settings["model_regions"]
    zone_num_map = {
        zone: f"z{number + 1}" for zone, number in zip(zones, range(len(zones)))
    }
    path_names = tx_df["transmission_path_name"].to_list()
    policy_nums = []

    # Loop through capacity reserve constraints (CapRes_*) and determine network
    # parameters for each
    for cap_res in settings.get("regional_capacity_reserves", {}) or {}:
        cap_res_num = int(cap_res.split("_")[-1])  # the number of the capres constraint
        policy_nums.append(cap_res_num)
        # TODO #179 fix reference to regional_capacity_reserves key of settings dict
        dest_regions = list(
            settings["regional_capacity_reserves"][cap_res].keys()
        )  # list of regions in the CapRes
        # May add ability to have different values by CapRes and line in the future
        tx_df[f"DerateCapRes_{cap_res_num}"] = settings.get(
            "cap_res_network_derate_default", 0.95
        )
        tx_df[f"CapRes_Excl_{cap_res_num}"] = label_cap_res_lines(
            path_names, dest_regions
        )

    policy_nums.sort()
    derate_cols = [f"DerateCapRes_{n}" for n in policy_nums]
    excl_cols = [f"CapRes_Excl_{n}" for n in policy_nums]

    return tx_df[original_cols + derate_cols + excl_cols].fillna(0)


def add_emission_policies(transmission_df, settings):
    """Add emission policies to the transmission dataframe

    Parameters
    ----------
    transmission_df : DataFrame
        Zone to zone transmission constraints
    settings : dict
        User-defined parameters from a settings file. Should have keys of `input_folder`
        (a Path object of where to find user-supplied data) and
        `emission_policies_fn` (the file to load).
    DistrZones : [type], optional
        Placeholder setting, by default None

    Returns
    -------
    DataFrame
        The emission policies provided by user next to the transmission constraints.
    """

    model_year = settings["model_year"]
    case_id = settings["case_id"]

    policies = load_policy_scenarios(settings)
    year_case_policy = policies.loc[(case_id, model_year), :]

    # Bug where multiple regions for a case will return this as a df, even if the policy
    # for this case applies to all regions (code below expects a Series)
    ycp_shape = year_case_policy.shape
    if ycp_shape[0] == 1 and len(ycp_shape) > 1:
        year_case_policy = year_case_policy.squeeze()  # convert to series

    zones = settings["model_regions"]
    zone_num_map = {
        zone: f"z{number + 1}" for zone, number in zip(zones, range(len(zones)))
    }

    zone_cols = ["Region description", "Network_zones"] + list(policies.columns)
    zone_df = pd.DataFrame(columns=zone_cols)
    zone_df["Region description"] = zones
    zone_df["Network_zones"] = zone_df["Region description"].map(zone_num_map)

    # Add code here to make DistrZones something else!
    # If there is only one region, assume that the policy is applied across all regions.
    if isinstance(year_case_policy, pd.Series):
        logger.info(
            "Only one zone was found in the emissions policy file."
            " The same emission policies are being applied to all zones."
        )
        for col, value in year_case_policy.iteritems():
            if col == "CO_2_Max_Mtons":
                zone_df.loc[:, col] = 0
                if value > 0:
                    zone_df.loc[0, col] = value
            else:
                zone_df.loc[:, col] = value
    else:
        for region, col in product(
            year_case_policy["region"].unique(), year_case_policy.columns
        ):
            zone_df.loc[
                zone_df["Region description"] == region, col
            ] = year_case_policy.loc[year_case_policy.region == region, col].values[0]

    zone_df = zone_df.drop(columns="region")

    network_df = pd.concat([zone_df, transmission_df], axis=1)

    return network_df


def add_misc_gen_values(
    gen_clusters: pd.DataFrame, settings: dict, resource_col: str = "Resource"
) -> pd.DataFrame:
    """Add parameter values from a CSV file to resources.

    Parameters
    ----------
    gen_clusters : pd.DataFrame
        Resource dataframe with columns "region" and `resource_col`.
    settings : dict
        Model settings, with parameters "input_folder" and "misc_gen_inputs_fn". The
        misc gen CSV file should have the column `resource_col`. If it has the column
        "region" then regional values will be applied, otherwise values for each resource
        will be applied across all regions.
    resource_col : str, optional
        Name of the column with resource name in both gen_clusters and the CSV file, by
        default "Resource".

    Returns
    -------
    pd.DataFrame
        A modified version of gen_clusters with new parameter values for resources.
    """
    path = Path(settings["input_folder"]) / settings["misc_gen_inputs_fn"]
    misc_values = pd.read_csv(path)
    misc_values[resource_col] = snake_case_col(misc_values[resource_col])

    context = f"Assigning misc generator values from the user-supplied file {path}."
    try:
        region_col = find_region_col(misc_values.columns, context)
    except ValueError:
        region_col = "region"
        misc_values["region"] = "all"
    regions = [r for r in misc_values[region_col].unique() if r.lower() != "all"]
    wrong_regions = [r for r in regions if r not in settings["model_regions"]]
    if wrong_regions:
        raise ValueError(
            f"The `misc_gen_inputs_fn` CSV has regions {wrong_regions}, which are not "
            f"valid model regions. Valid model regions are {settings['model_regions']}."
        )

    for region in settings["model_regions"]:
        _df = misc_values.loc[misc_values[region_col].str.lower() == "all", :]
        _df.loc[:, "region"] = region
        misc_values = misc_values.append(_df)

    misc_values = misc_values.loc[misc_values[region_col].str.lower() != "all", :]

    resource_len = 0
    for tech, _df in misc_values.groupby(resource_col):
        num_tech_regions = len(
            gen_clusters.loc[
                gen_clusters[resource_col].str.contains(tech, case=False)
            ].drop_duplicates(subset=["region"])
        )
        num_values = len(_df)
        if num_values < num_tech_regions:
            logger.warning(
                f"The `misc_gen_inputs_fn` CSV has {num_values} region(s) for the resource "
                f"'{tech}', but the resource is in {num_tech_regions} regions. Check "
                "your input file to ensure values are provided for all appropriate regions."
            )
    generic_resources = []
    for gen_resource in gen_clusters[resource_col].unique():
        for r in settings["model_regions"]:
            if r in gen_resource:
                gen_resource = gen_resource.replace(r + "_", "")
                generic_resources.append(gen_resource)
                continue
    generic_resources = set(generic_resources)
    missing_resources = []
    for resource in generic_resources:
        match = False
        for misc_resource in misc_values[resource_col].unique():
            if misc_resource in resource.lower():
                match = True
                continue
        if not match:
            missing_resources.append(resource)

    if missing_resources:
        logger.warning(
            f"The resources {missing_resources} are not included in your `misc_gen_inputs_fn` "
            "CSV file. This is a warning in case they should have parameters in that file."
        )

    misc_values = misc_values.reset_index(drop=True)
    value_cols = [
        col for col in misc_values.columns if col not in [region_col, resource_col]
    ]

    for idx, row in misc_values.iterrows():
        row_cols = row[value_cols].dropna().index
        gen_clusters.loc[
            (gen_clusters["region"] == row[region_col])
            & (gen_clusters[resource_col].str.contains(row[resource_col], case=False)),
            row_cols,
        ] = row[row_cols].values
    return gen_clusters


def reduce_time_domain(
    resource_profiles, load_profiles, settings, variable_resources_only=True
):
    demand_segments = load_demand_segments(settings)

    if settings.get("reduce_time_domain"):
        days = settings["time_domain_days_per_period"]
        time_periods = settings["time_domain_periods"]
        include_peak_day = settings["include_peak_day"]
        load_weight = settings["demand_weight_factor"]

        results, representative_point, _ = kmeans_time_clustering(
            resource_profiles=resource_profiles,
            load_profiles=load_profiles,
            days_in_group=days,
            num_clusters=time_periods,
            include_peak_day=include_peak_day,
            load_weight=load_weight,
            variable_resources_only=variable_resources_only,
        )

        reduced_resource_profile = results["resource_profiles"]
        reduced_resource_profile.index.name = "Resource"
        reduced_resource_profile.index = range(1, len(reduced_resource_profile) + 1)
        reduced_load_profile = results["load_profiles"]
        time_series_mapping = results["time_series_mapping"]

        time_index = pd.Series(data=reduced_load_profile.index + 1, name="Time_Index")
        sub_weights = pd.Series(
            data=[x * (days * 24) for x in results["ClusterWeights"]],
            name="Sub_Weights",
        )
        hours_per_period = pd.Series(data=[days * 24], name="Timesteps_per_Rep_Period")
        subperiods = pd.Series(data=[time_periods], name="Rep_Periods")
        reduced_load_output = pd.concat(
            [
                demand_segments,
                subperiods,
                hours_per_period,
                sub_weights,
                time_index,
                reduced_load_profile.round(0),
            ],
            axis=1,
        )

        return (
            reduced_resource_profile,
            reduced_load_output,
            time_series_mapping,
            representative_point,
        )

    else:
        time_index = pd.Series(data=range(1, 8761), name="Time_Index")
        sub_weights = pd.Series(data=[1], name="Sub_Weights")
        hours_per_period = pd.Series(data=[168], name="Timesteps_per_Rep_Period")
        subperiods = pd.Series(data=[1], name="Rep_Periods")

        # Not actually reduced
        load_output = pd.concat(
            [
                demand_segments,
                subperiods,
                hours_per_period,
                sub_weights,
                time_index,
                load_profiles.reset_index(drop=True),
            ],
            axis=1,
        )
        resource_profiles.index = time_index

        return resource_profiles, load_output, None, None


def network_line_loss(transmission: pd.DataFrame, settings: dict) -> pd.DataFrame:
    """Add line loss percentage for each network line between regions.

    Parameters
    ----------
    transmission : pd.DataFrame
        One network line per row with a column "distance_mile"
    settings : dict
        User-defined settings with a parameter "tx_line_loss_100_miles"

    Returns
    -------
    pd.DataFrame
        Same as input but with the new column 'Line_Loss_Percentage'
    """
    if "tx_line_loss_100_miles" not in settings:
        raise KeyError(
            "The parameter 'tx_line_loss_100_miles' is required in your settings file."
        )
    loss_per_100_miles = settings["tx_line_loss_100_miles"]
    if "distance_mile" in transmission.columns:
        distance_col = "distance_mile"
    elif "distance_km" in transmission.columns:
        distance_col = "distance_km"
        loss_per_100_miles *= 0.62137
        logger.info("Line loss per 100 miles was converted to km.")
    else:
        raise KeyError("No distance column is available in the transmission dataframe")

    transmission["Line_Loss_Percentage"] = (
        transmission[distance_col] / 100 * loss_per_100_miles
    )

    return transmission


def network_reinforcement_cost(
    transmission: pd.DataFrame, settings: dict
) -> pd.DataFrame:
    """Add transmission line reinforcement investment costs (per MW-mile-year)

    Parameters
    ----------
    transmission : pd.DataFrame
        One network line per row with columns "transmission_path_name" and
        "distance_mile".
    settings : dict
        User-defined settings with dictionary "transmission_investment_cost.tx" with
        keys "capex_mw_mile", "wacc", and "investment_years".

    Returns
    -------
    pd.DataFrame
        Same as input but with the new column 'Line_Reinforcement_Cost_per_MW_yr'
    """

    cost_dict = (
        settings.get("transmission_investment_cost", {})
        .get("tx", {})
        .get("capex_mw_mile")
    )
    if not cost_dict:
        raise KeyError(
            "No value for transmission reinforcement costs is included in the settings."
            " These costs are included under transmission_investment_costs.tx."
            "capex_mw_mile.<model_region>. See the `test_settings.yml` file for an "
            "example."
        )
    origin_region_cost = (
        transmission["transmission_path_name"].str.split("_to_").str[0].map(cost_dict)
    )
    dest_region_cost = (
        transmission["transmission_path_name"].str.split("_to_").str[-1].map(cost_dict)
    )

    # Average the costs per mile between origin and destination regions
    line_capex = (origin_region_cost + dest_region_cost) / 2
    line_wacc = (
        settings.get("transmission_investment_cost", {}).get("tx", {}).get("wacc")
    )
    if not line_wacc:
        raise KeyError(
            "No value for the transmission weighted average cost of capital (wacc) is "
            "included in the settings."
            "This numeric value is included under transmission_investment_costs.tx."
            "wacc. See the `test_settings.yml` file for an "
            "example."
        )
    line_inv_period = (
        settings.get("transmission_investment_cost", {})
        .get("tx", {})
        .get("investment_years")
    )
    if not line_inv_period:
        raise KeyError(
            "No value for the transmission investment period is included in the settings."
            "This numeric value is included under transmission_investment_costs.tx."
            "investment_years. See the `test_settings.yml` file for an example."
        )
    line_inv_cost = (
        investment_cost_calculator(line_capex, line_wacc, line_inv_period)
        * transmission["distance_mile"]
    )

    transmission["Line_Reinforcement_Cost_per_MWyr"] = line_inv_cost.round(0)

    return transmission


def network_max_reinforcement(
    transmission: pd.DataFrame, settings: dict
) -> pd.DataFrame:
    """Add the maximum amount that transmission lines between regions can be reinforced
    in a planning period.

    Parameters
    ----------
    transmission : pd.DataFrame
        One network line per row with the column "Line_Max_Flow_MW"
    settings : dict
        User-defined settings with the parameter "Line_Max_Reinforcement_MW"

    Returns
    -------
    pd.DataFrame
        A copy of the input transmission constraint dataframe with a new column
        `Line_Max_Reinforcement_MW`.
    """

    max_expansion = settings.get("tx_expansion_per_period")

    if not max_expansion and max_expansion != 0:
        raise KeyError(
            "No value for the transmission expansion allowed in this model period is "
            "included in the settings."
            "This numeric value is included under tx_expansion_per_period. See the "
            "`test_settings.yml` file for an example."
        )
    # if isinstance(max_expansion, dict):
    #     expansion_method = settings.get("tx_expansion_method")
    #     if not expansion_method:
    #         raise KeyError(
    #         "The transmission expansion parameter 'tx_expansion_per_period' is a "
    #         "dictionary. There should also be a settings parameter 'tx_expansion_method' "
    #         "with a dictionary of model region: <type> (either 'capacity' or 'fraction' "
    #         "but it isn't in your settings file."
    #     )
    #     for region, value in max_expansion.items():
    #         existing_tx = transmission.query("Region description == @region")["Line_Max_Flow_MW"]
    #         if expansion_method[region].lower() == "capacity":
    #             transmission.loc[transmission["Region description"] == region,
    #             "Line_Max_Reinforcement_MW"] = value
    #         elif expansion_method[region].lower() == "fraction":
    #             transmission.loc[transmission["Region description"] == region,
    #             "Line_Max_Reinforcement_MW"] = value * existing_tx
    #         else:
    #             raise KeyError(
    #             "The transmission expansion method parameter (tx_expansion_method) "
    #             "should have values of 'capacity' or 'fraction' for each model region. "
    #             f"The value provided was '{expansion_method[region]}'."
    #         )

    # else:
    transmission.loc[:, "Line_Max_Reinforcement_MW"] = (
        transmission.loc[:, "Line_Max_Flow_MW"] * max_expansion
    )
    transmission["Line_Max_Reinforcement_MW"] = transmission[
        "Line_Max_Reinforcement_MW"
    ].round(0)

    return transmission


def set_int_cols(df: pd.DataFrame, cols: list = None) -> pd.DataFrame:
    """Set values of some dataframe columns to integers.

    Parameters
    ----------
    df : pd.DataFrame
        Input dataframe.
    cols : list, optional
        Columns to set as integer, by default None. If none, will use
        `powergenome.GenX.INT_COLS`.

    Returns
    -------
    pd.DataFrame
        Input dataframe with some columns set as integer.
    """
    if not cols:
        cols = INT_COLS

    cols = [c for c in cols if c in df.columns]

    for col in cols:
        df[col] = df[col].fillna(0).astype(int)
    return df


def round_col_values(
    df: pd.DataFrame, col_round_val: Dict[str, int] = None
) -> pd.DataFrame:
    """Round values in columns to a specific sigfig.

    Parameters
    ----------
    df : pd.DataFrame
        Input dataframe.
    col_round_val : Dict, optional
        Dictionary with key values of column labels and integer values of the number of
        sigfigs, by default None.

    Returns
    -------
    pd.DataFrame
        Same dataframe as input but with rounded values in specified columns.
    """
    if not col_round_val:
        col_round_val = COL_ROUND_VALUES

    col_round_val = {k: v for k, v in col_round_val.items() if k in df.columns}

    for col, value in col_round_val.items():
        df[col] = df[col].fillna(0).round(value)
    return df


def calculate_partial_CES_values(gen_clusters, fuels, settings):
    gens = gen_clusters.copy()
    if settings.get("partial_ces"):
        assert set(fuels["Fuel"]) == set(gens["Fuel"])
        fuel_emission_map = fuels.copy()
        fuel_emission_map = fuel_emission_map.set_index("Fuel")

        gens["co2_emission_rate"] = gens["Heat_Rate_MMBTU_per_MWh"] * gens["Fuel"].map(
            fuel_emission_map["CO2_content_tons_per_MMBtu"]
        )

        # Make the partial CES credit equal to 1 ton minus the emissions rate, but
        # don't include coal plants

        partial_ces = 1 - gens["co2_emission_rate"]

        gens.loc[
            ~(gens["Resource"].str.contains("coal"))
            & (gens["STOR"] == 0)
            & (gens["FLEX"] == 0),
            # & ~(gens["Resource"].str.contains("battery"))
            # & ~(gens["Resource"].str.contains("load_shifting")),
            "CES",
        ] = partial_ces.round(3)
    # else:
    #     gen_clusters = add_genx_model_tags(gen_clusters, settings)

    gens["Zone"] = gens["Zone"]
    gens["Cap_Size"] = gens["Cap_size"]
    gens["Fixed_OM_Cost_per_MWyr"] = gens["Fixed_OM_Cost_per_MWyr"]
    gens["Fixed_OM_Cost_per_MWhyr"] = gens["Fixed_OM_Cost_per_MWhyr"]
    gens["Inv_Cost_per_MWyr"] = gens["Inv_Cost_per_MWyr"]
    gens["Inv_Cost_per_MWhyr"] = gens["Inv_Cost_per_MWhyr"]
    gens["Var_OM_Cost_per_MWh"] = gens["Var_OM_Cost_per_MWh"]
    # gens["Var_OM_Cost_per_MWh_In"] = gens["Var_OM_Cost_per_MWh_in"]
    gens["Start_Cost_per_MW"] = gens["Start_Cost_per_MW"]
    gens["Start_Fuel_MMBTU_per_MW"] = gens["Start_fuel_MMBTU_per_MW"]
    gens["Heat_Rate_MMBTU_per_MWh"] = gens["Heat_Rate_MMBTU_per_MWh"]
    gens["Min_Power"] = gens["Min_Power"]
    # gens["Self_Disch"] = gens["Self_disch"]
    # gens["Eff_Up"] = gens["Eff_up"]
    # gens["Eff_Down"] = gens["Eff_down"]
    # gens["Ramp_Up_Percentage"] = gens["Ramp_Up_percentage"]
    # gens["Ramp_Dn_Percentage"] = gens["Ramp_Dn_percentage"]
    # gens["Up_Time"] = gens["Up_time"]
    # gens["Down_Time"] = gens["Down_time"]
    # gens["Max_Flexible_Demand_Delay"] = gens["Max_DSM_delay"]
    gens["technology"] = gens["Resource"]
    gens["Resource"] = (
        gens["region"] + "_" + gens["technology"] + "_" + gens["cluster"].astype(str)
    )

    return gens


def check_min_power_against_variability(gen_clusters, resource_profile):
    min_gen_levels = resource_profile.min()

    assert len(min_gen_levels) == len(
        gen_clusters
    ), "The number of hourly resource profiles does not match the number of resources"
    # assert (
    #     gen_clusters["Min_Power"].isna().any() is False
    # ), (
    #     "At least one Min_Power value in 'gen_clusters' is null before checking against"
    #     " resource variability"
    # )

    min_gen_levels.index = gen_clusters.index

    gen_clusters["Min_Power"] = gen_clusters["Min_Power"].combine(min_gen_levels, min)

    # assert (
    #     gen_clusters["Min_Power"].isna().any() is False
    # ), (
    #     "At least one Min_Power value in 'gen_clusters' is null. Values were fine "
    #     "before checking against resoruce variability"
    # )

    return gen_clusters


def calc_emissions_ces_level(network_df, load_df, settings):
    # load_cols = [col for col in load_df.columns if "Load" in col]
    total_load = load_df.sum().sum()

    emissions_limit = settings["emissions_ces_limit"]

    if emissions_limit is not None:
        try:
            ces_value = 1 - (emissions_limit / total_load)
        except TypeError:
            print(emissions_limit, total_load)
            raise TypeError
        network_df["CES"] = ces_value
        network_df["CES"] = network_df["CES"].round(3)

        return network_df
    else:
        return network_df


def fix_min_power_values(
    resource_df: pd.DataFrame,
    gen_profile_df: pd.DataFrame,
    min_power_col: str = "Min_Power",
) -> pd.DataFrame:
    """Fix potentially erroneous min power values for resources with variable generation
    profiles. Any min power values that are higher than the lowest hourly generation
    will be adjusted down to match the lowest hourly generation.


    Parameters
    ----------
    resource_df : pd.DataFrame
        Records of generators/resources. Row order should match column order of
        `gen_profile_df`.
    gen_profile_df : pd.DataFrame
        Hourly generation values for all generators/resources. Column order should match
        row order in `resource_df`.
    min_power_col : str
        Column in `resource_df` that stores the minimum generation power of each
        resource. Default value is "Min_Power".

    Returns
    -------
    pd.DataFrame
        A modified version of `resource_df`. Any rows with minimum power larger than
        hourly generation are adjusted down to match the smallest hourly generation.
    """
    if min_power_col not in resource_df.columns:
        raise ValueError(
            f"When variable generation values against resource min power, the column "
            f"{min_power_col} was not found in the resource dataframe."
        )

    if resource_df.shape[0] != gen_profile_df.shape[1]:
        raise ValueError(
            "When trying to fix min power values, the number of resource dataframe rows"
            f" ({resource_df.shape[0]} rows) does not match the number of variable "
            f"profiles columns ({gen_profile_df.shape[1]} columns)."
        )

    resource_df = resource_df.reset_index(drop=True)
    resource_df.loc[:, "unadjusted_min_power"] = resource_df.loc[:, min_power_col]
    _gen_profile = gen_profile_df.copy(deep=True)
    _gen_profile
    gen_profile_min = _gen_profile.min().reset_index(drop=True)
    mask = (resource_df[min_power_col].fillna(0) > gen_profile_min).values

    logger.info(
        f"{sum(mask)} resources have {min_power_col} larger than hourly generation."
    )

    resource_df.loc[mask, min_power_col] = gen_profile_min[mask].round(3)

    return resource_df


def min_cap_req(settings: dict) -> pd.DataFrame:
    """Create a dataframe of minimum capacity requirements for GenX

    Parameters
    ----------
    settings : dict
        Dictionary with user settings. Should include the key `MinCapReg` with nested
        keys of `MinCapTag_*`, then further nested keys `description` and `min_mw`. The
        `MinCapTag_*` should also be listed as values under `model_tag_names`. Any
        technologies eligible for each of the `MinCapTag_*` should have `model_tag_values`
        of 1.

    Returns
    -------
    pd.DataFrame
        A dataframe with minimum capacity constraints formatted for GenX. If `MinCapReq`
        is not included in the settings dictionary it will return None.

    Raises
    ------
    KeyError
        If a `MinCapTag_*` is included under `MinCapReq` but not included in `model_tag_names`
        the function will raise an error.
    """

    c_num = []
    description = []
    min_mw = []

    # if settings.get("MinCapReq"):
    for cap_tag, values in (settings.get("MinCapReq", {}) or {}).items():
        if cap_tag not in settings.get("model_tag_names", []):
            raise KeyError(
                f"The minimum capacity tag {cap_tag} is listed in the settings "
                "'MinCapReq' but not under 'model_tag_names'. You must add it to "
                "'model_tag_names' for the column to appear in Generators_data.csv."
            )

        # It's easy to forget to add all the necessary column names to the
        # generators_columns list in settings.
        if cap_tag not in settings.get("generator_columns", []) and isinstance(
            settings.get("generator_columns"), list
        ):
            settings["generator_columns"].append(cap_tag)

        c_num.append(cap_tag.split("_")[1])
        description.append(values.get("description"))
        min_mw.append(values.get("min_mw"))

    min_cap_df = pd.DataFrame()
    min_cap_df["MinCapReqConstraint"] = c_num
    min_cap_df["Constraint_Description"] = description
    min_cap_df["Min_MW"] = min_mw

    if not min_cap_df.empty:
        return min_cap_df
    else:
        return None


def max_cap_req(settings: dict) -> pd.DataFrame:
    """Create a dataframe of maximum capacity requirements for GenX

    Parameters
    ----------
    settings : dict
        Dictionary with user settings. Should include the key `MinCapReg` with nested
        keys of `MaxCapTag_*`, then further nested keys `description` and `min_mw`. The
        `MaxCapTag_*` should also be listed as values under `model_tag_names`. Any
        technologies eligible for each of the `MaxCapTag_*` should have `model_tag_values`
        of 1.

    Returns
    -------
    pd.DataFrame
        A dataframe with maximum capacity constraints formatted for GenX. If `MaxCapReq`
        is not included in the settings dictionary it will return None.

    Raises
    ------
    KeyError
        If a `MaxCapTag_*` is included under `MaxCapReq` but not included in `model_tag_names`
        the function will raise an error.
    """

    c_num = []
    description = []
    max_mw = []

    # if settings.get("MaxCapReq"):
    for cap_tag, values in (settings.get("MaxCapReq", {}) or {}).items():
        if cap_tag not in settings.get("model_tag_names", []):
            raise KeyError(
                f"The maximum capacity tag {cap_tag} is listed in the settings "
                "'MaxCapReq' but not under 'model_tag_names'. You must add it to "
                "'model_tag_names' for the column to appear in Generators_data.csv."
            )

        # It's easy to forget to add all the necessary column names to the
        # generators_columns list in settings.
        if cap_tag not in settings.get("generator_columns", []) and isinstance(
            settings.get("generator_columns"), list
        ):
            settings["generator_columns"].append(cap_tag)

        c_num.append(cap_tag.split("_")[1])
        description.append(values.get("description"))
        max_mw.append(values.get("max_mw"))

    max_cap_df = pd.DataFrame()
    max_cap_df["MaxCapReqConstraint"] = c_num
    max_cap_df["Constraint_Description"] = description
    max_cap_df["Max_MW"] = max_mw

    if not max_cap_df.empty:
        return max_cap_df
    else:
        return None


def check_resource_tags(df: pd.DataFrame) -> pd.DataFrame:
    """Check complete generators dataframe to make sure each resource is assigned one,
    and only one, resource tag.

    Parameters
    ----------
    df : pd.DataFrame
        Resource clusters. Should have columns "technology" and "region" in addition
        to the resource tag columns expected by GenX.

    Returns
    -------
    pd.DataFrame
        An unaltered version of the input dataframe.
    """
    tags = [t for t in RESOURCE_TAGS if t in df.columns]
    df_copy = df.loc[:, ["technology", "region"] + tags].copy()
    df_copy[tags] = df_copy[tags].where(df_copy[tags] == 0, 1)
    if not (df_copy[tags].sum(axis=1) == 1).all():
        for idx, row in df_copy.iterrows():
            num_tags = row[tags].sum()
            if num_tags == 0:
                logger.warning(
                    "\n*************************\n"
                    f"The resource {row['technology']} in region {row['region']} does "
                    "not have any assigned resource tags. Check the 'model_tag_values' and "
                    "'regional_tag_values' parameters in your settings file to make sure"
                    "it is assigned one resource tag type from this list:\n\n"
                    f"{RESOURCE_TAGS}\n"
                )
            if num_tags > 1:
                s = row[tags]
                _tags = list(s[s == 1].index)
                logger.warning(
                    "\n*************************\n"
                    f"The resource {row['technology']} in region {row['region']} is "
                    f"assigned {num_tags} resource tags ({_tags}). Check the 'model_tag_values'"
                    " and 'regional_tag_values' parameters in your settings file to make"
                    " sure it is assigned only one resource tag type from this list:\n\n"
                    f"{RESOURCE_TAGS}\n"
                )

        raise ValueError(
            "Use the warnings above to fix the resource tags in your settings file."
        )
    return df


def hydro_energy_to_power(
    df: pd.DataFrame,
    default_factor: float = None,
    regional_factors: Dict[str, float] = {},
) -> pd.DataFrame:
    """Calculate the hydro energy to power ratio. Uses average hydro inflow rate and
    multiplied by a factor to calculate the rated number of hours of reservoir hydro
    storage at peak discharge power output. Value minimum is 1.

    Parameters
    ----------
    df : pd.DataFrame
        DataFrame of resources. Hydro resources should be identified with a "HYDRO" tag
        value of 1
    default_factor : float, optional
        Hydro factor used to scale average inflow rate, by default None
    regional_factors : Dict[str, float], optional
        Specific regional hydro factors used to scale average inflow rate, by default {}

    Returns
    -------
    pd.DataFrame
        Copy of the input dataframe with the new column "Hydro_Energy_to_Power_Ratio"
    """
    if "HYDRO" not in df.columns:
        logger.warning(
            "Generators do not have a column 'HYDRO', so no hydro energy to power ratio is calculated."
        )
        return df
    if not default_factor and not regional_factors:
        logger.warning(
            "No hydro factors have been included in the settings, so no hydro energy to power ratio is calculated."
        )
        return df
    hydro_mask = df["HYDRO"] == 1
    avg_inflow = (
        make_generator_variability(df).mean().reset_index(drop=True) * default_factor
    ).loc[hydro_mask]
    df.loc[hydro_mask, "Hydro_Energy_to_Power_Ratio"] = avg_inflow.where(
        avg_inflow > 1, 1
    )

    for region, factor in (regional_factors or {}).items():
        region_mask = df["region"] == region
        if region_mask.any():
            avg_inflow = (
                make_generator_variability(df).mean().reset_index(drop=True) * factor
            ).loc[hydro_mask & region_mask]
            df.loc[
                (df["HYDRO"] == 1) & region_mask, "Hydro_Energy_to_Power_Ratio"
            ] = avg_inflow.where(avg_inflow > 1, 1)
    df["Hydro_Energy_to_Power_Ratio"] = df["Hydro_Energy_to_Power_Ratio"].fillna(0)
    return df


def rename_gen_cols(
    df: pd.DataFrame, rename_cols: Dict[str, str] = None
) -> pd.DataFrame:
    """Rename columns in the generators data file.

    By default the only rename so far is "existing_mwh" to "Existing_Cap_MWh".

    Parameters
    ----------
    df : pd.DataFrame
        Final dataframe of generating resources.
    rename_cols : Dict[str, str], optional
        Additional rename pairs, by default None

    Returns
    -------
    pd.DataFrame
        Identical to input dataframe except for renamed columns.
    """

    rename = {
        "capacity_mwh": "Existing_Cap_MWh",
    }
    if rename_cols:
        rename.update(rename_cols)

    df = df.rename(columns=rename, errors="ignore")

    return df


# powergenome/cluster_method.py

"Different ways to cluster plants"

import numpy as np
import pandas as pd
from sklearn import cluster, preprocessing


def build_cluster_method_dict(settings):
    cluster_method_dict = {}

    for region in settings["model_regions"]:
        cluster_method_dict[region] = {}
        if region not in settings["cluster_by_owner_regions"]:
            for tech in settings["num_clusters"]:
                cluster_method_dict[region][tech] = cluster_kmeans
        else:
            if tech in settings["cluster_by_owner_regions"][region]:
                cluster_method_dict[region][tech] = cluster_by_owner
            else:
                cluster_method_dict[region][tech] = cluster_kmeans

    return cluster_method_dict


def cluster_kmeans(
    grouped: pd.DataFrame, region: str, tech: str, settings: dict
) -> pd.DataFrame:
    if region in settings.get("alt_num_clusters", {}):
        # if tech in settings["alt_num_clusters"][region]:
        n_clusters = settings["alt_num_clusters"][region].get(tech)
    else:
        n_clusters = settings["num_clusters"].get(tech)

    if n_clusters:
        clusters = cluster.KMeans(n_clusters=n_clusters, random_state=6).fit(
            preprocessing.StandardScaler().fit_transform(grouped)
        )
        grouped["cluster"] = clusters.labels_ + 1  # Change to 1-index for julia

    return grouped


def cluster_by_owner(grouped_units, weighted_ownership, plants, region, tech, settings):
    # I'm doing this by generator right now, but don't have generator info in the inputs.
    # Can calculate weighted percent ownership by pudl unit and use that in here.

    capacity_col = settings["capacity_col"]
    start_capacity = grouped_units[capacity_col].sum()

    owner_cols = [
        # "utility_id_eia",
        "plant_id_eia",
        "owner_utility_id_eia",
        "unit_id_pudl",
        # "owner_name",
        # "owner_state",
        "fraction_owned",
        "ownership_code",
    ]

    plants_cols = ["plant_id_eia", "utility_id_eia"]

    gens_ownership = grouped_units.merge(
        weighted_ownership[owner_cols], on=["plant_id_eia", "unit_id_pudl"], how="left"
    ).merge(plants[plants_cols], on=["plant_id_eia"], how="left")

    gens_ownership.loc[
        gens_ownership["owner_utility_id_eia"].isnull(), "owner_utility_id_eia"
    ] = gens_ownership.loc[
        gens_ownership["owner_utility_id_eia"].isnull(), "utility_id_eia"
    ]

    # First take care of plants where we have a utility_id to cluster by

    # Jointly owned plants
    mask = (
        (gens_ownership.ownership_code == "J")
        & (
            gens_ownership.owner_utility_id_eia.isin(
                settings["cluster_by_owner_regions"][region]["utility_ids_to_cluster"]
            )
        )
        & (gens_ownership.technology_description == tech)
    )

    gens_ownership.loc[mask, "plant_id_eia"] = (
        gens_ownership.loc[mask, "plant_id_eia"].astype(str)
        + "_"
        + gens_ownership.loc[mask, "owner_utility_id_eia"].astype(int).astype(str)
    )
    gens_ownership.loc[mask, capacity_col] *= gens_ownership.loc[mask, "fraction_owned"]
    gens_ownership.loc[mask, "cluster"] = (
        gens_ownership.loc[mask, "owner_utility_id_eia"].astype(int).astype(str)
    )

    # Single owner - respondent or someone else (S or W)
    mask = (
        (gens_ownership.ownership_code != "J")
        & (
            gens_ownership.owner_utility_id_eia.isin(
                settings["cluster_by_owner_regions"][region]["utility_ids_to_cluster"]
            )
        )
        & (gens_ownership.technology_description == tech)
    )
    gens_ownership.loc[mask, "cluster"] = (
        gens_ownership.loc[mask, "owner_utility_id_eia"].astype(int).astype(str)
    )

    # Now take care of plants where the cluster will be "other"
    other_mask = (
        (gens_ownership.ownership_code == "J")
        & (
            ~gens_ownership.owner_utility_id_eia.isin(
                settings["cluster_by_owner_regions"][region]["utility_ids_to_cluster"]
            )
        )
        & (gens_ownership.technology_description == tech)
    )
    gens_ownership.loc[other_mask, "plant_id_eia"] = (
        gens_ownership.loc[other_mask, "plant_id_eia"].astype(str) + "_" + "other"
    )
    gens_ownership.loc[other_mask, capacity_col] *= gens_ownership.loc[
        other_mask, "fraction_owned"
    ]
    gens_ownership.loc[other_mask, "cluster"] = "other"

    # Single owner - respondent or someone else (S or W)
    other_mask = (
        (gens_ownership.ownership_code != "J")
        & (
            ~gens_ownership.owner_utility_id_eia.isin(
                settings["cluster_by_owner_regions"][region]["utility_ids_to_cluster"]
            )
        )
        & (gens_ownership.technology_description == tech)
    )
    gens_ownership.loc[other_mask, "cluster"] = "other"

    gens_ownership = gens_ownership.drop_duplicates()

    end_capacity = gens_ownership[capacity_col].sum()

    assert np.allclose(start_capacity, end_capacity)
    # assert gens_ownership.loc[gens_ownership.cluster.isna(), :].empty is True

    return gens_ownership


def weighted_ownership_by_unit(units_model, gens_860, ownership, settings):
    owner_cols = [
        "utility_id_eia",
        "plant_id_eia",
        "generator_id",
        "owner_utility_id_eia",
        "owner_name",
        "owner_state",
        "fraction_owned",
    ]

    units_model_cols = [
        "plant_id_eia",
        "generator_id",
        "unit_id_pudl",
        settings["capacity_col"],
    ]

    gens_cols = ["plant_id_eia", "generator_id", "ownership_code"]

    owner_unit = (
        ownership[owner_cols]
        .merge(
            units_model.reset_index()[units_model_cols],
            on=["plant_id_eia", "generator_id"],
            how="right",
        )
        .merge(gens_860[gens_cols], on=["plant_id_eia", "generator_id"], how="left")
    )

    def w_ownership(df, capacity_col):
        weighted_ownership = np.average(df["fraction_owned"], weights=df[capacity_col])
        return weighted_ownership

    owner_unit.loc[:, "fraction_owned"] = owner_unit.loc[:, "fraction_owned"].fillna(1)
    weighted_ownership = (
        owner_unit.groupby(
            ["plant_id_eia", "unit_id_pudl", "owner_utility_id_eia", "ownership_code"],
            as_index=False,
        )
        .apply(w_ownership, settings["capacity_col"])
        .reset_index()
        .rename(columns={0: "fraction_owned"})
        .fillna(0.5)
    )

    weighted_ownership = weighted_ownership.merge(
        ownership[["plant_id_eia", "utility_id_eia"]], on=["plant_id_eia"], how="left"
    )

    return weighted_ownership


# powergenome/eia_opendata.py

"""
Load data from EIA's Open Data API. Requires an api key, which should be included in a
.env file (/powergenome/.env) with the format EIA_API_KEY=YOUR_API_KEY
"""

from itertools import product
import logging
import operator
from typing import Union

import pandas as pd
import requests
import zipfile

from powergenome.params import SETTINGS, DATA_PATHS
from powergenome.price_adjustment import inflation_price_adjustment
from powergenome.util import download_save, reverse_dict_of_lists

logger = logging.getLogger(__name__)

numeric = Union[int, float]


def download_bulk_file(fn: str, retry: bool = False):
    """Download and unzip a bulk data file from EIA

    Parameters
    ----------
    fn : str
        Name of the bulk data file
    retry : bool
        If True, redo the download/save
    """
    if ".zip" not in fn:
        fn = f"{fn}.zip"
    url = f"https://api.eia.gov/bulk/{fn}"

    save_path = DATA_PATHS["eia"] / "bulk_files"
    if not (save_path / fn).exists():
        download_save(url, save_path / fn)
    elif retry:
        download_save(url, save_path / fn)


def extract_bulk_series(aeo_year: int, name_str: str, columns: list = None):
    """Extract relevant data series from an AEO bulk data file.

    Parameters
    ----------
    aeo_year : int
        Year of AEO data
    name_str : str
        String match for something in the "name" field of the JSON
    columns : list, optional
        Names for columns in the dataframe written to file, by default None
    """
    bulk_data_dir = DATA_PATHS["eia"] / "bulk_files"
    save_data_dir = DATA_PATHS["eia"] / "open_data"
    save_data_dir.mkdir(exist_ok=True)
    fn = f"AEO{aeo_year}.zip"
    df = pd.read_json(bulk_data_dir / fn, lines=True)
    df = df.dropna(subset=["series_id"])
    filtered_df = df.query("name.str.contains(@name_str)")
    for row in filtered_df.itertuples():
        series = row.series_id
        _df = pd.DataFrame(row.data, columns=columns)
        _df.to_csv(save_data_dir / f"{series}.csv", index=False, float_format="%g")


def read_eia_api(series_id: str, api_key: str, columns: list = None):
    """Load EIA AEO data from the API and write it to file

    Parameters
    ----------
    series_id : str
        The AEO API series ID that uniquely identifies the data request.
    api_key : str
        A valid API key for EIA's open data portal
    columns : list
        The expected output dataframe columns

    Returns
    -------
    pd.DataFrame
        Data from EIA's AEO via their open data API.
    """
    data_dir = DATA_PATHS["eia"] / "open_data"
    data_dir.mkdir(parents=True, exist_ok=True)
    url = (
        f"https://api.eia.gov/series/?series_id={series_id}&api_key={api_key}&out=json"
    )
    r = requests.get(url)

    try:
        df = pd.DataFrame(r.json()["series"][0]["data"], columns=columns, dtype=float)
        df.to_csv(data_dir / f"{series_id}.csv", index=False)
        return df
    except KeyError:
        print(
            "There was an error creating a dataframe from your EIA AEO data request. "
            f"The constructed series ID is {series_id}. Check to make sure it looks "
            "correct. The data returned from EIA's API is: \n"
            f"{r.json()}"
        )


def load_aeo_series(
    series_id: str, api_key: str, columns: list = None, retry_bulk: bool = False
) -> pd.DataFrame:
    """Load EIA AEO data either from file (if it exists) from a bulk download file or
    from the API.

    Parameters
    ----------
    series_id : str
        The AEO API series ID that uniquely identifies the data request.
    api_key : str
        A valid API key for EIA's open data portal
    columns : list
        The expected output dataframe columns

    Returns
    -------
    pd.DataFrame
        Data from EIA's AEO via their open data API.
    """
    data_dir = DATA_PATHS["eia"] / "open_data"
    data_dir.mkdir(parents=True, exist_ok=True)
    if not (data_dir / f"{series_id}.csv").exists():
        try:
            year = series_id.split(".")[1]
            fn = f"AEO{year}.zip"
            download_bulk_file(fn)
            extract_bulk_series(
                aeo_year=year, name_str="Electricity Demand", columns=["year", "demand"]
            )
            extract_bulk_series(
                aeo_year=year,
                name_str="Energy Prices : Electric Power",
                columns=["year", "price"],
            )
        except (FileNotFoundError, zipfile.BadZipFile):
            try:
                download_bulk_file(fn, retry=True)
                extract_bulk_series(
                    aeo_year=year,
                    name_str="Electricity Demand",
                    columns=["year", "demand"],
                )
                extract_bulk_series(
                    aeo_year=year,
                    name_str="Energy Prices : Electric Power",
                    columns=["year", "price"],
                )
            except (FileNotFoundError, zipfile.BadZipFile):
                df = read_eia_api(series_id, api_key, columns)
                return df

    df = pd.read_csv(data_dir / f"{series_id}.csv")

    return df


def fetch_fuel_prices(settings: dict, inflate_price: bool = True) -> pd.DataFrame:
    """
    Get EIA AEO fuel prices for all regions, fuel types, and scenarios (series IDs)
    included in the settings.

    Parameters
    ----------
    settings : dict
        Should include the following keys:
            eia_aeo_year (int)
            eia_series_region_names (list)
            eia_series_fuel_names (list)
            eia_series_scenario_names (list)
    inflate_price: bool
        If True, adjust the AEO prices to the year "target_usd_year" from the settings.
        If False, do not adjust the AEO prices. Requires the additional settings keys
        "target_usd_year" and "aeo_fuel_usd_year".

    Returns
    -------
    pd.DataFrame
        All fuel price data from AEO for the product of regions, fuels, and scenarios
        included in the settings dictionary.

    Raises
    ------
    KeyError
        The settings parameter "eia_series_scenario_names" is missing.
    KeyError
        The AEO data year is not specified by either of the settings parameters
        "fuel_eia_aeo_year" or "eia_aeo_year".
    TypeError
        The parameter value of "fuel_eia_aeo_year" or "eia_aeo_year" is not an integer.

    Examples
    --------
    Prepare the settings dictionary
    >>> settings = {}
    >>> settings["eia_aeo_year"] = 2020
    >>> settings["aeo_fuel_usd_year"] = 2019
    >>> settings["eia_series_scenario_names"] = {"reference": "REF2020"}
    >>> settings["eia_series_fuel_names"] = {"coal": "STC"}
    >>> settings["eia_series_region_names"] = {"mountain": "MTN"}

    Find the fuel cost with inflating costs.

    >>> fuel_price = fetch_fuel_prices(settings, inflate_price=False)
    >>> print(fuel_price.head())
       year     price  fuel    region   scenario           full_fuel_name
    0  2050  1.501850  coal  mountain  reference  mountain_reference_coal
    1  2049  1.488098  coal  mountain  reference  mountain_reference_coal
    2  2048  1.508208  coal  mountain  reference  mountain_reference_coal
    3  2047  1.506809  coal  mountain  reference  mountain_reference_coal
    4  2046  1.497366  coal  mountain  reference  mountain_reference_coal

    If either of the keys "target_usd_year" or "aeo_fuel_usd_year" is missing, fuel prices
    cannot be inflated.

    >>> fuel_price = fetch_fuel_prices(settings)
    ************
    Unable to inflate fuel prices. Check your settings file to ensure the keys
    "target_usd_year" and "aeo_fuel_usd_year" exist and are valid integers.
    ************
    """
    API_KEY = SETTINGS["EIA_API_KEY"]

    if settings.get("fuel_eia_aeo_year"):
        aeo_year = settings.get("fuel_eia_aeo_year")
    else:
        aeo_year = settings.get("eia_aeo_year")

    if not aeo_year:
        raise KeyError(
            "The parameter 'fuel_eia_aeo_year' is not in your settings files. This is a "
            "required settings parameter when using fuel price data from EIA's AEO."
        )
    if not isinstance(aeo_year, int):
        raise TypeError(
            "The settings parameter 'fuel_eia_aeo_year' (or 'eia_aeo_year') must be an "
            f"integer, representing the AEO data year. Your parameter is {aeo_year}."
        )
    if not settings.get("eia_series_region_names"):
        logger.warning(
            "EIA fuel region names were not found in the settings ('eia_series_region_names'). "
            "Applying default values."
        )
        settings["eia_series_region_names"] = dict(
            mountain="MTN",
            pacific="PCF",
            west_south_central="WSC",
            east_south_central="ESC",
            south_atlantic="SOATL",
            west_north_central="WNC",
            east_north_central="ENC",
            middle_atlantic="MDATL",
            new_england="NEENGL",
        )
    if not settings.get("eia_series_fuel_names"):
        settings["eia_series_fuel_names"] = {
            "coal": "STC",
            "naturalgas": "NG",
            "distillate": "DFO",
            "uranium": "U",
        }
        logger.warning(
            "EIA fuel names were not found in the settings ('eia_series_fuel_names'). "
            "Applying default values:\n\n"
            f"{settings['eia_series_fuel_names']}"
        )
    if not settings.get("eia_series_scenario_names"):
        raise KeyError(
            "The settings parameter 'eia_series_scenario_names' is missing. This mapping "
            "of AEO scenario API names (e.g. REF2020) to plain english is required."
        )

    fuel_price_cases = product(
        settings.get("eia_series_region_names", {}).items(),
        settings.get("eia_series_fuel_names", {}).items(),
        settings.get("eia_series_scenario_names", {}).items(),
    )
    if not aeo_year or not fuel_price_cases:
        w = False
        for f in ["coal", "naturalgas", "distillate", "uranium"]:
            if f in settings.get("tech_fuel_map", {}).values():
                w = True
        if w:
            logger.warning(
                "Unable to get AEO fuel prices due to missing settings parameter 'eia_aeo_year', "
                "'eia_series_region_names', 'eia_series_fuel_names', or 'eia_series_scenario_names'. "
                "You have listed at least one AEO fuel in your settings 'tech_fuel_map' "
                "parameter, but no prices for these fuels are being included."
            )
        return pd.DataFrame(
            columns=["fuel", "region", "scenario", "full_fuel_name", "year"]
        )
    df_list = []
    for region, fuel, scenario in fuel_price_cases:
        region_name, region_series = region
        fuel_name, fuel_series = fuel
        scenario_name, scenario_series = scenario

        SERIES_ID = f"AEO.{aeo_year}.{scenario_series}.PRCE_REAL_ELEP_NA_{fuel_series}_NA_{region_series}_Y13DLRPMMBTU.A"

        df = load_aeo_series(
            series_id=SERIES_ID, api_key=API_KEY, columns=["year", "price"]
        )
        df["fuel"] = fuel_name
        df["region"] = region_name
        df["scenario"] = scenario_name
        df["full_fuel_name"] = df.region + "_" + df.scenario + "_" + df.fuel
        df["year"] = df["year"].astype(int)

        df_list.append(df)

    final = pd.concat(df_list, ignore_index=True)

    if inflate_price:
        try:
            fuel_price_base_year = settings["aeo_fuel_usd_year"]
            fuel_price_target_year = settings["target_usd_year"]
            final.loc[:, "price"] = inflation_price_adjustment(
                price=final.loc[:, "price"],
                base_year=fuel_price_base_year,
                target_year=fuel_price_target_year,
            )
        except (KeyError, TypeError):
            logger.warning(
                """
    ************
    Unable to inflate fuel prices. Check your settings file to ensure the keys
    "target_usd_year" and "aeo_fuel_usd_year" are valid integers.
    ************
                """
            )

    return final


def modify_fuel_prices(
    prices: pd.DataFrame,
    aeo_fuel_region_map: dict,
    regional_fuel_adjustments: dict = None,
) -> pd.DataFrame:
    """Modify the AEO fuel prices by model region or fuel within a model region.

    Parameters
    ----------
    prices : pd.DataFrame
        Fuel prices from AEO, with columns ['year', 'price', 'fuel', 'region', 'scenario',
        'full_fuel_name']
    aeo_fuel_region_map : dict
        Mapping of AEO census division fuel names to lists of model regions
    regional_fuel_adjustments : dict, optional
        Modifications of fuel prices by region or fuel within region, by default None

    Returns
    -------
    pd.DataFrame
        Full input dataframe with modified copies for model regions and fuels specified
        in `regional_fuel_adjustments`.

    Raises
    ------
    KeyError
        The required parameter 'aeo_fuel_region_map' is missing
    KeyError
        One or more model regions having fuel prices modified is not in `aeo_fuel_region_map`
    KeyError
        Invalid operator type
    KeyError
        Invalid fuel name
    KeyError
        Invalid operator type
    TypeError
        Fuel price modifiers are not a list or a dictionary of lists
    """

    if not regional_fuel_adjustments:
        return prices

    if not aeo_fuel_region_map:
        raise KeyError("The required parameter 'aeo_fuel_region_map' is missing.")

    allowed_operators = ["add", "mul", "truediv", "sub"]
    model_regions = list(regional_fuel_adjustments)
    model_aeo_region_map = reverse_dict_of_lists(aeo_fuel_region_map)
    if not all(r in model_aeo_region_map for r in model_regions):
        raise KeyError(
            "All model regions listed in the settings parameter 'regional_fuel_adjustments' "
            "should also be included in `aeo_fuel_region_map`. One or more regions was "
            "not found."
        )

    df_list = []
    for region, adj in regional_fuel_adjustments.items():
        aeo_region = model_aeo_region_map[region]
        if isinstance(adj, list):
            op, op_value = adj
            if op not in allowed_operators:
                raise KeyError(
                    f"The regional fuel price adjustment for {region} needs a valid "
                    f"operator from the list\n{allowed_operators}\n"
                    "in the format [<operator>, <value>].\n"
                )
            f = operator.attrgetter(op)
            df = prices.loc[prices["region"] == aeo_region, :]
            df.loc[:, "region"] = region
            df.loc[:, "price"] = f(operator)(df["price"], op_value)
            df.loc[:, "full_fuel_name"] = df["full_fuel_name"].str.replace(
                aeo_region, region
            )
            df_list.append(df)
        elif isinstance(adj, dict):
            for fuel, op_list in adj.items():
                if fuel not in prices["fuel"].unique():
                    raise KeyError(
                        f"The fuel '{fuel}' is listed under the region {region} in your settings "
                        "parameter 'regional_fuel_adjustments'. There was no AEO fuel "
                        "price fetched for this fuel so it cannot be modified."
                    )
                op, op_value = op_list
                if op not in allowed_operators:
                    raise KeyError(
                        f"The regional fuel price adjustment for '{fuel}' in {region} "
                        f"needs to be an operator from the list {allowed_operators}. "
                        f"You supplied '{op}', which is not a valid operator."
                    )
                f = operator.attrgetter(op)
                df = prices.loc[
                    (prices["region"] == aeo_region) & (prices["fuel"] == fuel.lower()),
                    :,
                ]
                df.loc[:, "region"] = region
                df.loc[:, "price"] = f(operator)(df["price"], op_value)
                df.loc[:, "full_fuel_name"] = df["full_fuel_name"].str.replace(
                    aeo_region, region
                )
                df_list.append(df)
        else:
            raise TypeError(
                "Fuel price modifiers in the settings parameter 'regional_fuel_adjustments' "
                "must be a list of the form '[<op>, <value>]', or a similar list for a "
                "specific fuel. "
                f"Your value look like '{adj}' for region '{region}'."
            )

    mod_prices = pd.concat([prices] + df_list, ignore_index=True, sort=False)

    return mod_prices


def add_user_fuel_prices(settings: dict, df: pd.DataFrame = None) -> pd.DataFrame:
    """Add user fuel prices to a dataframe of user prices from AEO (or elsewhere)

    Parameters
    ----------
    settings : dict
        If adding user prices, should have the key "user_fuel_price" with value of a
        dictionary matching user fuel names and prices. Prices can either be a single
        price for all regions or a price per region. For example this shows biomass with
        different prices in two regions and ZCF with the same price in all regions:

        settings["user_fuel_price"] = {
            "biomass": {"SC_VACA": 10, "PJM_DOM": 5},
            "ZCF": 15
        }

        If the keys "target_usd_year" and "user_fuel_usd_year" are also included, fuel
        prices will be corrected to the correct USD year. "user_fuel_usd_year" should
        be a dictionary with fuel name: USD year pairings. Only fuels included in this
        dictionary will have their prices changed to the target USD year.
    df : pd.DataFrame, optional
        A dataframe with fuel prices from AEO (or elsewhere), by default None. Should
        have columns ["year", "price", "fuel", "region", "scenario", "full_fuel_name"]

    Returns
    -------
    pd.DataFrame
        The combined dataframes of user prices and the other price dataframe provided
        as input. Columns are ["year", "price", "fuel", "region", "scenario", "full_fuel_name"].
    """

    if not settings.get("user_fuel_price"):
        if df is not None:
            return df
    cols = ["year", "price", "fuel", "region", "scenario", "full_fuel_name"]
    if df is not None and not df.empty:
        years = df["year"].unique()
    else:
        years = range(2020, 2051)
    fuel_data = {c: [] for c in cols}

    for fuel, val in settings["user_fuel_price"].items():
        if isinstance(val, dict):
            for region, price in val.items():
                fuel_name = f"{region}_{fuel}"
                fuel_data["year"].extend(years)
                fuel_data["price"].extend([price] * len(years))
                fuel_data["fuel"].extend([fuel] * len(years))
                fuel_data["region"].extend([region] * len(years))
                fuel_data["scenario"].extend(["user"] * len(years))
                fuel_data["full_fuel_name"].extend([fuel_name] * len(years))
        else:
            fuel_data["year"].extend(years)
            fuel_data["price"].extend([val] * len(years))
            fuel_data["fuel"].extend([fuel] * len(years))
            fuel_data["region"].extend([""] * len(years))
            fuel_data["scenario"].extend(["user"] * len(years))
            fuel_data["full_fuel_name"].extend([fuel] * len(years))

    user_fuel_price = pd.DataFrame(fuel_data)
    if settings.get("target_usd_year"):
        for fuel, year in (settings.get("user_fuel_usd_year", {}) or {}).items():
            user_fuel_price.loc[
                user_fuel_price["fuel"] == fuel, "price"
            ] = inflation_price_adjustment(
                user_fuel_price.loc[user_fuel_price["fuel"] == fuel, "price"],
                year,
                settings["target_usd_year"],
            )
    if df is not None:
        user_fuel_price = pd.concat([df, user_fuel_price])
    return user_fuel_price


def get_aeo_load(
    region: str,
    aeo_year: Union[str, numeric],
    scenario_series: str,
    sector: str = "ELEP",
) -> pd.DataFrame:
    """Find the electricity demand in a single AEO region. Use EIA API if data has not
    been previously saved.

    Parameters
    ----------
    region : str
        Short name of the AEO region
    aeo_year : Union[str, numeric]
        AEO data year
    scenario_series : str
        Short name of the AEO scenario

    Returns
    -------
    pd.DataFrame
        The demand data for a single region.

    Examples
    --------
    >>> texas_load = get_aeo_load("TRE", 2020, "REF2020")
    >>> print(texas_load.head())
       year      demand
    0  2050  489.009247
    1  2049  483.176544
    2  2048  477.624481
    3  2047  472.314972
    4  2046  466.875671
    """
    API_KEY = SETTINGS["EIA_API_KEY"]

    SERIES_ID = (
        f"AEO.{aeo_year}.{scenario_series}.CNSM_NA_{sector}_NA_ELC_NA_{region}_BLNKWH.A"
    )

    df = load_aeo_series(SERIES_ID, API_KEY, columns=["year", "demand"])
    df["year"] = df["year"].astype(int)

    return df


# powergenome/external_data.py

# Read in and add external inputs (user-supplied files) to PowerGenome outputs

import logging
import numpy as np
import pandas as pd
from pathlib import Path
from typing import Any, List

from powergenome.util import snake_case_col, remove_feb_29
from powergenome.price_adjustment import inflation_price_adjustment

logger = logging.getLogger(__name__)


def make_demand_response_profiles(
    path: Path, resource_name: str, year: int, scenario: str
) -> pd.DataFrame:
    """Read files with DR profiles across years and scenarios. Return the hourly
    load profiles for a single resource in the model year.

    Parameters
    ----------
    path : path-like
        Where to load the file from
    resource_name : str
        Name of of the demand response resource
    year : int
        Year of data to use from the user demand response file
    scenario : str
        Name of scenario to use from the user demand response file

    Returns
    -------
    DataFrame
        Hourly profiles of DR load for each region where the resource is available.
        Column names are the regions plus 'scenario'.
    """

    df = pd.read_csv(path, header=[0, 1, 2, 3])

    # Use the MultiIndex columns to just get columns with the correct resource listed
    # in the top row of the csv. The resource name is dropped from the columns.
    resource_df = df.loc[:, resource_name]

    assert year in set(
        resource_df.columns.get_level_values(0).astype(int)
    ), f"The model year is not in the years of data for DR resource {resource_name}"

    resource_df = resource_df.loc[:, str(year)]

    assert scenario in set(
        resource_df.columns.get_level_values(0)
    ), f"The scenario {scenario} is not included for DR resource {resource_name}"

    resource_df = resource_df.loc[:, scenario]
    resource_df = resource_df.reset_index(drop=True)

    if len(resource_df) == 8784:
        remove_feb_29(resource_df)

    return resource_df


def demand_response_resource_capacity(df, resource_name, settings):
    """Calculate the maximum capacity value to assign a demand response/DSM resource

    Parameters
    ----------
    df : DataFrame
        Hourly demand profile in each region of a single DR resource
    resource_name : str
        Name of the resource, which should match the settings file
    settings : dict
        User-defined parameters from a settings file

    Returns
    -------
    DataFrame
        Index of scenarios and columns of regions, values represent shiftable capacity.
    """

    year = settings["model_year"]
    fraction_shiftable = settings["flexible_demand_resources"][year][resource_name][
        "fraction_shiftable"
    ]

    # peak_load = df.groupby(["scenario"]).max()
    peak_load = df.max()
    shiftable_capacity = peak_load * fraction_shiftable

    return shiftable_capacity


def add_resource_max_cap_spur(
    new_resource_df: pd.DataFrame, settings: dict, capacity_col: str = "Max_Cap_MW"
) -> pd.DataFrame:
    """Load user supplied maximum capacity and spur line data for new resources. Add
    those values to the resources dataframe.

    Parameters
    ----------
    new_resource_df : pd.DataFrame
        New resources that can be built. Each row should be a single resource, with
        columns 'region' and 'technology'. The number of copies for a region/resource
        (e.g. more than one UtilityPV resource in a region) should match what is
        given in the user-supplied file.
    settings : dict
        User-defined parameters from a settings file. Should have keys of `input_folder`
        (a Path object of where to find user-supplied data) and
        `capacity_limit_spur_fn` (the file to load).
    capacity_col : str, optional
        The column that indicates maximum capacity constraints, by default "Max_Cap_MW"

    Returns
    -------
    pd.DataFrame
        A modified version of new_resource_df with spur_miles and the maximum
        capacity for a resource. Copies of a resource within a region should have a
        cluster name to uniquely identify them.
    """

    defaults = {
        "cluster": 1,
        "spur_miles": 0,
        capacity_col: -1,
        "interconnect_annuity": 0,
        "interconnect_capex_mw": 0,
    }
    # Prepare file
    path = Path(settings["input_folder"]) / settings["capacity_limit_spur_fn"]
    df = pd.read_csv(path)
    for col in "region", "technology":
        if col not in df:
            raise KeyError(f"The max capacity/spur file must have column {col}")
    for key, value in defaults.items():
        df[key] = df[key].fillna(value) if key in df else value
        new_resource_df[key] = (
            new_resource_df[key].fillna(value) if key in new_resource_df else value
        )
    # Update resources
    grouped_df = df.groupby(["region", "technology"])
    for (region, tech), _df in grouped_df:
        mask = (new_resource_df["region"] == region) & (
            new_resource_df["technology"]
            .str.replace("_\*", "_all")
            .str.contains(tech.replace("_\*", "_all"), case=False)
        )
        if mask.sum() > 1:
            resources = new_resource_df.loc[mask, "technology"].to_list()
            raise ValueError(
                f"Resource {tech} in region {region} from file "
                f"{settings['capacity_limit_spur_fn']} matches multiple resources:"
                f"\n{resources}"
            )
        else:
            pass
        for key, value in defaults.items():
            _key = "max_capacity" if key == capacity_col else key
            new_resource_df.loc[mask & (new_resource_df[key] == value), key] = _df[
                _key
            ].values
    logger.info(
        f"Inflating external interconnect annuity costs from 2017 to "
        f"{settings['target_usd_year']}"
    )
    new_resource_df["interconnect_annuity"] = inflation_price_adjustment(
        new_resource_df["interconnect_annuity"], 2017, settings["target_usd_year"]
    )
    return new_resource_df


def make_generator_variability(
    df: pd.DataFrame, remove_feb_29: bool = True
) -> pd.DataFrame:
    """Make a generator variability dataframe with normalized (0-1) hourly profiles
    for each resource in resource_df.

    Any resources that do not have a profile in column
    `profile` are assumed to have constant hourly profiles with a value of 1.
    February 29 is removed from any profiles of length 8784 (leap year).

    Parameters
    ----------
    df
        Dataframe with a single row for each resource. May have column `profile`.

    Returns
    -------
    pd.DataFrame
        A new dataframe with one column for each row of `df` and 8760 rows
        of generation profiles.

    Examples
    --------
    >>> df = pd.DataFrame({'profile': [[0] * 8760, np.ones(8784) / 2, None]})
    >>> make_generator_variability(df)
            0    1    2
    0     0.0  0.5  1.0
    1     0.0  0.5  1.0
    2     0.0  0.5  1.0
    3     0.0  0.5  1.0
    4     0.0  0.5  1.0
    ...   ...  ...  ...
    8755  0.0  0.5  1.0
    8756  0.0  0.5  1.0
    8757  0.0  0.5  1.0
    8758  0.0  0.5  1.0
    8759  0.0  0.5  1.0
    <BLANKLINE>
    [8760 rows x 3 columns]
    """

    def profile_len(x: Any) -> int:
        if isinstance(x, (list, np.ndarray)):
            return len(x)
        return 1

    def format_profile(
        x: Any, remove_feb_29: bool = True, hours: int = 8760
    ) -> np.ndarray:
        # from IPython import embed
        # embed()
        if isinstance(x, np.ndarray):
            if len(x) == 8784:
                # Remove February 29 from leap year
                if remove_feb_29:
                    return np.delete(x, slice(1416, 1440))
            return x
        if isinstance(x, list):
            return format_profile(np.array(x), remove_feb_29, hours)
        # Fill missing with default [1, ...]
        return np.ones(hours, dtype=float)

    if "profile" in df:
        hours = df["profile"].apply(profile_len).max()
        if remove_feb_29 and hours == 8784:
            hours = 8760
        kwargs = {"remove_feb_29": remove_feb_29, "hours": hours}
        profiles = np.column_stack(df["profile"].apply(format_profile, **kwargs).values)
        # Make sure values are not less than 0
        profiles = np.where(profiles >= 0, profiles, 0)
    # elif not remove_feb_29:
    #     profiles = np.ones((8760, len(df)), dtype=float)
    else:
        profiles = np.ones((8760, len(df)), dtype=float)
    return pd.DataFrame(profiles, columns=np.arange(len(df)).astype(str))


def load_policy_scenarios(settings: dict) -> pd.DataFrame:
    """Load the policy scenarios and copy cases where indicated. The policy file should
    start with columns `case_id` and `year`, and can contain an optional `copy_case_id`.
    Other columns should match the desired output format. The value `None` is included
    in na_values.

    Parameters
    ----------
    settings : dict
        User-defined parameters from a settings file. Should have keys of `input_folder`
        (a Path object of where to find user-supplied data) and
        `emission_policies_fn` (the file to load).

    Returns
    -------
    DataFrame
        Emission policies for each case_id/year.
    """

    path = Path(settings["input_folder"]) / settings["emission_policies_fn"]
    policies = pd.read_csv(path, na_values=["None", "none"])

    # Update the policies. The column `copy_case_id` can be used to copy values from
    # another policy to reduce human copy/paste errors.
    if "copy_case_id" in policies.columns:
        policies = copy_case_values(policies, match_cols=["case_id", "year", "region"])

    policies = policies.drop(columns="copy_case_id")
    policies = policies.set_index(["case_id", "year"])

    return policies


def copy_case_values(df, match_cols):
    """Copy values for one case to others in an external data file. Must have column
    "copy_case_id"

    Parameters
    ----------
    df : DataFrame
        Policies, settings, or some other type of data for each case/year.
    match_cols : list
        The columns used to match cases. Must include "case_id", can also include things
        like "year" and "region".

    Returns
    -------
    DataFrame
        Modified copy of the original with values copied to rows with valid
        "copy_case_id" values.
    """
    match_cols_len = len(match_cols)
    settings_cols = [
        col for col in df.columns if col not in match_cols + ["copy_case_id"]
    ]
    for row in df.itertuples(index=False):
        if not pd.isna(row.copy_case_id):
            # Create a dictionary with keys of each col from match_cols and values
            # from the row. Then replace case_id key with copy_case_id value. This way
            # we get the new case id but same region, year, etc.
            filter_dict = {
                col: value
                for col, value in zip(df.columns[:match_cols_len], row[:match_cols_len])
            }
            filter_dict["case_id"] = row.copy_case_id  # Set to the case we want to copy

            # Use the filter dictionary to filter the original df and assign the values
            # from that row to the row in question.
            # https://stackoverflow.com/a/34162576
            new_values = df.loc[
                (df[list(filter_dict)] == pd.Series(filter_dict)).all(axis=1),
                settings_cols,
            ].values

            filter_dict["case_id"] = row.case_id  # Set back to case we want to modify

            df.loc[
                (df[list(filter_dict)] == pd.Series(filter_dict)).all(axis=1),
                settings_cols,
            ] = new_values

    return df


def load_demand_segments(settings):
    """Load a file of demand segments as defined by the user in CSV.

    Parameters
    ----------
    settings : dict
        User defined PowerGenome settings. Must have the keys "input_folder" and
        "demand_segments_fn".

    Returns
    -------
    DataFrame
        Demand segments with columns such as "Voll", "Demand_segment", etc.
    """

    path = Path(settings["input_folder"]) / settings["demand_segments_fn"]
    demand_segments = pd.read_csv(path)

    return demand_segments


def load_user_genx_settings(settings):
    """Load a file with user-supplied GenX settings and copy values to cases where
    needed.

    Parameters
    ----------
    settings : dict
        User defined PowerGenome settings. Must have the keys "input_folder" and
        "case_genx_settings_fn".

    Returns
    -------
    DataFrame
        Values for GenX settings across cases and years as defined by the user. The
        index is ["case_id", "year"] and columns are GenX parameters.
    """

    path = Path(settings["input_folder"]) / settings["case_genx_settings_fn"]
    genx_case_settings = pd.read_csv(path)

    if "copy_case_id" in genx_case_settings.columns:
        genx_case_settings = copy_case_values(
            genx_case_settings, match_cols=["case_id", "year"]
        )

    genx_case_settings = genx_case_settings.drop(columns=["copy_case_id"])
    genx_case_settings = genx_case_settings.set_index(["case_id", "year"])

    return genx_case_settings


def overwrite_wind_pv_capacity(df, settings):
    """Use external data to overwrite the wind and solarpv capacity extracted from
    EIA860.

    Parameters
    ----------
    df : DataFrame
        Existing generators dataframe, with columns "region", "technology", and
        "Existing_Cap_MW". The technologies should include "Solar Photovoltaic"
        and "Onshore Wind Turbine".
    settings : dict
        User defined PowerGenome settings. Must have the keys "input_folder" and
        "region_wind_pv_cap_fn".

    Returns
    -------
    DataFrame
        Same as input dataframe but with new capacity values for technologies defined
        in the "region_wind_pv_cap_fn" file.
    """
    from powergenome.util import reverse_dict_of_lists

    idx = pd.IndexSlice

    path = settings["input_folder"] / settings["region_wind_pv_cap_fn"]

    wind_pv_ipm_region_capacity = pd.read_csv(path)

    region_agg_map = reverse_dict_of_lists(settings.get("region_aggregations", {}))

    # Set model_region as IPM_region to start
    wind_pv_ipm_region_capacity["model_region"] = wind_pv_ipm_region_capacity[
        "IPM_Region"
    ]
    # Change any aggregated regions to the user-defined model_region
    wind_pv_ipm_region_capacity.loc[
        wind_pv_ipm_region_capacity["IPM_Region"].isin(region_agg_map), "model_region"
    ] = wind_pv_ipm_region_capacity.loc[
        wind_pv_ipm_region_capacity["IPM_Region"].isin(region_agg_map), "IPM_Region"
    ].map(
        region_agg_map
    )
    wind_pv_model_region_capacity = wind_pv_ipm_region_capacity.groupby(
        ["model_region", "technology"]
    ).sum()

    df = df.reset_index()

    for region in df["region"].unique():
        for tech in ["Solar Photovoltaic", "Onshore Wind Turbine"]:
            if tech in df.query("region == @region")["technology"].to_list():
                try:
                    df.loc[
                        (df["region"] == region) & (df["technology"] == tech),
                        "Existing_Cap_MW",
                    ] = wind_pv_model_region_capacity.loc[
                        idx[region, tech], "nameplate_capacity_mw"
                    ]
                except:
                    pass

    df = df.set_index(["region", "technology", "cluster"])

    return df


def make_usr_demand_profiles(path, settings):
    idx = pd.IndexSlice
    year = settings["model_year"]
    scenario = settings.get("electrification")
    if not scenario:
        scenario_file = settings.get("scenario_definitions_fn")
        load_file = settings.get("regional_load_fn")
        raise KeyError(
            f"The scenario definitions file {scenario_file} must have a column "
            f"'electrification' with values that correspond to the second row of "
            f"{load_file}"
        )

    df = pd.read_csv(path, header=[0, 1, 2])
    scenario_df = df.loc[:, idx[str(year), scenario]]

    return scenario_df


def load_user_tx_costs(
    path: Path, model_regions: List[str], target_usd_year: int = None
) -> pd.DataFrame:
    """Load a user data file with cost and line loss of each interregional transmission
    line. Map the region names to zones (z1 to zM) and adjust the total cost columns
    to the dollar year specified in settings.

    Parameters
    ----------
    path : Path
        Path to the CSV file, which should have columns "start_region", "dest_region",
        "total_interconnect_annuity_mw", "total_interconnect_cost_mw", and "dollar_year".
    model_regions : List[str]
        List of model region names. Should be sorted to match order in other functions.
    target_usd_year : int, optional
        Desired final dollar year of cost columns, by default None. If None, no adjustment
        is made.

    Returns
    -------
    pd.DataFrame
        Cost and line loss of each potential transmission line between model regions.
        Contains columns "start_region", "dest_region", "zone_1", "zone_2",
        "total_interconnect_annuity_mw", "total_interconnect_cost_mw", "dollar_year",
        and "adjusted_dollar_year".
    """
    df = pd.read_csv(path)

    zones = model_regions
    zone_num_map = {
        zone: f"z{number + 1}" for zone, number in zip(zones, range(len(zones)))
    }
    df["zone_1"] = df["start_region"].map(zone_num_map)
    df["zone_2"] = df["dest_region"].map(zone_num_map)

    if target_usd_year:
        adjusted_annuities = []
        adjusted_costs = []
        for row in df.itertuples():
            adj_annuity = inflation_price_adjustment(
                row.total_interconnect_annuity_mw, row.dollar_year, target_usd_year
            )
            adjusted_annuities.append(adj_annuity)

            adj_cost = inflation_price_adjustment(
                row.total_interconnect_cost_mw, row.dollar_year, target_usd_year
            )
            adjusted_costs.append(adj_cost)
        df["total_interconnect_annuity_mw"] = adjusted_annuities
        df["total_interconnect_cost_mw"] = adjusted_costs
        df["adjusted_dollar_year"] = target_usd_year

    return df


def insert_user_tx_costs(tx_df: pd.DataFrame, user_costs: pd.DataFrame) -> pd.DataFrame:
    """Insert costs and line loss from the user. Can include more lines than were in the
    original transmission dataframe to create lines with zero existing capacity.

    Parameters
    ----------
    tx_df : pd.DataFrame
        Dataframe of interregional transmission lines created by PG. Must have the column
        "Network_Lines" (integer from 1 to N), a column for each network zone (z1 to zM),
        "Line_Max_Flow_MW", "Line_Min_Flow_MW", and "transmission_path_name".
    user_costs : pd.DataFrame
        Dataframe of costs and line loss created by the user. Should have columns
        "zone_1" and "zone_2" (possible values are z1 through zM), "total_interconnect_annuity_mw",
        "total_interconnect_cost_mw", and "total_line_loss_frac". Cost values should
        already be adjusted to the desired dollar year.

    Returns
    -------
    pd.DataFrame
        Supplemented interregional transmission lines
    """

    unused_lines = []
    for row in user_costs.itertuples():
        line_row = tx_df.loc[(tx_df[row.zone_1] != 0) & (tx_df[row.zone_2] != 0), :]
        assert not len(line_row) > 1

        if line_row.empty:
            unused_lines.append(row)
        tx_df.loc[
            line_row.index, "Line_Reinforcement_Cost_per_MWyr"
        ] = row.total_interconnect_annuity_mw
        tx_df.loc[
            line_row.index, "Line_Reinforcement_Cost_per_MW"
        ] = row.total_interconnect_cost_mw
        tx_df.loc[line_row.index, "Line_Loss_Percentage"] = row.total_line_loss_frac

    unused_line_df = pd.DataFrame(unused_lines)
    unused_line_df = unused_line_df.rename(
        columns={
            "total_interconnect_annuity_mw": "Line_Reinforcement_Cost_per_MWyr",
            "total_interconnect_cost_mw": "Line_Reinforcement_Cost_per_MW",
            "total_line_loss_frac": "Line_Loss_Percentage",
        }
    )

    zone_cols = [c for c in tx_df.columns if c[0] == "z"]
    unused_line_df[zone_cols] = 0
    for idx, row in unused_line_df.iterrows():
        unused_line_df.loc[idx, row["zone_1"]] = 1
        unused_line_df.loc[idx, row["zone_2"]] = -1
        unused_line_df.loc[
            idx, "transmission_path_name"
        ] = f"{row.start_region}_to_{row.dest_region}"

    cols = [c for c in tx_df.columns if c in unused_line_df.columns]
    tx_df = pd.concat([tx_df, unused_line_df[cols]], ignore_index=True)
    tx_df["Network_Lines"] = range(1, len(tx_df) + 1)
    tx_df["Line_Max_Flow_MW"] = tx_df["Line_Max_Flow_MW"].fillna(0)
    tx_df["Line_Min_Flow_MW"] = tx_df["Line_Max_Flow_MW"].fillna(0)

    return tx_df


# powergenome/extract_pudl_data.py

import argparse
import logging
import shutil
import sys
from datetime import datetime as dt

import pandas as pd
import powergenome
from powergenome.fuels import fuel_cost_table
from powergenome.generators import GeneratorClusters, load_ipm_shapefile
from powergenome.load_profiles import make_final_load_curves
from powergenome.params import DATA_PATHS
from powergenome.transmission import (
    agg_transmission_constraints,
    transmission_line_distance,
)
from powergenome.util import (
    get_git_hash,
    init_pudl_connection,
    load_settings,
    remove_fuel_scenario_name,
)

if not sys.warnoptions:
    import warnings

    warnings.simplefilter("ignore")


def parse_command_line(argv):
    """
    Parse command line arguments. See the -h option.

    :param argv: arguments on the command line must include caller file name.
    """
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "-sf",
        "--settings_file",
        dest="settings_file",
        type=str,
        default="pudl_data_extraction.yml",
        help="Specify a YAML settings file.",
    )
    parser.add_argument(
        "-rf",
        "--results_folder",
        dest="results_folder",
        type=str,
        default=dt.now().strftime("%Y-%m-%d %H.%M.%S"),
        help="Specify the results subfolder to write output",
    )
    parser.add_argument(
        "--no-current-gens",
        dest="current_gens",
        action="store_false",
        help="Don't load and cluster current generators.",
    )
    parser.add_argument(
        # "-g",
        "--no-gens",
        dest="gens",
        # type=bool,
        action="store_false",
        # default=True,
        help="Use flag to not calculate generator clusters.",
    )
    parser.add_argument(
        # "-l",
        "--no-load",
        dest="load",
        # type=bool,
        action="store_false",
        # default=True,
        help="Calculate hourly load. If False, file will not be written.",
    )
    parser.add_argument(
        # "-t",
        "--no-transmission",
        dest="transmission",
        # type=bool,
        action="store_false",
        # default=True,
        help="Calculate transmission constraints. If False, file will not be written.",
    )
    parser.add_argument(
        "-f",
        "--no-fuel",
        dest="fuel",
        # type=bool,
        action="store_false",
        # default=True,
        help=(
            "Create fuel table. If False, file will not be written."
            " Can not be created without the generators."
        ),
    )
    parser.add_argument(
        "-s",
        "--sort-gens",
        dest="sort_gens",
        # type=bool,
        action="store_true",
        # default=True,
        help=(
            "Sort generators alphabetically within region. Existing resources will still"
            " be separate from new resources."
        ),
    )
    arguments = parser.parse_args(argv[1:])
    return arguments


def main():
    args = parse_command_line(sys.argv)

    out_folder = DATA_PATHS["results"] / args.results_folder
    DATA_PATHS["results"].mkdir(exist_ok=True)
    out_folder.mkdir(exist_ok=True)

    # Create a logger to output any messages we might have...
    logger = logging.getLogger(powergenome.__name__)
    logger.setLevel(logging.INFO)
    handler = logging.StreamHandler()
    formatter = logging.Formatter(
        # More extensive test-like formatter...
        "%(asctime)s [%(levelname)8s] %(name)s:%(lineno)s %(message)s",
        # This is the datetime format string.
        "%Y-%m-%d %H:%M:%S",
    )
    handler.setFormatter(formatter)
    logger.addHandler(handler)

    filehandler = logging.FileHandler(out_folder / "log.txt")
    filehandler.setFormatter(formatter)
    logger.addHandler(filehandler)

    git_hash = get_git_hash()
    logger.info(f"Current git hash is {git_hash}")

    logger.info("Reading settings file")
    settings = load_settings(path=args.settings_file)

    # Copy the settings file to results folder
    shutil.copy(args.settings_file, out_folder)

    logger.info("Initiating PUDL connections")
    pudl_engine, pudl_out = init_pudl_connection(freq="YS")

    # Make sure everything in model_regions is either an aggregate region
    # or an IPM region. Will need to change this once we start using non-IPM
    # regions.
    ipm_regions = pd.read_sql_table("regions_entity_epaipm", pudl_engine)[
        "region_id_epaipm"
    ]
    all_valid_regions = ipm_regions.tolist() + list(
        settings.get("region_aggregations", {})
    )
    good_regions = [region in all_valid_regions for region in settings["model_regions"]]

    # assert all(good_regions), logger.warning(
    #     "One or more model regions is not valid. Check to make sure all"
    #     "regions are either in IPM or region_aggregations in the settings YAML file."
    # )
    if not all(good_regions):
        logger.warning(
            "One or more model regions is not valid. Check to make sure all regions "
            "are either in IPM or region_aggregations in the settings YAML file."
        )

    # Sort zones in the settings to make sure they are correctly sorted everywhere.
    settings["model_regions"] = sorted(settings["model_regions"])
    zones = settings["model_regions"]
    logger.info(f"Sorted zones are {', '.join(zones)}")
    zone_num_map = {
        zone: f"{number + 1}" for zone, number in zip(zones, range(len(zones)))
    }

    if args.gens:
        gc = GeneratorClusters(
            pudl_engine=pudl_engine,
            pudl_out=pudl_out,
            settings=settings,
            current_gens=args.current_gens,
            sort_gens=args.sort_gens,
        )
        gen_clusters = gc.create_all_generators()
        gen_clusters = remove_fuel_scenario_name(gen_clusters, settings)
        gen_clusters["Zone"] = gen_clusters["region"].map(zone_num_map)

    if args.load:
        load = make_final_load_curves(pudl_engine=pudl_engine, settings=settings)
        load.columns = "Load_MW_z" + load.columns.map(zone_num_map)

    if args.transmission:
        if args.gens is False:
            model_regions_gdf = load_ipm_shapefile(settings)
        else:
            model_regions_gdf = gc.model_regions_gdf
        transmission = agg_transmission_constraints(
            pudl_engine=pudl_engine, settings=settings
        )
        transmission = transmission.pipe(
            transmission_line_distance,
            ipm_shapefile=model_regions_gdf,
            settings=settings,
            units="mile",
        )

    if args.fuel and args.gens:
        fuels = fuel_cost_table(
            fuel_costs=gc.fuel_prices, generators=gc.all_resources, settings=settings
        )
        fuels["fuel_indices"] = range(1, len(fuels) + 1)
        fuels = remove_fuel_scenario_name(fuels, settings)

    logger.info(f"Write GenX input files to {args.results_folder}")
    if args.gens:
        gen_clusters.to_csv(
            out_folder / f"generator_clusters_{args.results_folder}.csv",
            index=False
            # float_format="%.3f",
        )
        # if args.all_units is True:
        # gc.all_units.to_csv(out_folder / f"all_units_{args.results_folder}.csv")

    if args.load:
        load.astype(int).to_csv(out_folder / f"load_curves_{args.results_folder}.csv")

    if args.transmission:
        transmission.to_csv(
            out_folder / f"transmission_constraints_{args.results_folder}.csv",
            float_format="%.1f",
        )

    if args.fuel and args.gens:
        fuels.to_csv(out_folder / f"Fuels_data_{args.results_folder}.csv", index=False)


if __name__ == "__main__":
    main()


# powergenome/fuels.py

"""
Load fuel prices needed for the model
"""

from asyncio.log import logger
import pandas as pd

from powergenome.eia_opendata import add_user_fuel_prices


def fuel_cost_table(fuel_costs, generators, settings):
    all_fuel_costs = add_user_fuel_prices(settings, fuel_costs)
    unique_fuels = generators["Fuel"].drop_duplicates()
    model_year_costs = all_fuel_costs.loc[
        all_fuel_costs["year"] == settings["model_year"], :
    ]
    fuel_df = pd.DataFrame(unique_fuels)

    fuel_price_map = {
        row.full_fuel_name: row.price
        for row in model_year_costs.itertuples(index=False, name="row")
    }

    emission_dict = settings.get("fuel_emission_factors", {}) or {}
    user_fuels = set(all_fuel_costs["fuel"]) - set(fuel_costs["fuel"])
    for u_f in user_fuels:
        if u_f not in emission_dict.keys():
            logger.warning(
                "\n\n**********************\n"
                f"The user fuel {u_f} does not have an emissions factor specified in "
                "the settings parameter 'fuel_emission_factors'. This is fine if the "
                "emission factor should be 0, otherwise be sure to add a value.\n"
            )
    fuel_emission_map = {}
    for full_fuel_name in fuel_price_map:
        if (
            full_fuel_name.split("_")[-1]
            in (settings.get("aeo_fuel_scenarios", {}) or {}).keys()
        ):
            base_fuel_name = full_fuel_name.split("_")[-1]
        elif (
            full_fuel_name.split("_")[-1]
            in (settings.get("user_fuel_price", {}) or {}).keys()
        ):
            base_fuel_name = full_fuel_name.split("_")[-1]
        else:
            base_fuel_name = full_fuel_name
        if base_fuel_name in emission_dict:
            fuel_emission_map[full_fuel_name] = emission_dict[base_fuel_name]
        else:
            fuel_emission_map[full_fuel_name] = 0

    ccs_fuels = (settings.get("ccs_fuel_map", {}) or {}).values()
    for ccs_fuel in ccs_fuels:
        fuels = generators.loc[
            generators["Fuel"].str.contains(ccs_fuel), "Fuel"
        ].unique()
        for f in fuels:
            # keep the non-ccs price
            base_name = ("_").join(f.split("_")[:-1])
            fuel_price_map[f] = fuel_price_map[base_name]
            fuel_emission_map[f] = fuel_emission_map[base_name]

    fuel_df["Cost_per_MMBtu"] = fuel_df["Fuel"].map(fuel_price_map)
    fuel_df["CO2_content_tons_per_MMBtu"] = fuel_df["Fuel"].map(fuel_emission_map)

    # Slow to loop through all of the rows this way but the df shouldn't be too long
    fuel_df = fuel_df.apply(adjust_ccs_fuels, axis=1, settings=settings)
    fuel_df = add_carbon_tax(fuel_df, settings)
    fuel_df["Cost_per_MMBtu"] = fuel_df["Cost_per_MMBtu"]
    fuel_df["CO2_content_tons_per_MMBtu"] = fuel_df["CO2_content_tons_per_MMBtu"]
    fuel_df.fillna(0, inplace=True)

    if settings.get("reduce_time_domain"):
        days = settings["time_domain_days_per_period"]
        time_periods = settings["time_domain_periods"]
        num_hours = days * time_periods * 24
    else:
        num_hours = 8760

    fuel_df_prices = pd.DataFrame(
        [fuel_df["Cost_per_MMBtu"]], index=range(1, num_hours + 1)
    )
    fuel_df_prices = fuel_df_prices.round(2)
    fuel_df_prices.columns = unique_fuels

    fuel_df_top = pd.DataFrame([fuel_df["CO2_content_tons_per_MMBtu"]])
    fuel_df_top = fuel_df_top.round(5)
    fuel_df_top.columns = unique_fuels
    fuel_df_top.index = [0]

    fuel_frames = [fuel_df_top, fuel_df_prices]
    fuel_df_new = pd.concat(fuel_frames)
    fuel_df_new.index.name = "Time_Index"
    return fuel_df_new


# def modify_fuel_new_genx():


def adjust_ccs_fuels(ccs_fuel_row, settings):
    base_fuel_name = None
    for ccs_fuel in (settings.get("ccs_fuel_map", {}) or {}).values():
        if ccs_fuel in ccs_fuel_row["Fuel"]:
            base_fuel_name = ccs_fuel
    if base_fuel_name:
        # USD/tonne disposal
        disposal_cost = settings.get("ccs_disposal_cost")
        if not disposal_cost:
            logger.warning(
                "You did not specify a CCS disposal cost, so it will be set to $0. "
                "Set a non-zero value with the settings parameter 'ccs_disposal_cost'."
            )
        disposal_cost = 0

        capture_rate = settings["ccs_capture_rate"][base_fuel_name]

        co2_captured = ccs_fuel_row["CO2_content_tons_per_MMBtu"] * capture_rate

        ccs_fuel_row["CO2_content_tons_per_MMBtu"] -= co2_captured
        ccs_fuel_row["Cost_per_MMBtu"] += co2_captured * disposal_cost

    else:
        pass

    return ccs_fuel_row


def add_carbon_tax(fuel_df, settings):
    ctax = settings.get("carbon_tax") or 0

    fuel_df.loc[:, "Cost_per_MMBtu"] = fuel_df.loc[:, "Cost_per_MMBtu"] + (
        fuel_df.loc[:, "CO2_content_tons_per_MMBtu"] * ctax
    )

    return fuel_df


# powergenome/generators.py

import collections
import logging
from numbers import Number
from typing import Dict, List, Union
import re
from zipfile import BadZipFile

import requests

import geopandas as gpd
import numpy as np
import pandas as pd
from pathlib import Path
import pudl
from bs4 import BeautifulSoup
from flatten_dict import flatten
import sqlalchemy
from powergenome.cluster_method import (
    cluster_by_owner,
    cluster_kmeans,
    weighted_ownership_by_unit,
)
from powergenome.eia_opendata import fetch_fuel_prices, modify_fuel_prices
from powergenome.external_data import (
    make_demand_response_profiles,
    demand_response_resource_capacity,
    add_resource_max_cap_spur,
)
from powergenome.load_profiles import make_distributed_gen_profiles
from powergenome.nrelatb import (
    atb_fixed_var_om_existing,
    atb_new_generators,
    fetch_atb_costs,
    fetch_atb_heat_rates,
    fetch_atb_offshore_spur_costs,
    investment_cost_calculator,
)
from powergenome.params import (
    DATA_PATHS,
    IPM_GEOJSON_PATH,
    build_resource_clusters,
)
from powergenome.price_adjustment import inflation_price_adjustment
from powergenome.resource_clusters import map_eia_technology
from powergenome.util import (
    download_save,
    find_region_col,
    map_agg_region_names,
    reverse_dict_of_lists,
    snake_case_col,
    regions_to_keep,
    snake_case_str,
    load_ipm_shapefile,
    remove_leading_zero,
)
from powergenome.GenX import rename_gen_cols
from scipy.stats import iqr
from sklearn import cluster, preprocessing
from xlrd import XLRDError

logger = logging.getLogger(__name__)


planned_col_map = {
    "Entity ID": "utility_id_eia",
    "Entity Name": "utility_name",
    "Plant ID": "plant_id_eia",
    "Plant Name": "plant_name",
    "Sector": "sector_name",
    "Plant State": "state",
    "Generator ID": "generator_id",
    "Unit Code": "unit_code",
    "Nameplate Capacity (MW)": "capacity_mw",
    "Net Summer Capacity (MW)": "summer_capacity_mw",
    "Net Winter Capacity (MW)": "winter_capacity_mw",
    "Technology": "technology_description",
    "Energy Source Code": "energy_source_code_1",
    "Prime Mover Code": "prime_mover_code",
    "Planned Operation Month": "planned_operating_month",
    "Planned Operation Year": "planned_operating_year",
    "Status": "operational_status",
    "Nameplate Energy Capacity (MWh)": "capacity_mwh",
    "DC Net Capacity (MW)": "dc_net_capacity_mw",
    "County": "county",
    "Latitude": "latitude",
    "Longitude": "longitude",
    "Google Map": "google_map",
    "Bing Map": "bing_map",
    "Balancing Authority Code": "balancing_authority_code",
}

op_status_map = {
    "(V) Under construction, more than 50 percent complete": "V",
    "(TS) Construction complete, but not yet in commercial operation": "TS",
    "(U) Under construction, less than or equal to 50 percent complete": "U",
    "(T) Regulatory approvals received. Not under construction": "T",
    "(P) Planned for installation, but regulatory approvals not initiated": "P",
    "(L) Regulatory approvals pending. Not under construction": "L",
    "(OT) Other": "OT",
    "(SB) Standby/Backup: available for service but not normally used": "SB",
    "(OP) Operating": "OP",
    "(OA) Out of service but expected to return to service in next calendar year": "OA",
    "(OS) Out of service and NOT expected to return to service in next calendar year": "OS",
}

TRANSMISSION_TYPES = ["spur", "offshore_spur", "tx"]


def fill_missing_tech_descriptions(
    df: pd.DataFrame, date_col: str = "report_date"
) -> pd.DataFrame:
    """
    EIA 860 records before 2014 don't have a technology description. If we want to
    include any of this data in the historical record (e.g. heat rates or capacity
    factors) then they need to be filled in.

    Parameters
    ----------
    df : dataframe
        A pandas dataframe with columns plant_id_eia, generator_id, and
        technology_description.
    date_col: str
        The column with date information, used to sort values from oldest to newest.
        Assumes that newer records will have a valid technology description for the
        generator.

    Returns
    -------
    dataframe
        Same data that came in, but with missing technology_description values filled
        in.
    """
    if (
        date_col not in df.columns
        and not df.loc[df["technology_description"].isnull(), :].empty
    ):
        logger.warning(
            "A dataframe with missing technology descriptions does not have the date column "
            f"{date_col}. The rows with missing technology descriptions look like:\n\n"
            f"{df.loc[df['technology_description'].isnull(), :]}\n\n"
        )
    start_len = len(df)
    df = df.sort_values(by=date_col)
    df_list = []
    missing_tech_plants = df.loc[df["technology_description"].isna(), :]
    if missing_tech_plants.empty:
        return df

    df = df.drop(index=missing_tech_plants.index)
    for _, _df in missing_tech_plants.groupby(
        ["plant_id_eia", "generator_id"], as_index=False
    ):
        _df["technology_description"].fillna(method="bfill", inplace=True)
        df_list.append(_df)
    results = pd.concat([df, pd.concat(df_list, ignore_index=True, sort=False)])

    if df.loc[df["technology_description"].isnull(), :].empty is False:
        logger.warning("Failed to fill some technology names.")

    end_len = len(results)
    assert (
        start_len == end_len
    ), "Somehow records were dropped when filling tech_descriptions"
    return results


def group_generators_at_plant(df, by=["plant_id_eia"], agg_fn={"capacity_mw": "sum"}):
    """
    Group generators at a plant. This is a flexible function that lets a user group
    by the desired attributes (e.g. plant id) and perform aggregated operations on each
    group.

    This function also might be a bit unnecessary given how simple it is.

    Parameters
    ----------
    df : dataframe
        Pandas dataframe with information on power plants.
    by : list, optional
        Columns to use for the groupby, by default ["plant_id_eia"]
    agg_fn : dict, optional
        Aggregation function to pass to groupby, by default {"capacity_mw": "sum"}

    Returns
    -------
    dataframe
        The grouped dataframe with aggregation functions applied.
    """

    df_grouped = df.groupby(by, as_index=False).agg(agg_fn)

    return df_grouped


def startup_fuel(df: pd.DataFrame, settings: dict) -> pd.DataFrame:
    """Add startup fuel consumption for generators

    Parameters
    ----------
    df : DataFrame
        All generator clusters. Must have a column "technology". Can include both EIA
        and NRELATB technology names.
    settings : dictionary
        User-defined settings loaded from a YAML file. Keys in "startup_fuel_use"
        must match those in "eia_atb_tech_map".

    Returns
    -------
    DataFrame
        Modified dataframe with the new column "Start_Fuel_MMBTU_per_MW".
    """
    df["Start_Fuel_MMBTU_per_MW"] = 0
    for eia_tech, fuel_use in (settings.get("startup_fuel_use") or {}).items():
        if not isinstance(settings["eia_atb_tech_map"][eia_tech], list):
            settings["eia_atb_tech_map"][eia_tech] = [
                settings["eia_atb_tech_map"][eia_tech]
            ]

        atb_tech = settings["eia_atb_tech_map"][eia_tech]
        atb_tech.append(eia_tech)
        for tech in atb_tech:
            df.loc[df["technology"] == tech, "Start_Fuel_MMBTU_per_MW"] = fuel_use
            df.loc[
                df["technology"].str.contains(tech, case=False),
                "Start_Fuel_MMBTU_per_MW",
            ] = fuel_use

    return df


def startup_nonfuel_costs(df: pd.DataFrame, settings: dict) -> pd.DataFrame:
    """Add inflation adjusted startup nonfuel costs per MW for generators

    Parameters
    ----------
    df : DataFrame
        Must contain a column "technology" with the names of each technology type.
    settings : dict
        Dictionary based on YAML settings file. Must contain the keys
        "startup_costs_type", "startup_vom_costs_mw", "existing_startup_costs_tech_map",
        etc.

    Returns
    -------
    DataFrame
        Modified df with new column "Start_Cost_per_MW"
    """
    logger.info("Adding non-fuel startup costs")
    target_usd_year = settings.get("target_usd_year")

    vom_costs = settings.get("startup_vom_costs_mw", {})
    vom_usd_year = settings.get("startup_vom_costs_usd_year")

    if target_usd_year and vom_usd_year:
        logger.info(
            f"Changing non-fuel VOM costs from {vom_usd_year} to " f"{target_usd_year}"
        )
        for key, cost in vom_costs.items():
            vom_costs[key] = inflation_price_adjustment(
                price=cost, base_year=vom_usd_year, target_year=target_usd_year
            )

    startup_type = settings.get("startup_costs_type")
    startup_costs = settings.get(startup_type, {})
    startup_costs_usd_year = settings.get("startup_costs_per_cold_start_usd_year")

    if target_usd_year and startup_costs_usd_year:
        logger.info(
            f"Changing non-fuel startup costs from {vom_usd_year} to {target_usd_year}"
        )
        for key, cost in startup_costs.items():
            startup_costs[key] = inflation_price_adjustment(
                price=cost,
                base_year=startup_costs_usd_year,
                target_year=target_usd_year,
            )

    df["Start_Cost_per_MW"] = 0

    for existing_tech, cost_tech in settings.get(
        "existing_startup_costs_tech_map", {}
    ).items():
        total_startup_costs = vom_costs[cost_tech] + startup_costs[cost_tech]
        df.loc[
            df["technology"].str.contains(existing_tech, case=False),
            "Start_Cost_per_MW",
        ] = total_startup_costs

    for new_tech, cost_tech in settings.get("new_build_startup_costs", {}).items():
        total_startup_costs = vom_costs[cost_tech] + startup_costs[cost_tech]
        df.loc[
            df["technology"].str.contains(new_tech), "Start_Cost_per_MW"
        ] = total_startup_costs
    df.loc[:, "Start_Cost_per_MW"] = df.loc[:, "Start_Cost_per_MW"]

    # df.loc[df["technology"].str.contains("Nuclear"), "Start_Cost_per_MW"] = "FILL VALUE"

    return df


def group_technologies(
    df: pd.DataFrame,
    group_technologies: bool = False,
    tech_groups: Dict[str, list] = {},
    regional_no_grouping: Dict[str, list] = {},
) -> pd.DataFrame:
    """
    Group different technologies together based on parameters in the settings file.
    An example would be to put a bunch of different technologies under the umbrella
    category of "biomass" or "peaker".

    Parameters
    ----------
    df : dataframe
        Pandas dataframe with
    settings : dictionary
        User-defined settings loaded from a YAML file. Must have key tech_groups.

    Returns
    -------
    dataframe
        Same as incoming dataframe but with grouped technology types
    """
    if not group_technologies:
        return df
    else:
        df["_technology"] = df["technology_description"]
        for tech, group in tech_groups.items():
            df.loc[df["technology_description"].isin(group), "_technology"] = tech

        for region, tech_list in regional_no_grouping.items():
            df.loc[
                (df["model_region"] == region)
                & (df["technology_description"].isin(tech_list)),
                "_technology",
            ] = df.loc[
                (df["model_region"] == region)
                & (df["technology_description"].isin(tech_list)),
                "technology_description",
            ]

        df.loc[:, "technology_description"] = df.loc[:, "_technology"]
        df = df.drop(columns=["_technology"])

        return df


def label_hydro_region(gens_860, pudl_engine, model_regions_gdf):
    """
    Label hydro facilities that don't have a region by default.

    Parameters
    ----------
    gens_860 : dataframe
        Infomation on all generators from PUDL
    pudl_engine : sqlalchemy.Engine
        A sqlalchemy connection for use by pandas
    model_regions_gdf : dataframe
        Geodataframe of the model regions

    Returns
    -------
    dataframe
        Plant id and region for any hydro that didn't originally have a region label.
    """

    plant_entity = pd.read_sql_table("plants_entity_eia", pudl_engine)

    model_hydro = gens_860.loc[
        gens_860["technology_description"] == "Conventional Hydroelectric"
    ].merge(plant_entity[["plant_id_eia", "latitude", "longitude"]], on="plant_id_eia")

    no_lat_lon = model_hydro.loc[
        (model_hydro["latitude"].isnull()) | (model_hydro["longitude"].isnull()), :
    ]
    if not no_lat_lon.empty:
        print(no_lat_lon["summer_capacity_mw"].sum(), " MW without lat/lon")
    model_hydro = model_hydro.dropna(subset=["latitude", "longitude"])

    # Convert the lon/lat values to geo points. Need to add an initial CRS and then
    # change it to align with the IPM regions
    model_hydro_gdf = gpd.GeoDataFrame(
        model_hydro,
        geometry=gpd.points_from_xy(model_hydro.longitude, model_hydro.latitude),
        crs="EPSG:4326",
    )

    if model_hydro_gdf.crs != model_regions_gdf.crs:
        model_hydro_gdf = model_hydro_gdf.to_crs(model_regions_gdf.crs)

    model_hydro_gdf = gpd.sjoin(model_regions_gdf, model_hydro_gdf)

    keep_cols = ["plant_id_eia", "region"]
    return model_hydro_gdf.loc[:, keep_cols]


def load_plant_region_map(
    gens_860,
    pudl_engine,
    pg_engine,
    settings,
    model_regions_gdf,
    table="plant_region_map_epaipm",
):
    """
    Load the region that each plant is located in.

    Parameters
    ----------
    pudl_engine : sqlalchemy.Engine
        A sqlalchemy connection for use by pandas
    settings : dictionary
        The dictionary of settings with a dictionary of region aggregations
    table : str, optional
        The SQL table to load, by default "plant_region_map_epaipm"

    Returns
    -------
    dataframe
        A dataframe where each plant has an associated "model_region" mapped
        from the original region labels.
    """
    # Load dataframe of region labels for each EIA plant id
    region_map_df = pd.read_sql_table(table, con=pg_engine)

    if settings.get("plant_region_map_fn"):
        user_region_map_df = pd.read_csv(
            Path(settings["input_folder"]) / settings["plant_region_map_fn"]
        )
        assert (
            "region" in user_region_map_df.columns
        ), f"The column 'region' must appear in {settings['plant_region_map_fn']}"
        assert (
            "plant_id_eia" in user_region_map_df.columns
        ), f"The column 'plant_id_eia' must appear in {settings['plant_region_map_fn']}"

        user_region_map_df = user_region_map_df.set_index("plant_id_eia")

        region_map_df.loc[
            region_map_df["plant_id_eia"].isin(user_region_map_df.index), "region"
        ] = region_map_df["plant_id_eia"].map(user_region_map_df["region"])

    # Label hydro using the IPM shapefile because NEEDS seems to drop some hydro
    all_hydro_regions = label_hydro_region(gens_860, pudl_engine, model_regions_gdf)

    region_map_df = pd.concat(
        [region_map_df, all_hydro_regions], ignore_index=True, sort=False
    ).drop_duplicates(subset=["plant_id_eia"], keep="first")

    # Settings has a dictionary of lists for regional aggregations. Need
    # to reverse this to use in a map method.
    keep_regions, region_agg_map = regions_to_keep(
        settings["model_regions"], settings.get("region_aggregations")
    )

    # Create a new column "model_region" with labels that we're using for aggregated
    # regions

    model_region_map_df = region_map_df.loc[
        region_map_df.region.isin(keep_regions), :
    ].drop(columns="id")

    model_region_map_df = map_agg_region_names(
        df=model_region_map_df,
        region_agg_map=region_agg_map,
        original_col_name="region",
        new_col_name="model_region",
    )

    # There are some cases of plants with generators assigned to different IPM regions.
    # If regions are aggregated there may be some duplicates in the results.
    model_region_map_df = model_region_map_df.drop_duplicates(
        subset=["plant_id_eia", "model_region"]
    )

    return model_region_map_df


def label_retirement_year(
    df,
    settings,
    age_col="operating_date",
    settings_retirement_table="retirement_ages",
    add_additional_retirements=True,
):
    """
    Add a retirement year column to the dataframe based on the year each generator
    started operating.

    Parameters
    ----------
    df : dataframe
        Dataframe of generators
    settings : dictionary
        The dictionary of settings with a dictionary of generator lifetimes
    age_col : str, optional
        The dataframe column to use when calculating the retirement year, by default
        "operating_date"
    settings_retirement_table : str, optional
        The settings dictionary key for another dictionary of generator retirement
        lifetimes, by default "retirement_ages"
    add_additional_retirements : bool, optional
        Logic to determine if additional retirements from the settings file should
        be checked. For example, this isn't necessary when adding proposed generators
        because we probably won't be setting an artifically early retirement year.
    """

    start_len = len(df)
    retirement_ages = settings.get(settings_retirement_table, {}) or {}
    if "retirement_year" not in df.columns:
        df["retirement_year"] = np.nan

    for tech, life in retirement_ages.items():
        try:
            df.loc[df.technology_description == tech, "retirement_year"] = (
                df.loc[df.technology_description == tech, age_col].dt.year + life
            )
        except AttributeError:
            # This is a bit hacky but for the proposed plants I have an int column
            df.loc[df.technology_description == tech, "retirement_year"] = (
                df.loc[df.technology_description == tech, age_col] + life
            )

    try:
        df.loc[~df["planned_retirement_date"].isnull(), "retirement_year"] = df.loc[
            ~df["planned_retirement_date"].isnull(), "planned_retirement_date"
        ].dt.year
    except KeyError:
        pass

    # Add additonal retirements from settings file
    if settings.get("additional_retirements") and add_additional_retirements:
        logger.info("Changing retirement dates based on settings file")
        model_year = settings["model_year"]
        start_ret_cap = df.loc[
            df["retirement_year"] <= model_year, settings["capacity_col"]
        ].sum()
        logger.info(f"Starting retirement capacity is {start_ret_cap} MW")
        i = 0
        ret_cap = 0
        for record in settings["additional_retirements"]:
            plant_id, gen_id, ret_year = record
            # gen ids are strings, not integers
            gen_id = str(gen_id)

            df.loc[
                (df["plant_id_eia"] == plant_id) & (df["generator_id"] == gen_id),
                "retirement_year",
            ] = ret_year

            i += 1
            ret_cap += df.loc[
                (df["plant_id_eia"] == plant_id) & (df["generator_id"] == gen_id),
                settings["capacity_col"],
            ].sum()

        end_ret_cap = df.loc[
            df["retirement_year"] <= model_year, settings["capacity_col"]
        ].sum()
        logger.info(f"Ending retirement capacity is {end_ret_cap} MW")
        if not end_ret_cap > start_ret_cap:
            logger.debug(
                "Adding retirements from settings didn't change the retiring capacity."
            )
        if end_ret_cap - start_ret_cap != ret_cap:
            logger.debug(
                f"Retirement diff is {end_ret_cap - start_ret_cap}, adding retirements "
                f"yields {ret_cap} MW"
            )
        logger.info(
            f"The retirement year for {i} plants, totaling {ret_cap} MW, was changed "
            "based on settings file parameters"
        )
    else:
        logger.info("No retirement dates changed based on the settings file")

    end_len = len(df)

    assert start_len == end_len

    return df


def label_small_hydro(df, settings, by=["plant_id_eia"]):
    """
    Use rules from the settings file to label plants below a certain size as small
    hydroelectric rather than conventional hydroelectric.

    Parameters
    ----------
    df : dataframe
        EIA 860 data on generators
    settings : dict
        User-defined parameters from a settings file
    by : list, optional
        What columns to use in the groupby function when summing capacity, by default
        ["plant_id_eia"]

    Returns
    -------
    dataframe
        If the user wants to label small hydro plants, some of the conventional
        hydro facilities will have their technology type changed to small hydro.
    """
    if not settings.get("small_hydro"):
        return df
    if "report_date" not in by and "report_date" in df.columns:
        # by.append("report_date")
        logger.warning("'report_date' is in the df but not used in the groupby")
    region_agg_map = reverse_dict_of_lists(settings.get("region_aggregations", {}))
    keep_regions = [
        x
        for x in settings["model_regions"] + list(region_agg_map)
        if x in settings["small_hydro_regions"]
    ]
    start_len = len(df)
    size_cap = settings["small_hydro_mw"]
    cap_col = settings.get("capacity_col")
    if not cap_col in df:
        cap_col = "capacity_mw"

    start_hydro_capacity = df.query(
        "technology_description=='Conventional Hydroelectric'"
    )[cap_col].sum()

    plant_capacity = (
        df.loc[
            (df["technology_description"] == "Conventional Hydroelectric")
            & (df["model_region"].isin(keep_regions))
        ]
        .groupby(by, as_index=False)[cap_col]
        .sum()
    )

    small_hydro_plants = plant_capacity.loc[
        plant_capacity[cap_col] <= size_cap, "plant_id_eia"
    ]

    df.loc[
        (df["technology_description"] == "Conventional Hydroelectric")
        & (df["plant_id_eia"].isin(small_hydro_plants)),
        "technology_description",
    ] = "Small Hydroelectric"

    end_len = len(df)
    small_hydro_capacity = df.query("technology_description=='Small Hydroelectric'")[
        cap_col
    ].sum()
    end_conv_hydro_capacity = df.query(
        "technology_description=='Conventional Hydroelectric'"
    )[cap_col].sum()

    assert start_len == end_len
    assert np.allclose(
        start_hydro_capacity, small_hydro_capacity + end_conv_hydro_capacity
    )

    return df


def load_generator_860_data(pudl_engine, data_years=[2017]):
    """
    Load EIA 860 generator data from the PUDL database

    Parameters
    ----------
    pudl_engine : sqlalchemy.Engine
        A sqlalchemy connection for use by pandas
    data_years : list, optional
        Years of data to load, by default [2017]

    Returns
    -------
    dataframe
        All of the generating units from PUDL
    """
    data_years = [str(y) for y in data_years]
    sql = f"""
        SELECT * FROM generators_eia860
        WHERE operational_status_code NOT IN ('RE', 'OS', 'IP', 'CN')
        AND strftime('%Y',report_date) in ({','.join(['?']*len(data_years))})
    """
    gens_860 = pd.read_sql_query(
        sql=sql,
        con=pudl_engine,
        params=data_years,
        parse_dates=["report_date", "planned_retirement_date"],
    )

    return gens_860


def supplement_generator_860_data(
    gens_860: pd.DataFrame,
    gens_entity: pd.DataFrame,
    bga: pd.DataFrame,
    model_region_map: pd.DataFrame,
    settings: dict,
):
    """
    Load data about each generating unit in the model area.

    Parameters
    ----------
    gens_860 : dataframe
        Information on all generating units for the given data years.
    pudl_engine : sqlalchemy.Engine
        A sqlalchemy connection for use by pandas
    settings : dictionary
        The dictionary of settings with a dictionary of region aggregations
    pudl_out : pudl.PudlTabl
        A PudlTabl object for loading pre-calculated PUDL analysis data
    model_region_map : dataframe
        A dataframe with columns 'plant_id_eia' and 'model_region' (aggregated regions)
    data_years : list, optional
        Years of data to include, by default [2017]

    Returns
    -------
    dataframe
        Data about each generator and generation unit that will be included in the
        model. Columns include:

        ['plant_id_eia', 'generator_id',
       'capacity_mw', 'energy_source_code_1',
       'energy_source_code_2', 'minimum_load_mw', 'operational_status_code',
       'planned_new_capacity_mw', 'switch_oil_gas', 'technology_description',
       'time_cold_shutdown_full_load_code', 'model_region', 'prime_mover_code',
       'operating_date', 'boiler_id', 'unit_id_eia', 'unit_id_pudl', 'unit_id_pg,
       'retirement_year']
    """

    initial_capacity = (
        gens_860.loc[gens_860["plant_id_eia"].isin(model_region_map["plant_id_eia"])]
        .groupby("technology_description")[settings["capacity_col"]]
        .sum()
    )

    # Add pudl unit ids, only include specified data years

    # Combine generator data that can change over time with static entity data
    # and only keep generators that are in a region of interest

    gen_cols = set(
        [
            # "report_date",
            "plant_id_eia",
            # "plant_name",
            "generator_id",
            # "balancing_authority_code",
            settings["capacity_col"],
            "capacity_mw",
            "energy_source_code_1",
            "energy_source_code_2",
            "minimum_load_mw",
            "operational_status_code",
            "planned_new_capacity_mw",
            "switch_oil_gas",
            "technology_description",
            "time_cold_shutdown_full_load_code",
            "planned_retirement_date",
        ]
    )

    entity_cols = ["plant_id_eia", "generator_id", "prime_mover_code", "operating_date"]

    bga_cols = [
        "plant_id_eia",
        "generator_id",
        "boiler_id",
        "unit_id_eia",
        "unit_id_pudl",
    ]

    # In this merge of the three dataframes we're trying to label each generator with
    # the model region it is part of, the prime mover and operating date, and the
    # PUDL unit codes (where they exist).
    gens_860_model = (
        pd.merge(
            gens_860[gen_cols],
            model_region_map.drop(columns="region"),
            on="plant_id_eia",
            how="inner",
        )
        .merge(
            gens_entity[entity_cols], on=["plant_id_eia", "generator_id"], how="inner"
        )
        .merge(bga[bga_cols], on=["plant_id_eia", "generator_id"], how="left")
    )
    gens_860_model["unit_id_pg"] = gens_860_model.loc[:, "unit_id_pudl"]
    gens_860_model.loc[gens_860_model.unit_id_pg.isnull(), "unit_id_pg"] = (
        gens_860_model.loc[gens_860_model.unit_id_pg.isnull(), "plant_id_eia"].astype(
            str
        )
        + "_"
        + gens_860_model.loc[gens_860_model.unit_id_pg.isnull(), "generator_id"].astype(
            str
        )
    ).to_numpy()

    # Where summer/winter capacity values are missing set equal to nameplate capacity,
    # but only if all generators within a unit are missing the capacity value
    check_units = gens_860_model.loc[
        gens_860_model[settings["capacity_col"]].isna()
    ].groupby(["plant_id_eia", "unit_id_pg"])
    for (plant_id, unit_id), _df in check_units:
        if _df[settings["capacity_col"]].isna().all():
            gens_860_model.loc[
                (gens_860_model["plant_id_eia"] == plant_id)
                & (gens_860_model["unit_id_pg"] == unit_id),
                settings["capacity_col"],
            ] = gens_860_model.loc[
                (gens_860_model["plant_id_eia"] == plant_id)
                & (gens_860_model["unit_id_pg"] == unit_id),
                "capacity_mw",
            ]

    merged_capacity = gens_860_model.groupby("technology_description")[
        settings["capacity_col"]
    ].sum()
    if not np.allclose(initial_capacity.sum(), merged_capacity.sum()):
        for i_idx, i_row in initial_capacity.iteritems():
            if not np.allclose(i_row, merged_capacity[i_idx]):
                logger.warning(
                    "********************************\n"
                    "When adding plant entity/boiler info to generators and filling missing"
                    " seasonal capacity values, technology"
                    f"{i_idx} changed capacity from {i_row} to {merged_capacity[i_idx]}"
                    "\n********************************"
                )

    return gens_860_model


def create_plant_gen_id(df):
    """Combine the plant id and generator id to form a unique combination

    Parameters
    ----------
    df : dataframe
        Must contain columns plant_id_eia and generator_id

    Returns
    -------
    dataframe
        Same as input but with the additional column plant_gen_id
    """

    df["plant_gen_id"] = (
        df["plant_id_eia"].astype(str) + "_" + df["generator_id"].astype(str)
    )

    return df


def remove_canceled_860m(df, canceled_860m):
    """Remove generators that 860m shows as having been canceled

    Parameters
    ----------
    df : dataframe
        All of the EIA 860 generators
    canceled_860m : dataframe
        From the 860m Canceled or Postponed sheet

    Returns
    -------
    dataframe
        Same as input, but possibly without generators that were proposed
    """
    df = create_plant_gen_id(df)
    canceled_860m = create_plant_gen_id(canceled_860m)

    canceled = df.loc[df["plant_gen_id"].isin(canceled_860m["plant_gen_id"]), :]

    not_canceled_df = df.loc[~df["plant_gen_id"].isin(canceled_860m["plant_gen_id"]), :]

    not_canceled_df = not_canceled_df.drop(columns="plant_gen_id")

    if not canceled.empty:
        assert len(df) == len(canceled) + len(not_canceled_df)

    return not_canceled_df.reset_index(drop=True)


def remove_retired_860m(df, retired_860m):
    """Remove generators that 860m shows as having been retired

    Parameters
    ----------
    df : dataframe
        All of the EIA 860 generators
    retired_860m : dataframe
        From the 860m Retired sheet

    Returns
    -------
    dataframe
        Same as input, but possibly without generators that have retired
    """

    df = create_plant_gen_id(df)
    retired_860m = create_plant_gen_id(retired_860m)

    retired = df.loc[df["plant_gen_id"].isin(retired_860m["plant_gen_id"]), :]

    not_retired_df = df.loc[~df["plant_gen_id"].isin(retired_860m["plant_gen_id"]), :]

    not_retired_df = not_retired_df.drop(columns="plant_gen_id")

    if not retired.empty:
        assert len(df) == len(retired) + len(not_retired_df)

    return not_retired_df.reset_index(drop=True)


def remove_future_retirements_860m(df, retired_860m):
    """Remove generators that 860m shows as having been retired

    Parameters
    ----------
    df : dataframe
        All of the EIA 860 generators
    retired_860m : dataframe
        From the 860m Retired sheet

    Returns
    -------
    dataframe
        Same as input, but possibly without generators that have retired
    """

    df = create_plant_gen_id(df)
    retired_860m = create_plant_gen_id(retired_860m)

    retired = df.loc[df["plant_gen_id"].isin(retired_860m["plant_gen_id"]), :]

    not_retired_df = df.loc[~df["plant_gen_id"].isin(retired_860m["plant_gen_id"]), :]

    not_retired_df = not_retired_df.drop(columns="plant_gen_id")

    if not retired.empty:
        assert len(df) == len(retired) + len(not_retired_df)

    return not_retired_df


def update_operating_date_860m(
    df: pd.DataFrame, operating_860m: pd.DataFrame
) -> pd.DataFrame:
    """Update the operating date of EIA generators using data from 860m.

    When the "operating_date" of a generator is nan, fill with operating year data
    from 860m.

    Parameters
    ----------
    df : pd.DataFrame
        Data on existing EIA generating units. Must have columns "plant_id_eia",
        "generator_id", and "operating_date".
    operating_860m : pd.DataFrame
        A dataframe of operating generating units from EIA 860m. Must have columns
        "plant_id_eia", "generator_id", and "operating_year".

    Returns
    -------
    pd.DataFrame
        The original "df" dataframe with missing operating dates filled using the operating
        year from 860m.
    """
    _df = df.set_index(["plant_id_eia", "generator_id"])
    _operating_860m = operating_860m.set_index(["plant_id_eia", "generator_id"])
    no_op_date = _df.loc[_df["operating_date"].isna(), :]
    no_op_date = pd.merge(
        no_op_date,
        _operating_860m["operating_year"],
        left_index=True,
        right_index=True,
        how="left",
        validate="1:1",
    )

    _df.loc[no_op_date.index, "operating_date"] = pd.to_datetime(
        no_op_date.loc[no_op_date["operating_year"].notna(), "operating_year"],
        format="%Y",
    )

    return _df.reset_index()


def load_923_gen_fuel_data(pudl_engine, pudl_out, model_region_map, data_years=[2017]):
    """
    Load generation and fuel data for each plant. EIA-923 provides these values for
    each prime mover/fuel combination at every generator. This data can be used to
    calculate the heat rate of generators at a single plant. Generators sharing a prime
    mover (e.g. multiple combustion turbines) will end up sharing the same heat rate.

    Parameters
    ----------
    pudl_engine : sqlalchemy.Engine
        A sqlalchemy connection for use by pandas
    pudl_out : pudl.PudlTabl
        A PudlTabl object for loading pre-calculated PUDL analysis data
    model_region_map : dataframe
        A dataframe with columns 'plant_id_eia' and 'model_region' (aggregated regions)
    data_years : list, optional
        Years of data to include, by default [2017]

    Returns
    -------
    dataframe
        Generation, fuel use, and heat rates of prime mover/fuel combos over all data
        years. Columns are:

        ['plant_id_eia', 'fuel_type', 'fuel_type_code_pudl',
       'fuel_type_code_aer', 'prime_mover_code', 'fuel_consumed_units',
       'fuel_consumed_for_electricity_units', 'fuel_consumed_mmbtu',
       'fuel_consumed_for_electricity_mmbtu', 'net_generation_mwh',
       'heat_rate_mmbtu_mwh']
    """
    if isinstance(data_years, (int, float)):
        data_years = [str(data_years)]
    data_years = [str(y) for y in data_years]

    # Load 923 generation and fuel data for one or more years.
    # Only load plants in the model regions.
    sql = f"""
        SELECT * FROM generation_fuel_eia923
        WHERE strftime('%Y',report_date) in ({','.join(['?']*len(data_years))})
    """
    gen_fuel_923 = pd.read_sql_query(
        sql, pudl_engine, params=data_years, parse_dates=["report_date"]
    )
    gen_fuel_923 = gen_fuel_923.loc[
        gen_fuel_923["plant_id_eia"].isin(model_region_map.plant_id_eia),
        :,
    ]

    insp = sqlalchemy.inspect(pudl_engine)
    if insp.has_table("generation_fuel_nuclear_eia923"):
        sql = f"""
            SELECT * FROM generation_fuel_nuclear_eia923
            WHERE strftime('%Y',report_date) in ({','.join(['?']*len(data_years))})
        """
        gen_fuel_nuclear_923 = pd.read_sql_query(
            sql, pudl_engine, params=data_years, parse_dates=["report_date"]
        )
        gen_fuel_nuclear_923 = gen_fuel_nuclear_923.loc[
            gen_fuel_nuclear_923["plant_id_eia"].isin(model_region_map.plant_id_eia),
            :,
        ]

        gen_fuel_923 = pd.concat(
            [gen_fuel_923, gen_fuel_nuclear_923], ignore_index=True
        )

    return gen_fuel_923


def modify_cc_prime_mover_code(df, gens_860):
    """Change combined cycle prime movers from CA and CT to CC.

    The heat rate of combined cycle plants that aren't included in PUDL heat rate by
    unit should probably be done with the combustion and steam turbines combined. This
    modifies the prime mover code of those two generator types so that they match. It
    doesn't touch the CS code, which is for single shaft combined units.

    Parameters
    ----------
    df : dataframe
        A dataframe with columns prime_mover_code, and plant_id_eia.
    gens_860 : dataframe
        EIA860 dataframe with technology_description, unit_id_pudl, plant_id_eia
        columns.

    Returns
    -------
    dataframe
        Modified 923 dataframe where prime mover codes at CC generators that don't have
        a PUDL unit id are modified from CA and CT to CC.
    """
    df.loc[
        (df["prime_mover_code"].isin(["CA", "CT"])),
        "prime_mover_code",
    ] = "CC"

    return df


def group_gen_by_year_fuel_primemover(df):
    """
    Group generation and fuel consumption by plant, prime mover, and fuel type. Only
    matters where multiple years of data are used, otherwise output should be the same
    as input.

    Parameters
    ----------
    df : dataframe
        Generation and fuel consumption data from EIA 923 for each plant, prime mover,
        and fuel type

    Returns
    -------
    dataframe
        Sum of generation and fuel consumption data (if multiple years).
    """

    # Group the data by plant, fuel type, and prime mover
    by = [
        "plant_id_eia",
        "fuel_type",
        "energy_source_code",
        "fuel_type_code_pudl",
        "fuel_type_code_aer",
        "prime_mover_code",
    ]
    by = [c for c in by if c in df.columns]
    sort = ["plant_id_eia", "fuel_type", "energy_source_code", "prime_mover_code"]
    sort = [c for c in sort if c in df.columns]

    annual_gen_fuel_923 = (
        (
            df.groupby(  # .drop(columns=["id", "nuclear_unit_id"])
                by=by, as_index=False
            )[
                [
                    "fuel_consumed_units",
                    "fuel_consumed_for_electricity_units",
                    "fuel_consumed_mmbtu",
                    "fuel_consumed_for_electricity_mmbtu",
                    "net_generation_mwh",
                ]
            ].sum()
        )
        .reset_index()
        .drop(columns="index")
        .sort_values(sort)
    )

    return annual_gen_fuel_923


def add_923_heat_rate(df):
    """
    Small function to calculate the heat rate of records with fuel consumption and net
    generation.

    Parameters
    ----------
    df : dataframe
        Must contain the columns net_generation_mwh and
        fuel_consumed_for_electricity_mmbtu

    Returns
    -------
    dataframe
        Same dataframe with new column of heat_rate_mmbtu_mwh
    """

    # Calculate the heat rate for each prime mover/fuel combination
    df["heat_rate_mmbtu_mwh"] = (
        df["fuel_consumed_for_electricity_mmbtu"] / df["net_generation_mwh"]
    )

    return df


def calculate_weighted_heat_rate(heat_rate_df):
    """
    Calculate the weighed heat rate when multiple years of data are used. Net generation
    in each year is used as the weights.

    Parameters
    ----------
    heat_rate_df : dataframe
        Currently the PudlTabl unit_hr method.

    Returns
    -------
    dataframe
        Heat rate weighted by annual generation for each plant and PUDL unit
    """

    def w_hr(df):
        weighted_hr = np.average(
            df["heat_rate_mmbtu_mwh"], weights=df["net_generation_mwh"]
        )
        return weighted_hr

    weighted_unit_hr = heat_rate_df.groupby(["plant_id_eia", "unit_id_pudl"]).apply(
        w_hr
    )
    weighted_unit_hr.name = "heat_rate_mmbtu_mwh"
    weighted_unit_hr = weighted_unit_hr.reset_index()

    return weighted_unit_hr


def plant_pm_heat_rates(annual_gen_fuel_923):
    """
    Calculate the heat rate by plant, prime mover, and fuel type. Values are saved
    as a dictionary.

    Parameters
    ----------
    annual_gen_fuel_923 : dataframe
        Data from the 923 generation and fuel use table. Heat rate for each row should
        already be calculated.

    Returns
    -------
    dict
        Keys are a tuple of plant id, prime mover, and fuel type. Values are the heat
        rate.
    """

    by = ["plant_id_eia", "prime_mover_code", "fuel_type", "energy_source_code"]
    by = [c for c in by if c in annual_gen_fuel_923.columns]
    annual_gen_fuel_923_groups = annual_gen_fuel_923.groupby(by)

    prime_mover_hr_map = {
        _: df["heat_rate_mmbtu_mwh"].values[0] for _, df in annual_gen_fuel_923_groups
    }

    return prime_mover_hr_map


def unit_generator_heat_rates(pudl_out, data_years):
    """
    Calculate the heat rate for each PUDL unit and generators that don't have a PUDL
    unit id.

    Parameters
    ----------
    pudl_out : pudl.PudlTabl
        A PudlTabl object for loading pre-calculated PUDL analysis data
    data_years : list
        Years of data to use

    Returns
    -------
    dataframe, dict
        A dataframe of heat rates for each pudl unit (columsn are ['plant_id_eia',
        'unit_id_pg', 'heat_rate_mmbtu_mwh']).
    """

    # Load the pre-calculated PUDL unit heat rates for selected years.
    # Remove rows without generation or with null values.
    unit_hr = pudl_out.hr_by_unit()
    unit_hr = unit_hr.loc[
        (unit_hr.report_date.dt.year.isin(data_years))
        & (unit_hr.net_generation_mwh > 0),
        :,
    ].dropna()

    weighted_unit_hr = calculate_weighted_heat_rate(unit_hr)

    return weighted_unit_hr


def group_units(df, settings):
    """
    Group by units within a region/technology/cluster. Add a unique unit code
    (plant plus generator) for any generators that aren't part of a unit.


    Returns
    -------
    dataframe
        Grouped generators with the total capacity, minimum load, and average heat
        rate for each.
    """

    by = ["plant_id_eia", "unit_id_pg"]
    # add a unit code (plant plus generator code) in cases where one doesn't exist
    df_copy = df.reset_index()

    # All units should have the same heat rate so taking the mean will just keep the
    # same value.
    grouped_units = df_copy.groupby(by).agg(
        {
            settings["capacity_col"]: "sum",
            "capacity_mwh": "sum",
            "minimum_load_mw": "sum",
            "heat_rate_mmbtu_mwh": "mean",
            "Fixed_OM_Cost_per_MWyr": "mean",
            "Var_OM_Cost_per_MWh": "mean",
        }
    )
    grouped_units = grouped_units.replace([np.inf, -np.inf], np.nan)
    grouped_units = grouped_units.fillna(grouped_units.mean())

    return grouped_units


def calc_unit_cluster_values(
    df: pd.DataFrame, capacity_col: str, technology: str = None
):
    """
    Calculate the total capacity, minimum load, weighted heat rate, and number of
    units/generators in a technology cluster.

    Parameters
    ----------
    df : dataframe
        A dataframe with units/generators of a single technology. One column should be
        'cluster', to label units as belonging to a specific cluster grouping.
    capacity_col: str
        Name of the column with capacity values (e.g. capacity_mw, summer_capacity_mw or
        winter_capacity_mw).
    technology : str, optional
        Name of the generating technology, by default None

    Returns
    -------
    dataframe
        Aggragate values for generators in a technology cluster
    """
    # if not clustering units no need to calulate cluster average values
    if df["cluster"].max() == len(df):
        df["Min_Power"] = df["minimum_load_mw"] / df[capacity_col]
        df["num_units"] = 1
        if technology:
            df["technology"] = technology

        return df.reset_index()

    # Define a function to compute the weighted mean.
    # The issue here is that the df name needs to be used in the function.
    # So this will need to be within a function that takes df as an input
    def wm(x):
        return np.average(x, weights=df.loc[x.index, capacity_col])

    if df["heat_rate_mmbtu_mwh"].isnull().values.any():
        # mean =
        # df["heat_rate_mmbtu_mwh"] = df["heat_rate_mmbtu_mwh"].fillna(
        #     df["heat_rate_mmbtu_mwh"].median()
        # )
        start_cap = df[capacity_col].sum()
        df = df.loc[~df["heat_rate_mmbtu_mwh"].isnull(), :]
        end_cap = df[capacity_col].sum()
        cap_diff = start_cap - end_cap
        logger.warning(f"dropped {cap_diff}MW because of null heat rate values")

    df_values = df.groupby("cluster", as_index=False).agg(
        {
            capacity_col: "mean",
            "capacity_mwh": "sum",
            "minimum_load_mw": "mean",
            "heat_rate_mmbtu_mwh": wm,
            "Fixed_OM_Cost_per_MWyr": wm,
            "Var_OM_Cost_per_MWh": wm,
        }
    )
    df_values.index = df_values["cluster"].values
    if df_values["heat_rate_mmbtu_mwh"].isnull().values.any():
        print(df)
        print(df_values)
    df_values["heat_rate_mmbtu_mwh_iqr"] = df.groupby("cluster").agg(
        {"heat_rate_mmbtu_mwh": iqr}
    )
    df_values["heat_rate_mmbtu_mwh_std"] = df.groupby("cluster").agg(
        {"heat_rate_mmbtu_mwh": "std"}
    )
    df_values["fixed_o_m_mw_std"] = df.groupby("cluster").agg(
        {"Fixed_OM_Cost_per_MWyr": "std"}
    )

    df_values["Min_Power"] = df_values["minimum_load_mw"] / df_values[capacity_col]

    df_values["num_units"] = df.groupby("cluster")["cluster"].count()

    if technology:
        df_values["technology"] = technology

    return df_values


def add_genx_model_tags(df, settings):
    """
    Each generator type needs to have certain tags for use by the GenX model. Each tag
    is a column, e.g. THERM for thermal generators. These columns and tag values are
    defined in the settings file and applied here. Tags are (usually?) boolean 0/1
    values.

    Parameters
    ----------
    df : dataframe
        Clusters of generators. The index should have a column 'technology', which
        is used to map tag values.
    settings : dict
        User-defined settings loaded from a YAML file.

    Returns
    -------
    dataframe
        The original generator cluster results with new columns for each model tag.
    """
    ignored = r"_"
    technology = df["technology"].str.replace(ignored, "")
    # Create a new dataframe with the same index
    default = settings.get("default_model_tag", 0)
    for tag_col in settings.get("model_tag_names", []):
        df[tag_col] = default
        if tag_col not in settings.get("generator_columns", []) and isinstance(
            settings.get("generator_columns"), list
        ):
            settings["generator_columns"].append(tag_col)

        try:
            for tech, tag_value in settings["model_tag_values"][tag_col].items():
                tech = re.sub(ignored, "", tech)
                mask = technology.str.contains(rf"^{tech}", case=False)
                df.loc[mask, tag_col] = tag_value
        except (KeyError, AttributeError) as e:
            logger.warning(f"No model tag values found for {tag_col} ({e})")

    # Change tags with specific regional values for a technology
    flat_regional_tags = flatten(settings.get("regional_tag_values", {}) or {})

    for tag_tuple, tag_value in flat_regional_tags.items():
        region, tag_col, tech = tag_tuple
        tech = re.sub(ignored, "", tech)
        mask = technology.str.contains(rf"^{tech}", case=False)
        df.loc[(df["region"] == region) & mask, tag_col] = tag_value

    return df


def download_860m(settings: dict) -> pd.ExcelFile:
    """Load the entire 860m file into memory as an ExcelFile object.

    Parameters
    ----------
    settings : dict
        User-defined settings loaded from a YAML file. This is where the EIA860m
        filename is defined as the parameter "eia_860m_fn".

    Returns
    -------
    pd.ExcelFile
        The ExcelFile object with all sheets from 860m.
    """
    fn = settings.get("eia_860m_fn")
    if not fn:
        logger.info("Trying to determine the most recent EIA860m file...")
        fn = find_newest_860m()

    engine = None
    ext = fn.split(".")[-1]
    if ext == "xlsx":
        engine = "openpyxl"
    elif ext == "xls":
        engine = "xlrd"

    # Only the most recent file will not have archive in the url
    url = f"https://www.eia.gov/electricity/data/eia860m/xls/{fn}"
    archive_url = f"https://www.eia.gov/electricity/data/eia860m/archive/xls/{fn}"

    local_file = DATA_PATHS["eia_860m"] / fn
    if local_file.exists():
        logger.info(f"Reading a local copy of the EIA860m file {fn}")
        eia_860m = pd.ExcelFile(local_file)
    else:
        logger.info(f"Downloading the EIA860m file {fn}")
        try:
            download_save(url, local_file)
            eia_860m = pd.ExcelFile(local_file, engine=engine)
        except (XLRDError, ValueError, BadZipFile):
            logger.warning("A more recent version of EIA-860m is available")
            download_save(archive_url, local_file)
            eia_860m = pd.ExcelFile(local_file, engine=engine)
        # write the file to disk

    return eia_860m


def find_newest_860m() -> str:
    """Scrape the EIA 860m page to find the most recently posted file.

    Returns
    -------
    str
        Name of most recently posted file
    """
    site_url = "https://www.eia.gov/electricity/data/eia860m/"
    r = requests.get(site_url)
    soup = BeautifulSoup(r.content, "lxml")
    table = soup.find("table", attrs={"class": "basic-table"})
    if not table:
        raise ValueError(
            "Could not determine the most recently posted EIA 860m file. EIA may have "
            "changed their HTML format, please post this as an issue on the PowerGenome "
            "github repository (https://github.com/PowerGenome/PowerGenome/issues/new)."
        )
    href = table.find("a")["href"]
    fn = href.split("/")[-1]

    return fn


def clean_860m_sheet(
    eia_860m: pd.ExcelFile, sheet_name: str, settings: dict
) -> pd.DataFrame:
    """Load a sheet from the 860m ExcelFile object and clean it.

    Parameters
    ----------
    eia_860m : ExcelFile
        Entire 860m file loaded into memory
    sheet_name : str
        Name of the sheet to load as a dataframe
    settings : dict
        User-defined settings loaded from a YAML file.

    Returns
    -------
    pd.DataFrame
        One of the sheets from 860m
    """

    df = eia_860m.parse(sheet_name=sheet_name, na_values=[" "])

    # Find skiprows and skipfooters, which changes across 860m versions.
    # NEW: drop rows with all NaN because EIA added a blank row before the footer.
    sr = 0

    for idx, row in df.iterrows():
        if row.iloc[0] == "Entity ID":
            sr = idx + 1
            break

    sf = 0

    for idx in list(range(-10, 0)):
        if isinstance(df.iloc[idx, 0], str):
            sf = -idx
            break
    df = eia_860m.parse(
        sheet_name=sheet_name, skiprows=sr, skipfooter=sf, na_values=[" "]
    )
    df = df.dropna(how="all")
    df = df.rename(columns=planned_col_map)
    df["plant_id_eia"] = df["plant_id_eia"].astype("Int64")

    if sheet_name in ["Operating", "Planned"]:
        df.loc[:, "operational_status_code"] = df.loc[:, "operational_status"].map(
            op_status_map
        )

    if sheet_name == "Planned":
        df = df.loc[
            df["operational_status_code"].isin(settings["proposed_status_included"]), :
        ]
    df.columns = snake_case_col(df.columns)

    return df


def load_860m(settings: dict) -> Dict[str, pd.DataFrame]:
    """Load the planned, canceled, and retired sheets from an EIA 860m file.

    Parameters
    ----------
    settings : dict
        User-defined settings loaded from a YAML file. This is where the EIA860m
        filename is defined.

    Returns
    -------
    Dict[str, pd.DataFrame]
        The 860m dataframes, with the keys 'planned', 'canceled', and 'retired'.
    """
    sheet_map = {
        "operating": "Operating",
        "planned": "Planned",
        "canceled": "Canceled or Postponed",
        "retired": "Retired",
    }

    fn = settings.get("eia_860m_fn")
    if not fn:
        fn = find_newest_860m()

    fn_name = Path(fn).stem

    data_dict = {}
    eia_860m_excelfile = None
    for name, sheet in sheet_map.items():
        pkl_path = DATA_PATHS["eia_860m"] / f"{fn_name}_{name}.pkl"
        if pkl_path.exists():
            data_dict[name] = pd.read_pickle(pkl_path)
            data_dict[name]["plant_id_eia"] = data_dict[name]["plant_id_eia"].astype(
                "Int64"
            )
            data_dict[name].columns = snake_case_col(data_dict[name].columns)
        else:
            if eia_860m_excelfile is None:
                eia_860m_excelfile = download_860m(settings)
            data_dict[name] = clean_860m_sheet(eia_860m_excelfile, sheet, settings)
            data_dict[name].to_pickle(pkl_path)

    return data_dict


def label_gen_region(
    df: pd.DataFrame, settings: dict, model_regions_gdf: gpd.GeoDataFrame
) -> pd.DataFrame:
    """Label the region that generators in a dataframe belong to based on their
    geographic location. This is done via geospaital join and may not always be accurate
    based on actual utility connections.

    Parameters
    ----------
    df : pd.DataFrame
        Generators that are not assigned to a model region.
    settings : dict
        Need the parameter `capacity_col` to determine which column has capacity.
    model_regions_gdf : gpd.GeoDataFrame
        Contains the name and geometry of each region being used in the study

    Returns
    -------
    pd.DataFrame
        [description]
    """

    no_lat_lon = df.loc[
        (df["latitude"].isnull()) | (df["longitude"].isnull()), :
    ].copy()
    if not no_lat_lon.empty:
        no_lat_lon_cap = no_lat_lon[settings["capacity_col"]].sum()
        logger.warning(
            "Some generators do not have lon/lat data. Check the source "
            "file to determine if they should be included in results. "
            f"\nThe affected generators account for {no_lat_lon_cap} in balancing "
            "authorities: "
            f"\n{no_lat_lon['balancing_authority_code'].tolist()}"
        )

    df = df.dropna(subset=["latitude", "longitude"])

    # Convert the lon/lat values to geo points. Need to add an initial CRS and then
    # change it to align with the IPM regions
    print("Creating gdf")
    gdf = gpd.GeoDataFrame(
        df.copy(),
        geometry=gpd.points_from_xy(df.longitude.copy(), df.latitude.copy()),
        crs="EPSG:4326",
    )
    if gdf.crs != model_regions_gdf.crs:
        gdf = gdf.to_crs(model_regions_gdf.crs)

    gdf = gpd.sjoin(model_regions_gdf.drop(columns="region"), gdf)

    return gdf


def import_new_generators(
    operating_860m: pd.DataFrame,
    gens_860: pd.DataFrame,
    settings: dict,
    model_regions_gdf: gpd.GeoDataFrame,
) -> pd.DataFrame:
    """Find the set of generating units in 860m that are not in the annual 860 data.

    This is especially important for new wind, solar, and battery units, which are built
    on short timelines. Format the data for inclusion with other existing existing units.

    Parameters
    ----------
    operating_860m : pd.DataFrame
        Operating generators from EIA 860m.
    gens_860 : pd.DataFrame
        The set of operating units from other sources (e.g. annual 860 data in PUDL).
    settings : dict
        Requires the parameter "model_regions". Optional parameters include
        "proposed_gen_heat_rates", "proposed_min_load", "capacity_col", "group_technologies",
        and "retirement_ages".
    model_regions_gdf : gpd.GeoDataFrame
        Geospatial representation of the model regions. Used to assign generators to
        model regions.

    Returns
    -------
    pd.DataFrame
        Set of operating generators that were not already in the gens_860 dataframe
    """
    _operating_860m = operating_860m.copy()
    _operating_860m["generator_id"] = _operating_860m["generator_id"].apply(
        remove_leading_zero
    )
    gens_860_id = list(zip(gens_860["plant_id_eia"], gens_860["generator_id"]))
    operating_860m_id = zip(
        _operating_860m["plant_id_eia"], _operating_860m["generator_id"]
    )

    new_mask = [g not in gens_860_id for g in operating_860m_id]
    new_operating = label_gen_region(
        _operating_860m.loc[new_mask, :], settings, model_regions_gdf
    )
    new_operating = new_operating.drop_duplicates(
        subset=["plant_id_eia", "generator_id"]
    )
    new_operating.loc[:, "heat_rate_mmbtu_mwh"] = new_operating.loc[
        :, "technology_description"
    ].map(settings.get("proposed_gen_heat_rates", {}) or {})

    # The default EIA heat rate for non-thermal technologies is 9.21
    new_operating.loc[
        new_operating["heat_rate_mmbtu_mwh"].isnull(), "heat_rate_mmbtu_mwh"
    ] = 9.21

    new_operating.loc[:, "minimum_load_mw"] = (
        new_operating["technology_description"].map(
            settings.get("proposed_min_load", {}) or {}
        )
        * new_operating[settings.get("capacity_col", "capacity_mw")]
    )

    # Assume anything else being built at scale is wind/solar and will have a Min_power
    # of 0
    new_operating.loc[new_operating["minimum_load_mw"].isnull(), "minimum_load_mw"] = 0

    new_operating = new_operating.set_index(
        ["plant_id_eia", "prime_mover_code", "energy_source_code_1"]
    )

    # Add a retirement year based on the planned start year
    label_retirement_year(
        df=new_operating,
        settings=settings,
        age_col="operating_year",
        add_additional_retirements=False,
    )
    if (
        new_operating.loc[new_operating["technology_description"].isnull(), :].empty
        is False
    ):
        plant_ids = list(
            new_operating.loc[new_operating["technology_description"].isnull(), :]
            .index.get_level_values("plant_id_eia")
            .to_numpy()
        )
        plant_capacity = new_operating.loc[
            new_operating["technology_description"].isnull(),
            settings.get("capacity_col", "capacity_mw"),
        ].sum()

        logger.warning(
            f"The EIA860 file has {len(plant_ids)} operating generator(s) without a technology "
            f"description. The plant IDs are {plant_ids}, and they have a combined "
            f"capcity of {plant_capacity} MW."
        )

    if settings.get("group_technologies"):
        new_operating = group_technologies(
            new_operating,
            settings["group_technologies"],
            settings.get("tech_groups", {}) or {},
            settings.get("regional_no_grouping", {}) or {},
        )
        print(new_operating["technology_description"].unique().tolist())

    keep_cols = [
        "model_region",
        "technology_description",
        "generator_id",
        settings.get("capacity_col", "capacity_mw"),
        "capacity_mwh",
        "minimum_load_mw",
        "operational_status_code",
        "heat_rate_mmbtu_mwh",
        "retirement_year",
        "operating_year",
        "state",
    ]

    return new_operating.reindex(columns=keep_cols)


def import_proposed_generators(
    planned: pd.DataFrame, settings: dict, model_regions_gdf: gpd.GeoDataFrame
) -> pd.DataFrame:
    """
    Load the most recent proposed generating units from EIA860m. Will also add
    any planned generators that are included in the settings file.

    Parameters
    ----------
    planned : pd.DataFrame
        Generators that are not assigned to a model region.
    settings : dict
        User defined parameters from a settings YAML file
    model_regions_gdf : gpd.GeoDataFrame
        Contains the name and geometry of each region being used in the study

    Returns
    -------
    pd.DataFrame
        All proposed generators.
    """

    # Some plants don't have lat/lon data. Log this now to determine if any action is
    # needed, then drop them from the dataframe.
    # no_lat_lon = planned.loc[
    #     (planned["latitude"].isnull()) | (planned["longitude"].isnull()), :
    # ].copy()
    # if not no_lat_lon.empty:
    #     no_lat_lon_cap = no_lat_lon[settings["capacity_col"]].sum()
    #     logger.warning(
    #         "Some generators do not have lon/lat data. Check the source "
    #         "file to determine if they should be included in results. "
    #         f"\nThe affected generators account for {no_lat_lon_cap} in balancing "
    #         "authorities: "
    #         f"\n{no_lat_lon['balancing_authority_code'].tolist()}"
    #     )

    # planned = planned.dropna(subset=["latitude", "longitude"])

    # # Convert the lon/lat values to geo points. Need to add an initial CRS and then
    # # change it to align with the IPM regions
    # print("Creating gdf")
    # planned_gdf = gpd.GeoDataFrame(
    #     planned.copy(),
    #     geometry=gpd.points_from_xy(planned.longitude.copy(), planned.latitude.copy()),
    #     crs="EPSG:4326",
    # )
    # if planned_gdf.crs != model_regions_gdf.crs:
    #     planned_gdf = planned_gdf.to_crs(model_regions_gdf.crs)

    # planned_gdf = gpd.sjoin(model_regions_gdf.drop(columns="IPM_Region"), planned_gdf)
    _planned = planned.copy()
    _planned["generator_id"] = _planned["generator_id"].apply(remove_leading_zero)
    planned_gdf = label_gen_region(_planned, settings, model_regions_gdf)
    planned_gdf = planned_gdf.drop_duplicates(subset=["plant_id_eia", "generator_id"])

    # Add planned additions from the settings file
    additional_planned = settings.get("additional_planned") or []
    for record in additional_planned:
        plant_id, gen_id, model_region = record
        plant_record = planned.loc[
            (planned["plant_id_eia"] == plant_id) & (planned["generator_id"] == gen_id),
            :,
        ]
        plant_record["model_region"] = model_region

        planned_gdf = planned_gdf.append(plant_record, sort=False)

    logger.info(
        f"{len(additional_planned)} generators were added to the planned list based on settings"
    )

    planned_gdf.loc[:, "heat_rate_mmbtu_mwh"] = planned_gdf.loc[
        :, "technology_description"
    ].map(settings["proposed_gen_heat_rates"])

    # The default EIA heat rate for non-thermal technologies is 9.21
    planned_gdf.loc[
        planned_gdf["heat_rate_mmbtu_mwh"].isnull(), "heat_rate_mmbtu_mwh"
    ] = 9.21

    planned_gdf.loc[:, "minimum_load_mw"] = (
        planned_gdf["technology_description"].map(settings["proposed_min_load"])
        * planned_gdf[settings["capacity_col"]]
    )

    # Assume anything else being built at scale is wind/solar and will have a Min_Power
    # of 0
    planned_gdf.loc[planned_gdf["minimum_load_mw"].isnull(), "minimum_load_mw"] = 0

    planned_gdf = planned_gdf.set_index(
        ["plant_id_eia", "prime_mover_code", "energy_source_code_1"]
    )

    if (
        planned_gdf.loc[planned_gdf["technology_description"].isnull(), :].empty
        is False
    ):
        plant_ids = list(
            planned_gdf.loc[planned_gdf["technology_description"].isnull(), :]
            .index.get_level_values("plant_id_eia")
            .to_numpy()
        )
        plant_capacity = planned_gdf.loc[
            planned_gdf["technology_description"].isnull(), settings["capacity_col"]
        ].sum()

        logger.warning(
            f"The EIA860 file has {len(plant_ids)} proposed generator(s) without a technology "
            f"description. The plant IDs are {plant_ids}, and they have a combined "
            f"capcity of {plant_capacity} MW."
        )

    # Add a retirement year based on the planned start year
    label_retirement_year(
        df=planned_gdf,
        settings=settings,
        age_col="planned_operating_year",
        add_additional_retirements=False,
    )

    if settings.get("group_technologies"):
        planned_gdf = group_technologies(
            planned_gdf,
            settings["group_technologies"],
            settings.get("tech_groups", {}) or {},
            settings.get("regional_no_grouping", {}) or {},
        )
        print(planned_gdf["technology_description"].unique().tolist())

    keep_cols = [
        "model_region",
        "technology_description",
        "generator_id",
        settings["capacity_col"],
        "minimum_load_mw",
        "operational_status_code",
        "heat_rate_mmbtu_mwh",
        "planned_operating_year",
        "retirement_year",
    ]

    return planned_gdf.loc[:, keep_cols]


def gentype_region_capacity_factor(
    pudl_engine, plant_region_map, settings, years_filter=None
):
    """
    Calculate the average capacity factor for all generators of a type/region. This
    uses all years of available data unless otherwise specified. The potential
    generation is calculated for every year a plant is in operation using the capacity
    type specified in settings (nameplate, summer, or winter) and the number of hours
    in each year.

    As of this time PUDL only has generation data back to 2011.

    Parameters
    ----------
    pudl_engine : sqlalchemy.Engine
        A sqlalchemy connection for use by pandas
    plant_region_map : dataframe
        A dataframe with the region for every plant
    settings : dictionary
        The dictionary of settings with a dictionary of region aggregations

    Returns
    -------
    DataFrame
        A dataframe with the capacity factor of every selected technology
    """
    data_years = [str(y) for y in settings["data_years"]]
    data_years.extend(settings.get("capacity_factor_default_year_filter", []))
    if type(settings.get("alt_year_filters")) is dict:
        for tech, value in settings["alt_year_filters"].items():
            if isinstance(value, list):
                data_years.extend(value)
            else:
                data_years.append(value)
    cap_col = settings["capacity_col"]

    # Include standby (SB) generators since they are in our capacity totals
    sql = f"""
        SELECT
            G.report_date,
            G.plant_id_eia,
            G.generator_id,
            SUM(G.capacity_mw) AS capacity_mw,
            SUM(G.summer_capacity_mw) as summer_capacity_mw,
            SUM(G.winter_capacity_mw) as winter_capacity_mw,
            G.technology_description,
            G.fuel_type_code_pudl
        FROM
            generators_eia860 G
        WHERE operational_status_code NOT IN ('RE', 'OS', 'IP', 'CN')
        AND strftime('%Y',report_date) in ({','.join(['?']*len(data_years))})
        GROUP BY
            G.report_date,
            G.plant_id_eia,
            G.technology_description,
            G.fuel_type_code_pudl,
            G.generator_id
        ORDER by G.plant_id_eia, G.report_date
    """

    plant_gen_tech_cap = pd.read_sql_query(
        sql, pudl_engine, params=data_years, parse_dates=["report_date"]
    )
    plant_gen_tech_cap = plant_gen_tech_cap.loc[
        plant_gen_tech_cap["plant_id_eia"].isin(plant_region_map["plant_id_eia"]), :
    ]

    plant_gen_tech_cap = fill_missing_tech_descriptions(plant_gen_tech_cap)
    plant_tech_cap = group_generators_at_plant(
        df=plant_gen_tech_cap,
        by=["plant_id_eia", "report_date", "technology_description"],
        agg_fn={cap_col: "sum"},
    )

    plant_tech_cap = plant_tech_cap.merge(
        plant_region_map, on="plant_id_eia", how="left"
    )

    label_small_hydro(plant_tech_cap, settings, by=["plant_id_eia", "report_date"])

    sql = """
        SELECT
            strftime('%Y', GF.report_date) AS report_date,
            GF.plant_id_eia,
            SUM(GF.net_generation_mwh) AS net_generation_mwh,
            GF.fuel_type_code_pudl
        FROM
            generation_fuel_eia923 GF
        GROUP BY
            strftime('%Y', GF.report_date),
            GF.plant_id_eia,
            GF.fuel_type_code_pudl
        ORDER by GF.plant_id_eia, strftime('%Y', GF.report_date)
    """
    generation = pd.read_sql_query(sql, pudl_engine, parse_dates={"report_date": "%Y"})

    if pudl.__version__ > "0.5.0":
        by = ["plant_id_eia"]
    else:
        by = {"plant_id_eia": "eia"}
    capacity_factor = pudl.helpers.clean_merge_asof(
        generation,
        plant_tech_cap,
        left_on="report_date",
        right_on="report_date",
        by=by,
    )

    if settings.get("group_technologies"):
        capacity_factor = group_technologies(
            capacity_factor,
            settings["group_technologies"],
            settings.get("tech_groups", {}) or {},
            settings.get("regional_no_grouping", {}) or {},
        )

    if years_filter is None:
        years_filter = {
            tech: settings["capacity_factor_default_year_filter"]
            for tech in plant_gen_tech_cap["technology_description"].unique()
        }
        if type(settings.get("alt_year_filters")) is dict:
            for tech, value in settings["alt_year_filters"].items():
                years_filter[tech] = value

        data_years = plant_gen_tech_cap["report_date"].dt.year.unique()

        # Use all years where the value is None

        for tech, value in years_filter.items():
            if value is None:
                years_filter[tech] = data_years

    df_list = []
    for tech, years in years_filter.items():
        _df = capacity_factor.loc[
            (capacity_factor["technology_description"] == tech)
            & (capacity_factor["report_date"].dt.year.isin(years)),
            :,
        ]
        df_list.append(_df)
    capacity_factor = pd.concat(df_list, sort=False)

    # get a unique set of dates to generate the number of hours
    dates = capacity_factor["report_date"].drop_duplicates()
    dates_to_hours = pd.DataFrame(
        data={
            "report_date": dates,
            "hours": dates.apply(
                lambda d: (
                    pd.date_range(d, periods=2, freq="YS")[1]
                    - pd.date_range(d, periods=2, freq="YS")[0]
                )
                / pd.Timedelta(hours=1)
            ),
        }
    )

    # merge in the hours for the calculation
    capacity_factor = capacity_factor.merge(dates_to_hours, on=["report_date"])
    capacity_factor["potential_generation_mwh"] = (
        capacity_factor[cap_col] * capacity_factor["hours"]
    )

    capacity_factor_tech_region = capacity_factor.groupby(
        ["model_region", "technology_description"], as_index=False
    )[["potential_generation_mwh", "net_generation_mwh"]].sum()

    # actually calculate capacity factor wooo!
    capacity_factor_tech_region["capacity_factor"] = (
        capacity_factor_tech_region["net_generation_mwh"]
        / capacity_factor_tech_region["potential_generation_mwh"]
    )

    capacity_factor_tech_region.rename(
        columns={"model_region": "region", "technology_description": "technology"},
        inplace=True,
    )

    logger.debug(capacity_factor_tech_region)

    return capacity_factor_tech_region


def add_fuel_labels(df, fuel_prices, settings):
    """Add a Fuel column with the approproriate regional fuel for each generator type

    Parameters
    ----------
    df : DataFrame
        Generator clusters dataframe with all existing and proposed technologies
    fuel_prices : DataFrame
        Prices of fuels from EIA AEO scenarios in each census region. Columns include
        ['year', 'price', 'fuel', 'region', 'scenario', 'full_fuel_name']
    settings : dictionary
        The dictionary of settings with fuel price variables

    Returns
    -------
    DataFrame
        Same as input, but with a new column "Fuel" that is either the name of the
        corresponding fuel (coal, natural_gas, uranium, or distillate) or "None".

    Raises
    ------
    KeyError
        The model region is not mapped to a fuel region in 'aeo_fuel_region_map'
    """

    df["Fuel"] = np.nan
    # This variable is called eia_tech but it can be any tech name or a mapping from
    # EIA technologies through to other techs via "eia_atb_tech_map"
    for eia_tech, fuel in (settings.get("tech_fuel_map") or {}).items():
        try:
            if eia_tech == "Natural Gas Steam Turbine":
                # No ATB natural gas steam turbine and I match it with coal for O&M
                # which would screw this up and list natural gas as a fuel for ATB
                # coal plants
                atb_tech = None
            else:
                if not isinstance(settings["eia_atb_tech_map"][eia_tech], list):
                    settings["eia_atb_tech_map"][eia_tech] = [
                        settings["eia_atb_tech_map"][eia_tech]
                    ]
                atb_tech = [
                    tech.split("_")[0] + "_"
                    for tech in settings["eia_atb_tech_map"][eia_tech]
                ]
        except KeyError:
            # No corresponding ATB technology
            atb_tech = None
        scenario = settings.get("aeo_fuel_scenarios", {}).get(fuel)
        model_year = settings["model_year"]
        if not scenario:
            if fuel not in settings.get("user_fuel_price", []) or []:
                raise KeyError(
                    f"The fuel type '{fuel}' is not in the settings parameters "
                    "'aeo_fuel_scenarios' or 'user_fuel_price'. All fuels listed in "
                    "'tech_fuel_map' must be included in one of these."
                )
            if isinstance(settings["user_fuel_price"][fuel], dict):
                for region, price in settings["user_fuel_price"][fuel].items():
                    fuel_name = f"{region}_{fuel}"
                    df.loc[
                        (
                            df["technology"].str.rstrip("_").str.lower()
                            == eia_tech.lower()
                        )
                        & (df["region"] == region),
                        "Fuel",
                    ] = fuel_name

                    if atb_tech is not None:
                        for tech in atb_tech:
                            df.loc[
                                (df["technology"].str.contains(tech, case=False))
                                & (df["region"] == region)
                                & (df["Fuel"].isna()),
                                "Fuel",
                            ] = fuel_name
            else:
                df.loc[
                    (df["technology"].str.rstrip("_").str.lower() == eia_tech.lower())
                    & (df["Fuel"].isna()),
                    "Fuel",
                ] = fuel

                if atb_tech is not None:
                    for tech in atb_tech:
                        df.loc[
                            (df["technology"].str.contains(tech, case=False))
                            & (df["Fuel"].isna()),
                            "Fuel",
                        ] = fuel
        else:
            for aeo_region, model_regions in settings["aeo_fuel_region_map"].items():
                fuel_name = ("_").join([aeo_region, scenario, fuel])
                assert (
                    fuel_prices.query(
                        "year==@model_year & full_fuel_name==@fuel_name"
                    ).empty
                    is False
                ), f"{fuel_name} doesn't show up in {model_year}"

                df.loc[
                    (df["technology"].str.contains(eia_tech, case=False))
                    & df["region"].isin(model_regions),
                    "Fuel",
                ] = fuel_name

                if atb_tech is not None:
                    for tech in atb_tech:
                        df.loc[
                            (df["technology"].str.contains(tech, case=False))
                            & (df["region"].isin(model_regions))
                            & (df["Fuel"].isna()),
                            "Fuel",
                        ] = fuel_name

    for ccs_tech, ccs_fuel in (settings.get("ccs_fuel_map") or {}).items():
        ccs_base_name = ("_").join(ccs_fuel.split("_")[:-1])
        if ccs_base_name in (settings.get("aeo_fuel_scenarios", {}) or {}).keys():
            scenario = settings["aeo_fuel_scenarios"][ccs_base_name]
            for aeo_region, model_regions in settings["aeo_fuel_region_map"].items():
                ccs_fuel_name = ("_").join([aeo_region, scenario, ccs_fuel])

                df.loc[
                    (df["technology"].str.contains(ccs_tech))
                    & df["region"].isin(model_regions),
                    "Fuel",
                ] = ccs_fuel_name
        elif ccs_base_name in (settings.get("user_fuel_price", {}) or {}).keys():
            if isinstance(settings["user_fuel_price"][ccs_base_name], dict):
                for region in settings["user_fuel_price"][ccs_base_name].keys():
                    ccs_fuel_name = ("_").join([region, ccs_fuel])
                    df.loc[
                        (df["technology"].str.contains(ccs_tech, case=False))
                        & (df["region"] == region),
                        "Fuel",
                    ] = ccs_fuel_name
            else:
                df.loc[
                    (df["technology"].str.contains(ccs_tech, case=False)),
                    "Fuel",
                ] = ccs_fuel
        else:
            logger.warning(
                f"The fuel {ccs_fuel} is included in settings parameter `ccs_fuel_map` "
                "but it can't be matched against an AEO or user fuel. CCS fuels should "
                "have the format <fuel name>_ccs<capture rate>, where the capture rate "
                "is optional. The <fuel name> should match a fuel from `aeo_fuel_scenarios' "
                "or `user_fuel_prices`."
            )

    # Replace AEO region name with model region in cases where users are modifying AEO price
    model_aeo_region_map = reverse_dict_of_lists(
        settings.get("aeo_fuel_region_map", {})
    )
    for region, adj in (settings.get("regional_fuel_adjustments", {}) or {}).items():
        aeo_region = model_aeo_region_map.get(region)
        if not aeo_region:
            raise KeyError(
                f"There is no mapping of the model region {region} to an AEO fuel region "
                "in the settings parameter 'aeo_fuel_region_map'."
            )
        if isinstance(adj, list):
            # Replace the aeo region name with model region for all resources
            df.loc[
                (df["Fuel"].str.contains(aeo_region))
                & (df["Fuel"].notna())
                & (df["region"].str.lower() == region.lower()),
                "Fuel",
            ] = df.loc[
                (df["Fuel"].str.contains(aeo_region))
                & (df["Fuel"].notna())
                & (df["region"].str.lower() == region.lower()),
                "Fuel",
            ].str.replace(
                aeo_region, region
            )
        if isinstance(adj, dict):
            # Replace the aeo region name with model region only for select fuels
            for fuel, op in adj.items():
                df.loc[
                    (df["Fuel"].str.contains(aeo_region))
                    & (df["Fuel"].str.contains(fuel))
                    & (df["Fuel"].notna())
                    & (df["region"].str.lower() == region.lower()),
                    "Fuel",
                ] = df.loc[
                    (df["Fuel"].str.contains(aeo_region))
                    & (df["Fuel"].str.contains(fuel))
                    & (df["Fuel"].notna())
                    & (df["region"].str.lower() == region.lower()),
                    "Fuel",
                ].str.replace(
                    aeo_region, region
                )

    df.loc[df["Fuel"].isna(), "Fuel"] = "None"

    return df


def calculate_transmission_inv_cost(resource_df, settings, offshore_spur_costs=None):
    """Calculate the transmission investment cost for each new resource.

    Parameters
    ----------
    resource_df : DataFrame
        Each row represents a single resource within a region. Should have columns
        `region` and `<type>_miles`, where transmission <type> is one of
        'spur', 'offshore_spure', or 'tx'.
    settings : dict
        A dictionary of user-supplied settings. Must have key
        `transmission_investment_cost` with the format:
            - <type>
                - `capex_mw_mile` (float)
                - `wacc` (float)
                - `investment_years` (int)
            - ...
    offshore_spur_costs : DataFrame
        Offshore spur costs per mile in the format
        `technology` ('OffShoreWind'), `tech_detail`, `cost_case`, and `capex_mw_mile`.
        Only used if `settings.transmission_investment_cost.capex_mw_mile` is missing.

    Returns
    -------
    DataFrame
        Modified copy of the input dataframe with new columns '<type>_capex' and
        '<type>_inv_mwyr' for each column `<type>_miles`.

    Raises
    ------
    KeyError
        Settings missing transmission types present in resources.
    KeyError
        Settings missing required keys.
    KeyError
        Setting capex_mw_mile missing regions present in resources.
    TypeError
        Setting capex_mw_mile is neither a dictionary nor a numeric value.
    """
    SETTING = "transmission_investment_cost"
    KEYS = ["wacc", "investment_years", "capex_mw_mile"]
    ttypes = settings.get(SETTING, {})
    # Check coverage of transmission types in resources
    resource_ttypes = [x for x in TRANSMISSION_TYPES if f"{x}_miles" in resource_df]
    missing_ttypes = list(set(resource_ttypes) - set(ttypes))
    if missing_ttypes:
        raise KeyError(f"{SETTING} missing transmission line types {missing_ttypes}")
    # Apply calculation for each transmission type
    regions = resource_df["region"].unique()
    use_offshore_spur_costs = False
    for ttype, params in ttypes.items():
        if ttype not in resource_ttypes:
            continue
        if (
            ttype == "offshore_spur"
            and offshore_spur_costs is not None
            and not params.get("capex_mw_mile")
        ):
            use_offshore_spur_costs = True
            # Build technology: capex_mw_mile map
            params = params.copy()
            params["capex_mw_mile"] = (
                offshore_spur_costs.assign(
                    technology=offshore_spur_costs[
                        ["technology", "tech_detail", "cost_case"]
                    ]
                    .astype(str)
                    .agg("_".join, axis=1)
                )
                .set_index("technology")["capex_mw_mile"]
                .to_dict()
            )
        # Check presence of required keys
        missing_keys = list(set(KEYS) - set(params))
        if missing_keys:
            raise KeyError(f"{SETTING}.{ttype} missing required keys {missing_keys}")
        if isinstance(params["capex_mw_mile"], dict):
            if use_offshore_spur_costs:
                capex_mw_mile = resource_df["technology"].map(params["capex_mw_mile"])
            else:
                # Check coverage of regions in resources
                missing_regions = list(set(regions) - set(params["capex_mw_mile"]))
                if missing_regions:
                    raise KeyError(
                        f"{SETTING}.{ttype}.capex_mw_mile missing regions {missing_regions}"
                    )
                capex_mw_mile = resource_df["region"].map(params["capex_mw_mile"])
        elif isinstance(params["capex_mw_mile"], Number):
            capex_mw_mile = params["capex_mw_mile"]
        else:
            raise TypeError(
                f"{SETTING}.{ttype}.capex_mw_mile should be numeric or a dictionary"
                f" of <region>: <capex>, not {params['capex_mw_mile']}"
            )
        resource_df[f"{ttype}_capex"] = (
            capex_mw_mile.fillna(0) * resource_df[f"{ttype}_miles"]
        )
        resource_df[f"{ttype}_inv_mwyr"] = investment_cost_calculator(
            resource_df[f"{ttype}_capex"], params["wacc"], params["investment_years"]
        )
    return resource_df


def add_transmission_inv_cost(
    resource_df: pd.DataFrame, settings: dict
) -> pd.DataFrame:
    """Add tranmission investment costs to plant investment costs

    Parameters
    ----------
    resource_df
        Each row represents a single resource within a region. Should have columns
        `Inv_Cost_per_MWyr` and transmission costs.
            - one or more `<type>_inv_mwyr`,
                where <type> is 'spur', 'offshore_spur', or 'tx'.
            - `interconnect_annuity`
    settings
        User settings. If `transmission_investment_cost.use_total` is present and true,
        `interconnect_annuity` is used over `<type>_inv_mwys` if present, not null,
        and not zero.

    Returns
    -------
    DataFrame
        A modified copy of the input dataframe where 'Inv_Cost_per_MWyr' represents the
        combined plant and transmission investment costs. The new column
        `plant_inv_cost_mwyr` represents just the plant investment costs.
    """
    use_total = (
        settings.get("transmission_investment_cost", {}).get("use_total", False)
        and "interconnect_annuity" in resource_df
    )
    resource_df["plant_inv_cost_mwyr"] = resource_df["Inv_Cost_per_MWyr"]
    columns = [
        c for c in [f"{t}_inv_mwyr" for t in TRANSMISSION_TYPES] if c in resource_df
    ]
    cost = resource_df[columns].sum(axis=1)
    if use_total:
        total = resource_df["interconnect_annuity"]
        has_total = ~total.isna() & total != 0
        cost[has_total] = total[has_total]
    if cost.isna().any() or (cost == 0).any():
        logger.warning(
            "Transmission investment costs are missing or zero for some resources"
            " and will not be included in the total investment costs."
        )
    resource_df["Inv_Cost_per_MWyr"] += cost
    return resource_df


def save_weighted_hr(weighted_unit_hr, pudl_engine):
    pass


def add_dg_resources(
    pg_engine: sqlalchemy.engine.Engine,
    settings: dict,
    gen_df: pd.DataFrame = pd.DataFrame(),
) -> pd.DataFrame:
    """Add distributed generation resources as rows in a generators dataframe

    Parameters
    ----------
    pg_engine : sqlalchemy.engine.Engine
        Connection to database with hourly generation values. Needed if installed DG
        capacity is calculated as a fraction of load.
    settings : dict
        Settings dictionary with parameters "model_year", "input_folder", "distributed_gen_profiles_fn",
        "distributed_gen_method", "distributed_gen_values", and "avg_distribution_loss".
    gen_df : pd.DataFrame, optional
        A dataframe with other generating resources, by default pd.DataFrame()

    Returns
    -------
        A modified version of the input dataframe with distributed generation resources
        for each region where a generation profile has been supplied in the
        "distributed_gen_profiles_fn" file. Each dg resource is one row and includes
        values for the columns "technology", "region", "Existing_Cap_MW", and "profile".
    """
    dg_profiles = make_distributed_gen_profiles(pg_engine, settings)
    df = pd.DataFrame(
        columns=["technology", "region", "cluster", "Existing_Cap_MW", "profile"],
        index=range(len(dg_profiles.columns)),
    )

    for idx, (region, s) in enumerate(dg_profiles.iteritems()):
        cap = s.max()
        df.loc[idx, "profile"] = (s / cap).round(3).to_numpy()
        df.loc[idx, "Existing_Cap_MW"] = cap.round(0).astype(int)
    df["technology"] = "distributed_generation"
    df["region"] = dg_profiles.columns
    df["cluster"] = 1

    return pd.concat([gen_df, df], ignore_index=True)


def energy_storage_mwh(
    df: pd.DataFrame,
    energy_storage_duration: Dict[str, Union[float, Dict[str, float]]],
    tech_col: str,
    cap_col: str,
    energy_col: str,
) -> pd.DataFrame:
    """Convert resource capacity (MW) to MWh using a dictionary with storage duration
    by technology name.

    Parameters
    ----------
    df : pd.DataFrame
        Resource dataframe with columns specified by `tech_col`, `cap_col`, and
        `energy_col`
    energy_storage_duration : Dict[str, Union[float, Dict[str, float]]]
        Keys are technology names, values are either the duration of storage (float) or
        a dictionary with region keys and storage duration values
    tech_col : str
        Dataframe column with technology names
    cap_col : str
        Dataframe column with technology capacity (power)
    energy_col : str
        Dataframe column to fill with technology energy storage

    Returns
    -------
    pd.DataFrame
        Modified dataframe with energy storage values
    """
    context = "Setting energy storage MWh in existing generators."
    region_col = find_region_col(df.columns, context)
    all_regions = df[region_col].unique()

    if energy_col not in df.columns:
        df[energy_col] = 0

    storage_techs = list(df.loc[df[energy_col] > 0, tech_col].unique())
    partial_storage = list(
        df.loc[
            (df[tech_col].isin(storage_techs)) & ~(df[energy_col] > 0),
            tech_col,
        ].unique()
    )
    missing_techs = [t for t in partial_storage if t not in energy_storage_duration]
    if missing_techs:
        logger.warning(
            "\n\n**************************\n"
            f"The storage technology(ies) {missing_techs} have some existing generators "
            "with energy capacity (MWh) values and some where the energy capacity is "
            "missing. You have not included these technologies in the settings parameter "
            "'energy_storage_duration', which is used to fill missing energy capacity "
            "data.\n\nNOTE: This is not a comprehensive list of technologies that *should* "
            "be included in 'energy_storage_duration'. Technologies without any existing "
            "energy capacity data might also be missing.\n"
            "**************************\n"
        )
    for tech, val in energy_storage_duration.items():
        if isinstance(val, dict):
            tech_regions = val.keys()
            model_tech_regions = df.loc[
                snake_case_col(df[tech_col]).str.contains(snake_case_str(tech)),
                region_col,
            ].to_list()
            if not all(r in tech_regions for r in model_tech_regions):
                missing_regions = [
                    r for r in model_tech_regions if r not in tech_regions
                ]
                logger.warning(
                    f"The regions {missing_regions} are missing from technology {tech} "
                    "in the settings parameter 'energy_storage_duration'. This technology "
                    "will not have any energy storage capacity in these regions."
                )
            for region, v in val.items():
                if region not in all_regions:
                    logger.warning(
                        f"The settings parameter 'energy_storage_duration', technology '{tech}' "
                        f"has the region '{region}', which is not one of your model regions."
                    )
                df.loc[
                    (snake_case_col(df[tech_col]).str.contains(snake_case_str(tech)))
                    & (df[region_col] == region)
                    & ~(df[energy_col] > 0),
                    energy_col,
                ] = (
                    df[cap_col] * v
                )
        else:
            df.loc[
                snake_case_col(df[tech_col]).str.contains(snake_case_str(tech))
                & ~(df[energy_col] > 0),
                energy_col,
            ] = (
                df[cap_col] * val
            )
    df[energy_col] = df[energy_col].fillna(0)
    return df


def load_plants_860(
    pudl_engine: sqlalchemy.engine.Engine, data_years: List[int] = [2020]
) -> pd.DataFrame:
    """Load database table with EIA860 information on plants

    Parameters
    ----------
    pudl_engine : sqlalchemy.engine.Engine
        Connection to PUDL database
    data_years : List[int], optional
        Year of data to keep, by default [2020]

    Returns
    -------
    pd.DataFrame
        Includes all columns from the database table
    """
    data_years = [str(y) for y in data_years]
    s = f"""
    SELECT * from plants_eia860
    WHERE strftime('%Y',report_date) in ({','.join(['?']*len(data_years))})
    """
    plants = pd.read_sql_query(
        s, pudl_engine, params=data_years, parse_dates=["report_date"]
    )

    return plants


def load_demand_response_efs_profile(
    resource: str,
    electrification_stock_fn: str,
    model_year: int,
    electrification_scenario: str,
    model_regions: list,
    region_aggregations: dict = {},
    path_in: Path = None,
) -> pd.DataFrame:
    """Load the demand profile of a single flexible resource in all model regions.

    Parameters
    ----------
    resource : str
        Name of the flexible resource.
    electrification_stock_fn : str
        Name of the data file with stock values for each year.
    model_year : int
        Planning period or model year. Used to select stock values for flexible resources
        and their demand profiles.
    electrification_scenario : str
        Name of a scenario from the stock data file.
    model_regions : list
        Names of the model regions, including those that are aggregated from multiple
        base regions.
    region_aggregations : dict, optional
        A list of base regions for each aggregated model region, by default {}. For
        example, {"CA_N": ["WEC_BANC", "WEC_CALN"]}.
    path_in : Path, optional
        Folder where stock and incremental factor (profile) data are located, by default
        None.

    Returns
    -------
    pd.DataFrame
        Flexible demand profiles for the selected resource in each model region. Column
        names are model regions.
    Raises
    ------
    KeyError
        The resource is not a valid name.
    """
    from powergenome.load_construction import electrification_profiles

    keep_regions, region_agg_map = regions_to_keep(model_regions, region_aggregations)
    elec_profiles = electrification_profiles(
        stock_fn=electrification_stock_fn,
        year=model_year,
        elec_scenario=electrification_scenario,
        regions=keep_regions,
        path_in=path_in,
    )

    if resource not in elec_profiles.resource.unique():
        raise KeyError(
            f"No profile was available for the flexible resource '{resource}' specified "
            "in your settings file under 'flexible_demand_resources'. Available "
            f"resources include: {list(elec_profiles.resource.unique())}."
        )
    dr_profile = elec_profiles.loc[
        elec_profiles["resource"] == resource, ["time_index", "region", "load_mw"]
    ].pivot(index="time_index", columns="region")
    dr_profile.columns = dr_profile.columns.droplevel()
    for model_region, base_regs in region_aggregations.items():
        base_regs = [r for r in base_regs if r != model_region]
        dr_profile[model_region] = dr_profile.reindex(
            columns=base_regs, fill_value=0
        ).sum(axis=1)
        dr_profile = dr_profile.drop(columns=base_regs, errors="ignore")

    return dr_profile


def add_860m_storage_mwh(
    gen_df: pd.DataFrame, operating_860m: pd.DataFrame, storage_techs: List[str] = None
) -> pd.DataFrame:
    idx_cols = ["plant_id_eia", "generator_id"]
    gen_df = gen_df.set_index(idx_cols)

    if not storage_techs:
        storage_techs = [
            "Batteries",
            "Natural Gas with Compressed Air Storage",
            "Flywheels",
        ]

    storage_860m = operating_860m.loc[
        operating_860m.technology_description.isin(storage_techs), :
    ].set_index(idx_cols)
    gen_df.loc[storage_860m.index, "capacity_mwh"] = storage_860m["capacity_mwh"]
    # gen_df = pd.merge(
    #     gen_df, storage_860m[idx_cols + ["capacity_mwh"]], how="left", on=idx_cols
    # )

    return gen_df.reset_index()


class GeneratorClusters:
    """
    This class is used to determine genererating units that will likely be operating
    in a given year, clusters them according to parameters for the settings file,
    and determines the average operating characteristics of each cluster. Structuring
    this as a class isn't strictly necessary but makes it easier to access generator
    data part-way through the process.
    """

    def __init__(
        self,
        pudl_engine,
        pudl_out,
        pg_engine,
        settings,
        current_gens=True,
        supplement_with_860m=True,
        sort_gens=False,
        plant_region_map_table="plant_region_map_epaipm",
        settings_agg_key="region_aggregations",
    ):
        """

        Parameters
        ----------
        pudl_engine : sqlalchemy.Engine
            A sqlalchemy connection for use by pandas
        pudl_out : pudl.PudlTabl
            A PudlTabl object for loading pre-calculated PUDL analysis data
        settings : dictionary
            The dictionary of settings with a dictionary of region aggregations
        """
        self.pudl_engine = pudl_engine
        self.pudl_out = pudl_out
        self.pg_engine = pg_engine
        self.settings = settings
        self.current_gens = current_gens
        self.sort_gens = sort_gens
        self.model_regions_gdf = load_ipm_shapefile(self.settings)
        self.weighted_unit_hr = None
        self.supplement_with_860m = supplement_with_860m
        self.cluster_builder = build_resource_clusters(
            self.settings.get("RESOURCE_GROUPS")
        )

        if self.current_gens:
            self.data_years = self.settings["data_years"]

            self.gens_860 = load_generator_860_data(self.pudl_engine, self.data_years)
            self.gens_entity = pd.read_sql_table(
                "generators_entity_eia", self.pudl_engine
            )

            bga = self.pudl_out.bga_eia860()
            self.bga = bga.loc[
                bga.report_date.dt.year.isin(self.data_years), :
            ].drop_duplicates(["plant_id_eia", "generator_id"])

            logger.info("Loading map of plants to IPM regions")
            self.plant_region_map = load_plant_region_map(
                self.gens_860,
                self.pudl_engine,
                self.pg_engine,
                self.settings,
                self.model_regions_gdf,
                table=plant_region_map_table,
            )

            self.gen_923 = load_923_gen_fuel_data(
                self.pudl_engine,
                self.pudl_out,
                model_region_map=self.plant_region_map,
                data_years=self.data_years,
            )

            self.eia_860m = load_860m(self.settings)
            self.operating_860m = self.eia_860m["operating"]
            self.planned_860m = self.eia_860m["planned"]
            self.canceled_860m = self.eia_860m["canceled"]
            self.retired_860m = self.eia_860m["retired"]

            # self.ownership = load_ownership_eia860(self.pudl_engine, self.data_years)
            self.plants_860 = load_plants_860(self.pudl_engine, self.data_years)
            # self.utilities_eia = load_utilities_eia(self.pudl_engine)
        else:
            self.existing_resources = pd.DataFrame()
        self.fuel_prices = fetch_fuel_prices(self.settings).pipe(
            modify_fuel_prices,
            self.settings.get("aeo_fuel_region_map"),
            self.settings.get("regional_fuel_adjustments"),
        )
        self.atb_hr = fetch_atb_heat_rates(self.pg_engine, self.settings)
        self.coal_fgd = pd.read_csv(DATA_PATHS["coal_fgd"])

    def fill_na_heat_rates(self, s):
        """Fill null heat rate values with the median of the series. Not many null
        values are expected.

        Parameters
        ----------
        df : DataFrame
            Must contain the column 'heat_rate_mmbtu_mwh'

        Returns
        -------
        Dataframe
            Same as input but with any null values replaced by the median.
        """
        if s.isnull().any():
            median_hr = s.median()
            return s.fillna(median_hr)
        else:
            return s
        # median_hr = df["heat_rate_mmbtu_mwh"].median()
        # df["heat_rate_mmbtu_mwh"].fillna(median_hr, inplace=True)

        # return df

    def create_demand_response_gen_rows(self):
        """Create rows for demand response/management resources to include in the
        generators file.

        Returns
        -------
        DataFrame
            One row for each region/DSM resource with values in all columns filled.
        """
        year = self.settings["model_year"]
        df_list = []
        self.demand_response_profiles = {}

        if not self.settings.get("flexible_demand_resources"):
            logger.warning(
                "A demand response file is included in extra inputs but the parameter "
                "`flexible_demand_resources` is not in the settings file. No demand "
                "response resources will be included with the generators."
            )
            return pd.DataFrame()
        if year not in self.settings["flexible_demand_resources"].keys():
            logger.warning(
                f"The model year {year} is not included in your 'flexible_demand_resources' "
                "parameter. No flexible demand resources will be included for this "
                "planning period."
            )
        if self.settings["flexible_demand_resources"][year] is None:
            logger.warning(
                f"Your 'flexible_demand_resources' settings parameter has the value 'None' "
                "for this planning period. No flexible demand resources will be included."
            )
        for resource, parameters in (
            self.settings["flexible_demand_resources"].get(year, {}).items()
        ):
            _df = pd.DataFrame(
                index=self.settings["model_regions"],
                columns=list(self.settings["generator_columns"]) + ["profile"],
            )
            _df = _df.drop(columns="Resource")
            _df["technology"] = resource
            _df["region"] = self.settings["model_regions"]
            if self.settings.get("demand_response_fn"):
                dr_path = (
                    Path.cwd()
                    / self.settings["input_folder"]
                    / self.settings["demand_response_fn"]
                )
                dr_profile = make_demand_response_profiles(
                    dr_path,
                    resource,
                    self.settings["model_year"],
                    self.settings["demand_response"],
                )
            elif self.settings.get("electrification_stock_fn"):
                if not self.settings.get("electrification_scenario"):
                    logger.warning(
                        "You have provided a parameter value for 'electrification_stock_fn' "
                        "but not 'electrification_scenario'. No flexible demand resources "
                        "can be included without a valid electrification scenario."
                    )
                    pass
                keep_regions, region_agg_map = regions_to_keep(
                    self.settings["model_regions"],
                    self.settings.get("region_aggregations", {}) or {},
                )
                dr_profile = load_demand_response_efs_profile(
                    resource,
                    self.settings.get("electrification_stock_fn"),
                    self.settings["model_year"],
                    self.settings.get("electrification_scenario"),
                    keep_regions,
                    self.settings.get("region_aggregations", {}) or {},
                    self.settings.get("EFS_DATA"),
                )
            self.demand_response_profiles[resource] = dr_profile
            # Add hourly profile to demand response rows
            dr_cf = dr_profile / dr_profile.max()
            dr_regions = [r for r in dr_cf.columns if r in _df.index]
            _df = _df.loc[dr_regions, :]
            _df["profile"] = list(dr_cf[dr_regions].values.T)

            dr_capacity = demand_response_resource_capacity(
                dr_profile, resource, self.settings
            )

            # This is to solve a bug with only one region. Need to come back and solve
            # in a better fashion.
            if len(dr_capacity) > 1:
                dr_capacity_scenario = dr_capacity.squeeze()
            else:
                dr_capacity_scenario = dr_capacity
            _df["Existing_Cap_MW"] = _df["region"].map(dr_capacity_scenario)

            if not parameters.get("parameter_values"):
                logger.warning(
                    "No model parameter values are provided in the settings file for "
                    f"the flexible demand resource '{resource}'. If another resource"
                    " has values under "
                    "`flexible_demand_resource.<year>.<resource_type>.parameter_values`, "
                    f"those columns will have a value of 0 for '{resource}'."
                )
            for col, value in parameters.get("parameter_values", {}).items():
                _df[col] = value

            df_list.append(_df)

        dr_rows = pd.concat(df_list)
        dr_rows["New_Build"] = -1
        dr_rows["Fuel"] = "None"
        dr_rows["cluster"] = 1
        dr_rows = dr_rows.fillna(0)

        return dr_rows

    def create_region_technology_clusters(self, return_retirement_capacity=False):
        """
        Calculation of average unit characteristics within a technology cluster
        (capacity, minimum load, heat rate) and the number of units in the cluster.

        Parameters
        ----------
        plant_region_map_table : str, optional
            Name of the table with region names for each plant, by default
            "plant_region_map_epaipm"
        settings_agg_key : str, optional
            Name of the settings dictionary key with regional aggregations, by default
            "region_aggregations"
        return_retirement_capacity : bool, optional
            If retired generators should be retured as a second dataframe, by default
            False

        Returns
        -------
        dataframe

        """
        if self.gens_860.technology_description.isna().any():
            self.gens_860 = fill_missing_tech_descriptions(self.gens_860)
        self.gens_860_model = (
            self.gens_860.pipe(
                supplement_generator_860_data,
                self.gens_entity,
                self.bga,
                self.plant_region_map,
                self.settings,
            )
            .pipe(remove_canceled_860m, self.canceled_860m)
            .pipe(remove_retired_860m, self.retired_860m)
            .pipe(label_retirement_year, self.settings, add_additional_retirements=True)
            .pipe(label_small_hydro, self.settings, by=["plant_id_eia"])
            .pipe(
                group_technologies,
                self.settings.get("group_technologies"),
                self.settings.get("tech_groups", {}) or {},
                self.settings.get("regional_no_grouping", {}) or {},
            )
            .pipe(update_operating_date_860m, self.operating_860m.copy())
        )
        self.gens_860_model = self.gens_860_model.pipe(
            modify_cc_prime_mover_code, self.gens_860_model
        )
        self.gens_860_model.drop_duplicates(inplace=True)

        self.annual_gen_hr_923 = (
            self.gen_923.pipe(modify_cc_prime_mover_code, self.gens_860_model)
            .pipe(group_gen_by_year_fuel_primemover)
            .pipe(add_923_heat_rate)
        )

        # Add heat rates to the data we already have from 860
        logger.info("Loading heat rate data for units and generator/fuel combinations")
        self.prime_mover_hr_map = plant_pm_heat_rates(self.annual_gen_hr_923)
        if self.weighted_unit_hr is None:
            self.weighted_unit_hr = unit_generator_heat_rates(
                self.pudl_out, self.data_years
            )
        else:
            logger.info("Using unit heat rates from previous round.")
        self.weighted_unit_hr["unit_id_pg"] = self.weighted_unit_hr[
            "unit_id_pudl"
        ].astype("object")

        # Merge the PUDL calculated heat rate data and set the index for easy
        # mapping using plant/prime mover heat rates from 923
        hr_cols = ["plant_id_eia", "unit_id_pg", "heat_rate_mmbtu_mwh"]
        idx = ["plant_id_eia", "prime_mover_code", "energy_source_code_1"]
        self.units_model = self.gens_860_model.merge(
            self.weighted_unit_hr[hr_cols],
            on=["plant_id_eia", "unit_id_pg"],
            how="left",
        ).set_index(idx)

        logger.info(
            f"Units model technologies are "
            f"{self.units_model.technology_description.unique().tolist()}"
        )
        # print(units_model.head())

        logger.info(
            "Assigning technology/fuel heat rates where unit heat rates are not "
            "available"
        )
        self.units_model.loc[
            self.units_model.heat_rate_mmbtu_mwh.isnull(), "heat_rate_mmbtu_mwh"
        ] = self.units_model.loc[
            self.units_model.heat_rate_mmbtu_mwh.isnull()
        ].index.map(
            self.prime_mover_hr_map
        )

        self.units_model.loc[
            self.units_model.heat_rate_mmbtu_mwh > 35, "heat_rate_mmbtu_mwh"
        ] = self.units_model.loc[self.units_model.heat_rate_mmbtu_mwh > 35].index.map(
            self.prime_mover_hr_map
        )

        # Set heat rates < 6.36 (ATB NGCC) or > 35 mmbtu/MWh to nan. Don't change heat rates of 0,
        # which is when there is positive generation and no fuel use (pumped storage)
        self.units_model.loc[
            (
                (self.units_model.heat_rate_mmbtu_mwh < 6.36)
                & ~(
                    self.units_model.technology_description.str.contains(
                        "pumped", case=False
                    )
                )
            )
            | (self.units_model.heat_rate_mmbtu_mwh > 35),
            "heat_rate_mmbtu_mwh",
        ] = np.nan

        # Fill any null heat rate values for each tech
        for tech in self.units_model["technology_description"].unique():
            self.units_model.loc[
                self.units_model.technology_description == tech, "heat_rate_mmbtu_mwh"
            ] = self.fill_na_heat_rates(
                self.units_model.loc[
                    self.units_model.technology_description == tech,
                    "heat_rate_mmbtu_mwh",
                ]
            )
        # assert (
        #     self.units_model["heat_rate_mmbtu_mwh"].isnull().any() is False
        # ), "There are still some null heat rate values"
        # from IPython import embed

        # embed()
        logger.info(
            f"Units model technologies are "
            f"{self.units_model.technology_description.unique().tolist()}"
        )
        if self.supplement_with_860m:
            logger.info(
                f"Before adding proposed generators, {len(self.units_model)} units with "
                f"{self.units_model[self.settings['capacity_col']].sum()} MW capacity"
            )
            self.proposed_gens = import_proposed_generators(
                planned=self.planned_860m,
                settings=self.settings,
                model_regions_gdf=self.model_regions_gdf,
            )
            self.new_860m_gens = import_new_generators(
                operating_860m=self.operating_860m,
                gens_860=self.gens_860_model,
                settings=self.settings,
                model_regions_gdf=self.model_regions_gdf,
            )
            # Add new/proposed generators to plant_region_map
            self.plant_region_map = pd.concat(
                [
                    self.plant_region_map,
                    self.proposed_gens.reset_index()[
                        ["plant_id_eia", "model_region"]
                    ].drop_duplicates(),
                    self.new_860m_gens.reset_index()[
                        ["plant_id_eia", "model_region"]
                    ].drop_duplicates(),
                ]
            ).drop_duplicates(subset=["plant_id_eia", "model_region"])
            # embed()
            logger.info(
                f"Proposed gen technologies are "
                f"{self.proposed_gens.technology_description.unique().tolist()}"
            )
            logger.info(
                f"{self.proposed_gens[self.settings['capacity_col']].sum()} MW proposed"
            )
            self.units_model = pd.concat(
                [self.proposed_gens, self.units_model, self.new_860m_gens], sort=False
            )

        # Create a pudl unit id based on plant and generator id where one doesn't exist.
        # This is used later to match the cluster numbers to plants
        self.units_model.reset_index(inplace=True)
        if self.supplement_with_860m:
            self.units_model = add_860m_storage_mwh(
                self.units_model,
                self.new_860m_gens.reset_index(),
                storage_techs=["Batteries"],
            )
        else:
            self.units_model["capacity_mwh"] = 0
        if self.settings.get("energy_storage_duration"):
            self.units_model = energy_storage_mwh(
                self.units_model,
                self.settings["energy_storage_duration"],
                "technology_description",
                self.settings["capacity_col"],
                "capacity_mwh",
            )
        else:
            logger.warning(
                "\n\n*******************************\n"
                "The parameter 'energy_storage_duration' is not included in your settings "
                "file. This parameter converts MW to energy capacity (MWh) for existing "
                "generators where data is otherwise not available. For example, EIA-860m "
                "generally does not have MWh for battery storage units installed in the "
                "current calendar year. The energy capacity of existing storage clusters "
                "may be incorrect if you do not include this parameter!\n"
                "*******************************\n"
            )
        self.units_model["plant_id_eia"] = self.units_model["plant_id_eia"].astype(
            "Int64"
        )
        self.units_model.loc[self.units_model.unit_id_pg.isnull(), "unit_id_pg"] = (
            self.units_model.loc[
                self.units_model.unit_id_pg.isnull(), "plant_id_eia"
            ].astype(str)
            + "_"
            + self.units_model.loc[
                self.units_model.unit_id_pg.isnull(), "generator_id"
            ].astype(str)
        ).values
        self.units_model.set_index(idx, inplace=True)

        logger.info("Calculating plant O&M costs")
        techs = self.settings["num_clusters"].keys()
        self.units_model = (
            self.units_model.rename(columns={"technology_description": "technology"})
            .query("technology.isin(@techs).values")
            .pipe(
                atb_fixed_var_om_existing,
                self.atb_hr,
                self.settings,
                self.pg_engine,
                self.coal_fgd,
            )
        )

        # logger.info(
        #     f"After adding proposed, units model technologies are "
        #     f"{self.units_model.technology_description.unique().tolist()}"
        # )
        logger.info(
            f"After adding proposed generators, {len(self.units_model)} units with "
            f"{self.units_model[self.settings['capacity_col']].sum()} MW capacity"
        )

        techs = list(self.settings["num_clusters"])

        num_clusters = {}
        for region in self.settings["model_regions"]:
            num_clusters[region] = self.settings["num_clusters"].copy()

        if self.settings.get("alt_num_clusters"):
            for region in self.settings["alt_num_clusters"]:
                for tech, cluster_size in self.settings["alt_num_clusters"][
                    region
                ].items():
                    num_clusters[region][tech] = cluster_size

        region_tech_grouped = self.units_model.loc[
            (self.units_model.technology.isin(techs))
            & ~(self.units_model.retirement_year <= self.settings["model_year"]),
            :,
        ].groupby(["model_region", "technology"])

        self.retired = self.units_model.loc[
            ~(self.units_model.retirement_year > self.settings["model_year"]), :
        ]

        # gens_860 lost the ownership code... refactor this!
        # self.all_gens_860 = load_generator_860_data(self.pudl_engine, self.data_years)
        # Getting weighted ownership for each unit, which will be used below.
        # self.weighted_ownership = weighted_ownership_by_unit(
        #     self.units_model, self.all_gens_860, self.ownership, self.settings
        # )

        # For each group, cluster and calculate the average size/min load/heat rate
        # logger.info("Creating technology clusters by region")
        logger.info("Creating technology clusters by region")
        unit_list = []
        self.cluster_list = []
        alt_cluster_method = self.settings.get("alt_cluster_method") or {}

        for _, df in region_tech_grouped:
            region, tech = _
            grouped = group_units(df, self.settings)

            # This is bad. Should be setting up a dictionary of objects that picks the
            # correct clustering method. Can't keep doing if statements as the number of
            # methods grows. CHANGE LATER.
            if not alt_cluster_method:
                # Allow users to set value as None and not cluster units.
                if num_clusters[region][tech] is None:
                    grouped["cluster"] = np.arange(len(grouped)) + 1
                    unit_list.append(grouped)
                    _df = calc_unit_cluster_values(
                        grouped, self.settings["capacity_col"], tech
                    )
                    _df["region"] = region

                    self.cluster_list.append(_df)
                    continue
                if num_clusters[region][tech] > 0:
                    cluster_cols = [
                        "Fixed_OM_Cost_per_MWyr",
                        # "Var_OM_Cost_per_MWh",
                        # "minimum_load_mw",
                        "heat_rate_mmbtu_mwh",
                    ]
                    if len(grouped) < num_clusters[region][tech]:
                        s = f"""
    *****************************
    The technology {tech} in region {region} has only {len(grouped)} operating units,
    which is less than the {num_clusters[region][tech]} clusters you specified.
    The number of clusters has been set equal to the number of units.
    *****************************
                            """
                        logger.info(s)
                        num_clusters[region][tech] = len(grouped)
                    clusters = cluster.KMeans(
                        n_clusters=num_clusters[region][tech], random_state=6
                    ).fit(
                        preprocessing.StandardScaler().fit_transform(
                            grouped[cluster_cols]
                        )
                    )

                    grouped["cluster"] = (
                        clusters.labels_ + 1
                    )  # Change to 1-index for julia

            else:
                if (
                    region in alt_cluster_method
                    and tech in alt_cluster_method[region]["technology_description"]
                ):
                    grouped = cluster_by_owner(
                        df,
                        self.weighted_ownership,
                        # self.ownership,
                        self.plants_860,
                        region,
                        tech,
                        self.settings,
                    )

                elif num_clusters[region][tech] > 0:
                    clusters = cluster.KMeans(
                        n_clusters=num_clusters[region][tech], random_state=6
                    ).fit(preprocessing.StandardScaler().fit_transform(grouped))

                    grouped["cluster"] = (
                        clusters.labels_ + 1
                    )  # Change to 1-index for julia

            # Saving individual unit data for later analysis (if needed)
            unit_list.append(grouped)

            # Don't add technologies with specified 0 clusters
            if num_clusters[region][tech] != 0:
                _df = calc_unit_cluster_values(
                    grouped, self.settings["capacity_col"], tech
                )
                _df["region"] = region
                _df["plant_id_eia"] = (
                    grouped.reset_index().groupby("cluster")["plant_id_eia"].apply(list)
                )
                _df["unit_id_pg"] = (
                    grouped.reset_index().groupby("cluster")["unit_id_pg"].apply(list)
                )

                self.cluster_list.append(_df)

        # Save some data about individual units for easy access
        self.all_units = pd.concat(unit_list, sort=False)
        self.all_units = pd.merge(
            self.units_model.reset_index(),
            self.all_units,
            on=["plant_id_eia", "unit_id_pg"],
            how="left",
        ).merge(
            self.plants_860[["plant_id_eia", "utility_id_eia"]],
            on=["plant_id_eia"],
            how="left",
        )

        logger.info("Finalizing generation clusters")
        self.results = pd.concat(self.cluster_list)
        logger.info(
            f"Results technologies are {self.results.technology.unique().tolist()}"
        )

        # if self.settings.get("region_wind_pv_cap_fn"):
        #     from powergenome.external_data import overwrite_wind_pv_capacity

        #     logger.info("Setting existing wind/pv using external file")
        #     self.results = overwrite_wind_pv_capacity(self.results, self.settings)
        for col in ["region", "technology", "cluster"]:
            if col in self.results.index.names:
                self.results = self.results.reset_index()

        self.results = self.results.set_index(["region", "technology", "cluster"])
        self.results.rename(
            columns={
                self.settings["capacity_col"]: "Cap_Size",
                "heat_rate_mmbtu_mwh": "Heat_Rate_MMBTU_per_MWh",
            },
            inplace=True,
        )

        # Calculate average capacity factors
        if type(self.settings.get("capacity_factor_techs")) is list:
            self.capacity_factors = gentype_region_capacity_factor(
                self.pudl_engine,
                self.units_model[["plant_id_eia", "model_region"]],
                self.settings,
            )

            self.results = pd.merge(
                self.results.reset_index(),
                self.capacity_factors[["region", "technology", "capacity_factor"]],
                on=["region", "technology"],
                how="left",
            )

            if self.settings.get("derate_capacity"):
                derate_techs = self.settings["derate_techs"]
                self.results.loc[:, "unmodified_cap_size"] = self.results.loc[
                    :, "Cap_Size"
                ].copy()
                self.results.loc[
                    self.results["technology"].isin(derate_techs), "Cap_Size"
                ] = (
                    self.results.loc[
                        self.results["technology"].isin(derate_techs),
                        "unmodified_cap_size",
                    ]
                    * self.results.loc[
                        self.results["technology"].isin(derate_techs), "capacity_factor"
                    ]
                )

        # Round Cap_size to prevent GenX error.
        self.results = self.results.round(3)
        self.results["Cap_Size"] = self.results["Cap_Size"]
        self.results["Existing_Cap_MW"] = self.results.Cap_Size * self.results.num_units
        if self.settings.get("derate_capacity"):
            self.results["unmodified_existing_cap_mw"] = (
                self.results["unmodified_cap_size"] * self.results["num_units"]
            )

        if self.settings.get("region_wind_pv_cap_fn"):
            from powergenome.external_data import overwrite_wind_pv_capacity

            logger.info("Setting existing wind/pv using external file")
            self.results = overwrite_wind_pv_capacity(self.results, self.settings)

        if self.settings.get("dg_as_resource"):
            logger.info(
                "\n **************** \nDistributed generation is being added as generating"
                " resources. The capacity of DG in each region is increased by "
                f"{self.settings.get('avg_distribution_loss', 0):%} to account for no "
                "distribution losses.\n"
            )
            self.results = add_dg_resources(
                self.pg_engine, self.settings, self.results.reset_index()
            )
        else:
            self.results["profile"] = None

        # Add fixed/variable O&M based on NREL atb
        self.results = (
            self.results.reset_index()
            # .pipe(
            #     atb_fixed_var_om_existing, self.atb_costs, self.atb_hr, self.settings
            # )
            # .pipe(atb_new_generators, self.atb_costs, self.atb_hr, self.settings)
            .pipe(startup_fuel, self.settings)
            .pipe(add_fuel_labels, self.fuel_prices, self.settings)
            .pipe(startup_nonfuel_costs, self.settings)
            .pipe(add_genx_model_tags, self.settings)
        )

        if self.sort_gens:
            logger.info("Sorting new resources alphabetically.")
            self.results = self.results.sort_values(["region", "technology"])

        # self.results = self.results.rename(columns={"technology": "Resource"})
        self.results["Resource"] = (
            self.results["region"]
            + "_"
            + snake_case_col(self.results["technology"])
            + "_"
            + self.results["cluster"].astype(str)
        )

        # Add variable resource profiles
        self.results = self.results.reset_index(drop=True)
        for i, row in enumerate(self.results.itertuples()):
            params = map_eia_technology(row.technology)
            if not params:
                # EIA technology not supported
                continue
            params.update({"existing": True})
            groups = self.cluster_builder.find_groups(**params)
            if not groups:
                # No matching resource groups
                continue
            if len(groups) > 1:
                # Multiple matching resource groups
                raise ValueError(
                    f"Multiple existing resource groups match EIA technology"
                    + row.technology
                )
            group = groups[0]
            if group.profiles is None:
                # Resource group has no profiles
                continue
            if row.region in (self.settings.get("region_aggregations", {}) or {}):
                ipm_regions = self.settings.get("region_aggregations", {})[row.region]
            else:
                ipm_regions = [row.region]
            metadata = group.metadata.read()
            if not metadata["ipm_region"].isin(ipm_regions).any():
                # Resource group has no resources in selected IPM regions
                continue
            clusters = group.get_clusters(
                ipm_regions=ipm_regions,
                max_clusters=1,
                utc_offset=self.settings.get("utc_offset", 0),
            )
            self.results["profile"][i] = clusters["profile"][0]

        self.results = rename_gen_cols(self.results)

        return self.results

    def create_new_generators(self):
        self.offshore_spur_costs = fetch_atb_offshore_spur_costs(
            self.pg_engine, self.settings
        )
        self.atb_costs = fetch_atb_costs(
            self.pg_engine, self.settings, self.offshore_spur_costs
        )

        self.new_generators = atb_new_generators(
            self.atb_costs, self.atb_hr, self.settings, self.cluster_builder
        )

        if not self.new_generators.empty:
            self.new_generators = (
                self.new_generators.pipe(startup_fuel, self.settings)
                .pipe(add_fuel_labels, self.fuel_prices, self.settings)
                .pipe(startup_nonfuel_costs, self.settings)
            )

            if self.sort_gens:
                logger.info("Sorting new resources alphabetically.")
                self.new_generators = self.new_generators.sort_values(
                    ["region", "technology"]
                )

            if self.settings.get("capacity_limit_spur_fn"):
                self.new_generators = self.new_generators.pipe(
                    add_resource_max_cap_spur, self.settings
                )
            else:
                logger.warning("No settings parameter for max capacity/spur file")
            self.new_generators = self.new_generators.pipe(
                calculate_transmission_inv_cost, self.settings, self.offshore_spur_costs
            ).pipe(add_transmission_inv_cost, self.settings)

        if self.settings.get("demand_response_fn") or self.settings.get(
            "electrification_stock_fn"
        ):
            dr_rows = self.create_demand_response_gen_rows()
            self.new_generators = pd.concat(
                [self.new_generators, dr_rows], sort=False, ignore_index=True
            )
        self.new_generators = add_genx_model_tags(self.new_generators, self.settings)
        if "cluster" not in self.new_generators.columns:
            self.new_generators["cluster"] = 1
        self.new_generators["Resource"] = (
            self.new_generators["region"]
            + "_"
            + snake_case_col(self.new_generators["technology"])
            + "_"
            + self.new_generators["cluster"].astype(str)
        )

        return self.new_generators

    def create_all_generators(self):
        if self.current_gens:
            self.existing_resources = self.create_region_technology_clusters()

        self.new_resources = self.create_new_generators()

        self.all_resources = pd.concat(
            [self.existing_resources, self.new_resources], ignore_index=True, sort=False
        )

        self.all_resources = self.all_resources.round(3)
        self.all_resources["Cap_Size"] = self.all_resources["Cap_Size"]
        self.all_resources["Heat_Rate_MMBTU_per_MWh"] = self.all_resources[
            "Heat_Rate_MMBTU_per_MWh"
        ]

        self.all_resources = self.all_resources.reset_index(drop=True)
        self.all_resources["variable_CF"] = 0.0
        for i, p in enumerate(self.all_resources["profile"]):
            if isinstance(p, (collections.abc.Sequence, np.ndarray)):
                self.all_resources.loc[i, "variable_CF"] = np.mean(p)

        # Set Min_Power of wind/solar to 0
        if "VRE" in self.all_resources.columns:
            self.all_resources.loc[self.all_resources["VRE"] == 1, "Min_Power"] = 0

        self.all_resources["R_ID"] = np.arange(len(self.all_resources)) + 1

        if self.current_gens:
            logger.info(
                f"Capacity of {self.all_resources['Existing_Cap_MW'].sum()} MW in final clusters"
            )

        return self.all_resources


# powergenome/load_construction.py

from typing import Dict, List, Union
import numpy as np
import pandas as pd
import os
import logging
from functools import lru_cache

from pathlib import Path


from powergenome.util import (
    deep_freeze_args,
    find_region_col,
    snake_case_col,
)
from powergenome.params import DATA_PATHS, SETTINGS

logger = logging.getLogger(__name__)


us_state_abbrev = {
    "Alabama": "AL",
    "Alaska": "AK",
    "American Samoa": "AS",
    "Arizona": "AZ",
    "Arkansas": "AR",
    "California": "CA",
    "Colorado": "CO",
    "Connecticut": "CT",
    "Delaware": "DE",
    "District Of Columbia": "DC",
    "District of Columbia": "DC",
    "Florida": "FL",
    "Georgia": "GA",
    "Guam": "GU",
    "Hawaii": "HI",
    "Idaho": "ID",
    "Illinois": "IL",
    "Indiana": "IN",
    "Iowa": "IA",
    "Kansas": "KS",
    "Kentucky": "KY",
    "Louisiana": "LA",
    "Maine": "ME",
    "Maryland": "MD",
    "Massachusetts": "MA",
    "Michigan": "MI",
    "Minnesota": "MN",
    "Mississippi": "MS",
    "Missouri": "MO",
    "Montana": "MT",
    "Nebraska": "NE",
    "Nevada": "NV",
    "New Hampshire": "NH",
    "New Jersey": "NJ",
    "New Mexico": "NM",
    "New York": "NY",
    "North Carolina": "NC",
    "North Dakota": "ND",
    "Northern Mariana Islands": "MP",
    "Ohio": "OH",
    "Oklahoma": "OK",
    "Oregon": "OR",
    "Pennsylvania": "PA",
    "Puerto Rico": "PR",
    "Rhode Island": "RI",
    "South Carolina": "SC",
    "South Dakota": "SD",
    "Tennessee": "TN",
    "Texas": "TX",
    "Utah": "UT",
    "Vermont": "VT",
    "Virgin Islands": "VI",
    "Virginia": "VA",
    "Washington": "WA",
    "West Virginia": "WV",
    "Wisconsin": "WI",
    "Wyoming": "WY",
}


def load_region_pop_frac(
    path_in: Path,
    fn: str = "ipm_state_pop_weight.parquet",
) -> pd.DataFrame:
    # TODO #178 finalize state pop weight file and filename
    # read in state proportions
    # how much state load should be distributed to GenXRegion

    if (path_in / fn).suffix == ".csv":
        pop = pd.read_csv(path_in / fn)
    elif (path_in / fn).suffix == ".parquet":
        pop = pd.read_parquet(path_in / fn)
    pop["state"] = pop["state"].map(us_state_abbrev)
    context = "Loading region population fraction file for EFS load construction."
    region_col = find_region_col(pop.columns, context)
    pop = pop.rename(columns={region_col: "region"})
    pop_cols = ["region", "state", "state_prop"]
    return pop[pop_cols]


running_sectors = {
    "res_space_heat_cool": ("Residential", "space heating and cooling"),
    "res_water_heat": ("Residential", "water heating"),
    "comm_space_heat_cool": ("Commercial", "space heating and cooling"),
    "comm_water_heat": ("Commercial", "water heating"),
    "trans_light_duty": ("Transportation", "light-duty vehicles"),
    "trans_medium_duty": ("Transportation", "medium-duty trucks"),
    "trans_heavy_duty": ("Transportation", "heavy-duty trucks"),
    "trans_bus": ("Transportation", "transit buses"),
}


def CreateOutputFolder(case_folder):
    path = case_folder / "extra_outputs"
    if not os.path.exists(path):
        os.makedirs(path)


def create_subsector_ts(
    sector: str,
    subsector: str,
    year: int,
    scenario_stock: pd.DataFrame,
    path_in: Path = None,
) -> pd.DataFrame:
    if not path_in:
        try:
            path_in = Path(SETTINGS["EFS_DATA"])
        except TypeError:
            logger.warning("The variable 'EFS_DATA' is not included in your .env file.")

    timeseries = pd.read_parquet(
        path_in / f"{sector}_{subsector}_Incremental_Factor.parquet"
    )
    timeseries.columns = snake_case_col(timeseries.columns)
    timeseries = timeseries.rename(columns={"localhourid": "time_index"})
    ts_cols = ["state", "year", "time_index", "unit", "factor_type1", "factor_type2"]
    timeseries = timeseries.loc[:, ts_cols]
    stock_temp = scenario_stock.loc[
        (scenario_stock["sector"] == sector)
        & (scenario_stock["subsector"] == subsector),
        ["state", "year", "agg_stock_type1", "agg_stock_type2"],
    ]

    factor_years = timeseries["year"].unique()
    if not year in factor_years:
        diff = np.array(factor_years - year)
        index = diff[np.where(diff <= 0)].argmax()
        year_approx = factor_years[index]
        timeseries = timeseries.loc[timeseries["year"] == year_approx, :]
        timeseries["year"] = year
        logger.warning(
            f"No incremental factor available for {sector}/{subsector} "
            f"in year {year}: using factors from year {year_approx} instead."
        )

    timeseries = pd.merge(timeseries, stock_temp, on=["state", "year"], validate="m:1")
    timeseries = timeseries.assign(
        load_mw=timeseries["agg_stock_type1"] * timeseries["factor_type1"]
        + timeseries["agg_stock_type2"] * timeseries["factor_type2"]
    )
    return timeseries[["state", "time_index", "load_mw"]].dropna()


def state_demand_to_region(
    df: pd.DataFrame, pop: pd.DataFrame, col_name: str
) -> pd.DataFrame:
    temp = pd.merge(df, pop, on=["state"], how="left")
    temp["load_mw"] *= temp["state_prop"]
    temp = temp.groupby(["time_index", "region"], as_index=False)["load_mw"].sum()
    return temp


@deep_freeze_args
@lru_cache()
def electrification_profiles(
    stock_fn: str,
    year: int,
    elec_scenario: str,
    regions: List[str],
    path_in: Path = None,
) -> pd.DataFrame:
    """Create demand profiles for potentially flexible resources that will be electrified.

    Parameters
    ----------
    stock_fn : str
        Name of the data file with stock values for each year.
    year : int
        Planning period or model year. Used to select stock values for flexible resources
        and their demand profiles.
    elec_scenario : str
        Name of a scenario from the stock data file.
    regions : List[str]
        All of the base regions that will have profiles created.
    path_in : Path, optional
        Folder where stock and incremental factor (profile) data are located, by default
        None.

    Returns
    -------
    pd.DataFrame
        Columns include "time_index", "region", "load_mw", and "resource".

    Raises
    ------
    KeyError
        No "EFS_DATA" key is included in the .env file.
    FileNotFoundError
        The "EFS_DATA" or "path_in" folder does not exist.
    ValueError
        The "electrificication_scenario" parameter does not match scenarios in the stock
        file.
    """
    if not path_in:
        try:
            path_in = Path(SETTINGS["EFS_DATA"])
        except TypeError:
            raise KeyError("The variable 'EFS_DATA' is not included in your .env file.")
    if not path_in.is_dir():
        raise FileNotFoundError(
            f"The folder with EFS/flexible demand data ({str(path_in)}) was not found."
        )

    pop_files = path_in.glob("*pop_weight*")
    newest_pop_file = max(pop_files, key=os.path.getctime)
    pop = load_region_pop_frac(path_in=path_in, fn=newest_pop_file.name)
    pop = pop.loc[pop["region"].isin(regions), :]
    states = pop["state"].unique()
    scenario_stock = pd.read_parquet(path_in / stock_fn)
    scenario_stock.columns = snake_case_col(scenario_stock.columns)
    valid_scenarios = list(scenario_stock.scenario.unique())
    if elec_scenario not in valid_scenarios:
        raise ValueError(
            f"Your 'electrification_scenario' parameter value '{elec_scenario}' was not "
            f"found in the 'electrification_stock_fn' file '{stock_fn}'. Valid scenarios "
            f"are: {valid_scenarios}."
        )
    scenario_stock = scenario_stock.loc[
        (scenario_stock["year"] == year)
        & (scenario_stock["scenario"] == elec_scenario)
        & (scenario_stock["state"].isin(states)),
        :,
    ]

    subsector_ts_dfs = []
    for name, (sector, subsector) in running_sectors.items():
        df = create_subsector_ts(sector, subsector, year, scenario_stock, path_in).pipe(
            state_demand_to_region, pop, name
        )

        df["resource"] = name
        subsector_ts_dfs.append(df)

    elec_profiles = pd.concat(subsector_ts_dfs)

    return elec_profiles


# powergenome/load_profiles.py

"""
Hourly demand profiles
"""

import logging
from inspect import signature
from pathlib import Path
from typing import Dict, List, Tuple

import numpy as np
import pandas as pd
import sqlalchemy as sa

from powergenome.load_construction import electrification_profiles
from powergenome.eia_opendata import get_aeo_load
from powergenome.external_data import make_demand_response_profiles
from powergenome.util import (
    find_region_col,
    map_agg_region_names,
    regions_to_keep,
    reverse_dict_of_lists,
    remove_feb_29,
)

logger = logging.getLogger(__name__)


def filter_load_by_region(load_source):  # "decorator factory"
    """If regional load options are given, return the columns listed in
    settings["regional_load_source"][load_source].

    If settings["regional_load_source"] exists and settings["regional_load_source"][load_source]
    is null, return None.

    If settings["regional_load_source"] DNE, return the load profile if the load_source is EFS,
    else return None. This makes EFS the default load type/source.
    """

    def decorator(func):
        def wrapper(*args, **kwargs):
            ## retrieve settings:
            # if kwarg:
            settings = kwargs.get("settings", None)
            # if arg:
            if settings is None:
                settings_arg_position = list(signature(func).parameters).index(
                    "settings"
                )
                settings = args[settings_arg_position]

            regional_load_sources = settings.get("regional_load_source")

            if regional_load_sources is not None:
                regions = None
                if load_source == regional_load_sources:
                    # if only one load profile sources are specified, use for all regions
                    regions = settings.get("model_regions")
                elif (
                    isinstance(regional_load_sources, dict)
                    and load_source in regional_load_sources.keys()
                ):
                    # if multiple load profiles sources are specified, find the proper regions
                    regions = regional_load_sources.get(load_source)

                # pd.reindex will return the entire DataFrame if regions=None,
                # We want the opposite; return None if regions = None
                if regions is not None:
                    load_profile = func(*args, **kwargs)
                    load_profile = load_profile.reindex(columns=regions)
                else:
                    load_profile = None
            else:
                load_profile = None
                if load_source == "EFS":
                    load_profile = func(*args, **kwargs)

            return load_profile

        return wrapper

    return decorator


def make_load_curves(
    pg_engine: sa.engine.base.Engine,
    settings: dict,
    pg_table: str = "load_curves_nrel_efs",
) -> pd.DataFrame:
    """Read base load profiles from database and grow the load to a future year.

    Parameters
    ----------
    pg_engine : sa.engine.base.Engine
        Engine to connect to a PowerGenome database
    settings : dict
        User parameter settings. Required keys are "model_regions", "future_load_region_map",
        and "historical_load_region_maps" (if load data represent a year before
        2019). Optional keys include "region_aggregations", "electrification_stock_fn",
        "electrification_scenario", and "alt_growth_rate".
    pg_table : str, optional
        Name of the database table with load profiles, by default "load_curves_nrel_efs"

    Returns
    -------
    pd.DataFrame
        Wide dataframe of load profiles for each model region. The only change from
        base year data in the database is load growth.

    Raises
    ------
    KeyError
        No table in the database with the supplied name.
    """
    # IPM regions to keep. Regions not in this list will be dropped from the
    # dataframe
    keep_regions, region_agg_map = regions_to_keep(
        settings["model_regions"], settings.get("region_aggregations")
    )

    # I'd rather use a sql query and only pull the regions of interest but
    # sqlalchemy doesn't allow table names to be parameterized.
    logger.info("Loading load curves from PUDL")
    inst = sa.inspect(pg_engine)
    if not inst.has_table(pg_table):
        raise KeyError(
            f"There is no load curves table with the name {pg_table} in the 'PG_DB' "
            "database specified in your .env file."
        )
    table_cols = [c["name"] for c in inst.get_columns(pg_table)]
    context = f"Load curves table ({pg_table} in database {pg_engine}."
    region_col = find_region_col(table_cols, context)
    if "sector" in table_cols or "subsector" in table_cols:
        if settings.get("electrification_stock_fn") and settings.get(
            "electrification_scenario"
        ):
            # This is a default list of sector/subsectors that are considered "base" demand
            # and are not affected by stock levels of electric technologies (e.g. EVs and heat pumps)
            # NOTE: This should be parameratized so it can be changed by the user, especially
            # if load data is from a source other than NREL EFS
            base_sector_subsectors = [
                ("commercial", "other"),
                ("residential", "other"),
                ("residential", "clothes and dish washing/drying"),
                ("industrial", "machine drives"),
                ("industrial", "process heat"),
                ("industrial", "other"),
            ]
            s = f"""
                    SELECT year, {region_col} as region, time_index, sector, sum(load_mw) as load_mw
                    FROM {pg_table}
                    WHERE {region_col} in ({','.join(['?']*len(keep_regions))})
                    AND
                    ({' OR '.join(["(sector=? and subsector=?)"]*len(base_sector_subsectors))})
                    GROUP BY year, region, sector, time_index
                    """
            params = keep_regions + [
                item for sublist in base_sector_subsectors for item in sublist
            ]
            load_curves = pd.read_sql_query(sql=s, con=pg_engine, params=params)
        else:
            s = f"""
                    SELECT year, {region_col} as region, time_index, sector, sum(load_mw) as load_mw
                    FROM {pg_table}
                    WHERE {region_col} in ({','.join(['?']*len(keep_regions))})
                    GROUP BY year, region, sector, time_index
                    """
            params = keep_regions
            load_curves = pd.read_sql_query(sql=s, con=pg_engine, params=params)
    else:
        # With no sector or subsector columns, assume that table has total load in each hour
        s = f"""
            SELECT year, {region_col} as region, time_index, load_mw
            FROM {pg_table}
            WHERE {region_col} in ({','.join(['?']*len(keep_regions))})
            """
        params = keep_regions
        load_curves = pd.read_sql_query(sql=s, con=pg_engine, params=params)

    # Increase demand to account for load growth
    load_curves = add_load_growth(load_curves, settings)

    load_curves.loc[
        load_curves.region.isin(region_agg_map), "region"
    ] = load_curves.region.map(region_agg_map)

    logger.info("Aggregating load curves in grouped regions")
    load_curves_agg = load_curves.groupby(["region", "time_index"])["load_mw"].sum()

    lc_wide = load_curves_agg.unstack(level=0)
    if lc_wide.columns.nlevels > 1:
        lc_wide.columns = lc_wide.columns.droplevel()

    if len(lc_wide) == 8784:
        lc_wide = remove_feb_29(lc_wide)

    # Shift load from UTC
    for col in lc_wide:
        lc_wide[col] = np.roll(lc_wide[col].values, settings.get("utc_offset", 0))

    lc_wide.index.name = "time_index"
    if lc_wide.index.min() == 0:
        lc_wide.index = lc_wide.index + 1

    return lc_wide


def add_load_growth(load_curves: pd.DataFrame, settings: dict) -> pd.DataFrame:
    """Multiply hourly load profiles by AEO or user growth factors.

    If the base data year is from more than one year before the AEO data year then load
    is first grown to that point, then the AEO data year is used to calculate load growth
    to the model planning year.

    Parameters
    ----------
    load_curves : pd.DataFrame
        Tidy dataframe of load curves with columns "region", "load_mw", and optionally
        "sector".
    settings : dict
        User settings parameters. Should include "historical_load_region_map",
        "future_load_region_map", and "model_year". Optional parameters include
        "aeo_sector_map" (mapping load sectors to AEO API sector names), and
        "alt_growth_rate" (either single growth rates for each region or sector-level
        growth rates within each region, where the sector names match those in the load
        profile).

    Returns
    -------
    pd.DataFrame
        Modified version of input dataframe to account for load growth from base year
        to model planning year.
    """
    keep_regions, region_agg_map = regions_to_keep(
        settings["model_regions"], settings.get("region_aggregations")
    )
    hist_region_map = reverse_dict_of_lists(settings["historical_load_region_map"])
    future_region_map = reverse_dict_of_lists(settings["future_load_region_map"])
    aeo_sector_map = settings.get("aeo_sector_map")
    if not settings.get("aeo_sector_map"):
        aeo_sector_map = {
            "commercial": "COMM",
            "industrial": "IDAL",
            "residential": "RESD",
            "transportation": "TRN",
        }
    if "sector" in load_curves.columns:
        load_sectors = set(load_curves.sector.unique())
        aeo_sectors = set(aeo_sector_map)
        if not all([s in aeo_sectors for s in load_sectors]):
            missing_sectors = list(load_sectors - aeo_sectors)
            logger.warning(
                "*********************\n"
                f"The load sectors {missing_sectors} are in your load data but are not "
                "mapped to EIA AEO sectors. The hourly values for these sectors will not "
                "be changed unless you added a growth rate for this sector to all regions "
                "in the settings parameter 'alt_growth_rate'."
                "*********************\n"
            )

    outer_list = []
    for year, df in load_curves.groupby("year"):
        growth_scenario = settings.get("growth_scenario", "REF2020")
        load_aeo_year = settings.get("load_eia_aeo_year") or settings.get(
            "eia_aeo_year", 2020
        )
        while year < load_aeo_year - 1:
            year, df = grow_historical_load(
                df=df,
                year=year,
                aeo_data_year=load_aeo_year,
                keep_regions=keep_regions,
                hist_region_map=hist_region_map,
                future_region_map=future_region_map,
                aeo_sector_map=aeo_sector_map,
                alt_growth_rate=settings.get("alt_growth_rate"),
            )

        df_list = []
        if "sector" in df.columns:
            for sector, _df in df.groupby("sector"):
                load_growth_dict = {
                    region: get_aeo_load(
                        region=future_region_map[region],
                        aeo_year=load_aeo_year,
                        scenario_series=growth_scenario,
                        sector=aeo_sector_map[sector],
                    ).set_index("year")
                    for region in keep_regions
                }

                load_growth_start_map = {
                    region: _df.loc[year, "demand"]
                    for region, _df in load_growth_dict.items()
                }

                load_growth_end_map = {
                    region: _df.loc[settings["model_year"], "demand"]
                    for region, _df in load_growth_dict.items()
                }

                growth_factor = {
                    region: load_growth_end_map[region] / load_growth_start_map[region]
                    for region in keep_regions
                }

                years_growth = settings["model_year"] - year
                for region, rate in (settings.get("alt_growth_rate") or {}).items():
                    if isinstance(rate, dict):
                        if rate.get(sector):
                            growth_factor[region] = (1 + rate["sector"]) ** years_growth
                        else:
                            raise KeyError(
                                f"You specified a sector specific alt_growth_rate for the "
                                f"region '{region}'. The demand data has a sector {sector}, "
                                f"which you did not specify a rate for. Without a sector "
                                "specific growth rate the demand will not be increased."
                            )
                for region in keep_regions:
                    _df.loc[_df["region"] == region, "load_mw"] *= growth_factor[region]
                df_list.append(_df)
        else:
            load_growth_dict = {
                region: get_aeo_load(
                    region=future_region_map[region],
                    aeo_year=load_aeo_year,
                    scenario_series=growth_scenario,
                ).set_index("year")
                for region in keep_regions
            }

            load_growth_start_map = {
                region: _df.loc[year, "demand"]
                for region, _df in load_growth_dict.items()
            }

            load_growth_end_map = {
                region: _df.loc[settings["model_year"], "demand"]
                for region, _df in load_growth_dict.items()
            }

            growth_factor = {
                region: load_growth_end_map[region] / load_growth_start_map[region]
                for region in keep_regions
            }

            years_growth = settings["model_year"] - year
            for region, rate in (settings.get("alt_growth_rate") or {}).items():
                if isinstance(rate, float):
                    growth_factor[region] = (1 + rate) ** years_growth
            for region in keep_regions:
                df.loc[df["region"] == region, "load_mw"] *= growth_factor[region]
            df_list.append(df)

        annual_load = pd.concat(df_list, ignore_index=True)
        outer_list.append(annual_load)

    load_curves = pd.concat(outer_list, ignore_index=True)

    return load_curves


def grow_historical_load(
    df: pd.DataFrame,
    year: int,
    aeo_data_year: int,
    keep_regions: List[str],
    hist_region_map: Dict[str, str],
    future_region_map: Dict[str, str],
    aeo_sector_map: Dict[str, str] = None,
    alt_growth_rate: Dict[str, float] = None,
) -> Tuple[int, pd.DataFrame]:
    """Grow historical load up to either 2019 or the year before "aeo_data_year".

    If the data year is less than 2019 than use "hist_region_map" to grow load to 2019.
    Otherwise use "future_region_map" to grow load from the data year to the the year
    before "aeo_data_year". This two step process is needed ecause the AEO EMM regions
    changed from AEO2019 to AEO2020.

    The function returns both the modified input dataframe and an updated "year"
    parameter.

    Parameters
    ----------
    df : pd.DataFrame
        Tidy dataframe of hourly load. Should include the columns "region" and "load_mw".
        Can also include the column "sector" with load disaggregated by sector.
    year : int
        Basis year of the load data.
    aeo_data_year : int
        Year of AEO data that will be used to grow load out to future planning periods.
    keep_regions : List[str]
        Regions that are included in the model. Only load from these regions is modified.
    hist_region_map : Dict[str, str]
        A map of load regions to AEO EMM regions from AEO2019 and earlier.
    future_region_map : Dict[str, str]
        A map of load regions to AEO EMM regions from AEO2020 and later.
    aeo_sector_map : Dict[str, str]
        A mapping of sector names from the load data to names used by EIA, by default
        None.
    alt_growth_rate : Dict[str, float], optional
        Alternative growth rates provided by the user, by default None.

    Returns
    -------
    int
        Updated data year
    pd.Dataframe
        Updated load data
    """

    old_aeo_list = []
    if year < 2019:
        region_map = hist_region_map
        start_aeo_year = year + 2
        end_aeo_year = 2019
        end_data_year = 2019
    else:
        region_map = future_region_map
        start_aeo_year = year + 1
        end_aeo_year = aeo_data_year
        end_data_year = end_aeo_year - 1
    if "sector" in df.columns:
        if not aeo_sector_map:
            raise KeyError(
                "The load data provided has the column 'sector' but no mapping of sectors "
                "to AEO sector names was provided."
            )
        for sector, _df in df.groupby("sector"):
            hist_demand_start = {
                region: get_aeo_load(
                    region=region_map[region],
                    aeo_year=start_aeo_year,
                    scenario_series=f"REF{start_aeo_year}",
                    sector=aeo_sector_map[sector],
                )
                .set_index("year")
                .loc[year, "demand"]
                for region in keep_regions
            }

            hist_demand_end = {
                region: get_aeo_load(
                    region=region_map[region],
                    aeo_year=end_aeo_year,
                    scenario_series=f"REF{end_aeo_year}",
                    sector=aeo_sector_map[sector],
                )
                .set_index("year")
                .loc[end_data_year, "demand"]
                for region in keep_regions
            }
            growth_factor = {
                region: hist_demand_end[region] / hist_demand_start[region]
                for region in keep_regions
            }

            years_growth = 1
            for region, rate in (alt_growth_rate or {}).items():
                if isinstance(rate, dict) and rate.get(sector):
                    growth_factor[region] = (1 + rate["sector"]) ** years_growth
            for region in keep_regions:
                _df.loc[_df["region"] == region, "load_mw"] *= growth_factor[region]
            old_aeo_list.append(_df)

    else:
        hist_demand_start = {
            region: get_aeo_load(
                region=region_map[region],
                aeo_year=start_aeo_year,
                scenario_series=f"REF{start_aeo_year}",
            )
            .set_index("year")
            .loc[year, "demand"]
            for region in keep_regions
        }
        hist_demand_end = {
            region: get_aeo_load(
                region=region_map[region],
                aeo_year=end_aeo_year,
                scenario_series=f"REF{end_aeo_year}",
            )
            .set_index("year")
            .loc[end_data_year, "demand"]
            for region in keep_regions
        }
        growth_factor = {
            region: hist_demand_end[region] / hist_demand_start[region]
            for region in keep_regions
        }

        years_growth = 1
        for region, rate in (alt_growth_rate or {}).items():
            if isinstance(rate, float):
                growth_factor[region] = (1 + rate) ** years_growth
        for region in keep_regions:
            df.loc[df["region"] == region, "load_mw"] *= growth_factor[region]
        old_aeo_list.append(df)

    df = pd.concat(old_aeo_list, ignore_index=True)

    year = end_data_year
    return year, df


def add_demand_response_resource_load(load_curves, settings):
    dr_path = Path(settings["input_folder"]) / settings["demand_response_fn"]
    dr_types = list(
        settings["flexible_demand_resources"][settings["model_year"]].keys()
    )

    dr_curves = make_demand_response_profiles(
        dr_path, list(dr_types)[0], settings["model_year"], settings["demand_response"]
    )

    if len(dr_types) > 1:
        for dr in dr_types[1:]:
            _dr_curves = make_demand_response_profiles(
                dr_path, dr, settings["model_year"], settings["demand_response"]
            )
            dr_curves = dr_curves + _dr_curves

    for col in dr_curves.columns:
        try:
            load_curves.loc[:, col] += dr_curves[col].values
        except KeyError:
            pass

    load_curves.index.name = "time_index"
    load_curves.index = load_curves.index + 1

    return load_curves


def subtract_distributed_generation(load_curves, pg_engine, settings):
    dg_profiles = make_distributed_gen_profiles(pg_engine, settings)
    dg_profiles.index = dg_profiles.index + 1

    for col in dg_profiles.columns:
        load_curves.loc[:, col] = load_curves.loc[:, col] - (
            dg_profiles[col].values * 1 + settings["avg_distribution_loss"]
        )

    return load_curves


@filter_load_by_region(load_source="USER")
def load_usr_demand_profiles(settings):
    """Temp function. Loads user demand profiles if the file name is provided, else returns None.
    If only specified regions are to be used (settings["regional_load_source"]["USER"]), then
    reindex to use only those regions. Else, returns all regions in the regional load file.
    """
    logger.info("Loading user supplied load profile.")
    regional_load_fn = settings.get("regional_load_fn")

    if regional_load_fn is not None:
        from powergenome.external_data import make_usr_demand_profiles

        lp_path = settings["input_folder"] / regional_load_fn
        hourly_load_profiles = make_usr_demand_profiles(lp_path, settings)

        if len(hourly_load_profiles) == 8784:
            remove_feb_29(hourly_load_profiles)

        hourly_load_profiles.index.name = "time_index"
        hourly_load_profiles.index = pd.RangeIndex(
            start=1, stop=len(hourly_load_profiles) + 1, step=1
        )

        regional_load_sources = settings.get("regional_load_source")
        if regional_load_sources is not None:
            if regional_load_sources == "USER":
                cols = settings.get("model_regions")
            else:
                cols = regional_load_sources.get("USER")
            if not all([col in hourly_load_profiles.columns for col in cols]):
                raise KeyError(
                    f"One or more of the regions {cols} is not included in your "
                    f"user-supplied load curves file {regional_load_fn}."
                )
            hourly_load_profiles = hourly_load_profiles.reindex(columns=cols)

        return hourly_load_profiles

    else:
        logger.info("User supplied load profile not found.")
        return None


def make_final_load_curves(
    pg_engine: sa.engine.base.Engine,
    settings: dict,
):
    """Create final load profiles from base year including growth, dg, and flexible loads

    Parameters
    ----------
    pg_engine : sa.engine.base.Engine
        Engine to connect to a PowerGenome database
    settings : dict
        User parameter settings. Required keys are "model_regions", "future_load_region_map",
        and "historical_load_region_maps" (if load data represent a year before
        2019). Optional keys include "load_source_table_name", "demand_response_fn",
        "distributed_gen_profiles_fn", "dg_as_resource", "region_aggregations",
        "electrification_stock_fn", "electrification_scenario", and "alt_growth_rate".

    Returns
    -------
    pd.DataFrame
        Wide dataframe with one column of load profiles for each model region

    Raises
    ------
    ValueError
        When all load curves are null.
    """

    logger.info("Loading load curves")
    user_load_curves = load_usr_demand_profiles(settings)

    if user_load_curves is not None and all(
        [r in user_load_curves.columns for r in settings["model_regions"]]
    ):
        load_curves_before_dr = user_load_curves

    else:
        load_sources = settings.get("load_source_table_name")
        if load_sources is None:
            s = """
            *****************************
            Regional load data sources have not been specified. Defaulting to EFS load data.
            Check your settings file, and please specify the preferred source for load data
            (FERC, EFS, USER) either for each region or for the entire system with the setting
            "regional_load_source".
            *****************************
            """
            logger.warning(s)
            load_sources = {"EFS": "load_curves_nrel_efs"}

        # `filter_load_by_region` is a decorator factory that generates a decorator
        # when given the parameter `load_source`. This decorator creates a wrapper
        # for the function `make_load_curves`, which is passed the args from the final
        # parentheses.
        load_curves_before_dr = [
            filter_load_by_region(load_source)(make_load_curves)(
                pg_engine, settings, load_table
            )
            for load_source, load_table in load_sources.items()
        ]
        load_curves_before_dr.append(user_load_curves)
        load_curves_before_dr = [df for df in load_curves_before_dr if df is not None]
        if not all(
            [
                len(load_curves_before_dr[0].index.intersection(df.index))
                == load_curves_before_dr[0].shape[0]
                for df in load_curves_before_dr
                if df is not None
            ]
        ):
            raise ValueError(
                "One or more of your load curve data sources does not have a matching time index."
            )

        try:
            load_curves_before_dr = pd.concat(load_curves_before_dr, axis=1)
        except ValueError:
            raise ValueError("All load curves are null.")

    if settings.get("demand_response_fn"):
        if settings.get("regional_load_includes_demand_response"):
            load_curves_before_dg = load_curves_before_dr
        else:
            load_curves_before_dg = add_demand_response_resource_load(
                load_curves_before_dr, settings
            )
    elif settings.get("electrification_stock_fn") and settings.get(
        "electrification_scenario"
    ):
        load_curves_before_dg = load_curves_before_dr.copy()
        keep_regions, region_agg_map = regions_to_keep(
            settings["model_regions"], settings.get("region_aggregations", {}) or {}
        )

        flex_profiles = electrification_profiles(
            settings.get("electrification_stock_fn"),
            settings["model_year"],
            settings.get("electrification_scenario"),
            keep_regions,
            settings.get("EFS_DATA"),
        )
        flex_profiles = map_agg_region_names(
            flex_profiles, region_agg_map, "region", "model_region"
        )
        for region in load_curves_before_dg.columns:
            region_flex_load = (
                flex_profiles.query("model_region==@region")
                .groupby("time_index")["load_mw"]
                .sum()
            )
            if not region_flex_load.empty:
                load_curves_before_dg[region] += region_flex_load
    else:
        load_curves_before_dg = load_curves_before_dr

    if settings.get("distributed_gen_profiles_fn") and not settings.get(
        "dg_as_resource"
    ):
        final_load_curves = subtract_distributed_generation(
            load_curves_before_dg, pg_engine, settings
        )
    else:
        final_load_curves = load_curves_before_dg

    final_load_curves = final_load_curves.astype(int)

    # change order to match model regions
    model_regions = settings.get("model_regions")
    if not all(r in final_load_curves.columns for r in model_regions):
        missing_regions = set(final_load_curves.columns) - set(model_regions)
        logger.warning(
            "You have supplied regional load in an external file, but the load for some "
            f"regions is missing. The regions {missing_regions} are not included in the file. "
            "The load for these regions will not be included in output files."
        )
    final_load_curves = final_load_curves.reindex(columns=model_regions)

    return final_load_curves


def make_distributed_gen_profiles(pg_engine, settings):
    """Create 8760 annual generation profiles for distributed generation in regions.
    Uses a distribution loss parameter in the settings file when DG generation is
    defined a fraction of delivered load.

    Parameters
    ----------
    dg_profiles_path : path-like
        Where to load the file from
    pg_engine : sqlalchemy.Engine
        A sqlalchemy connection for use by pandas. Needed to create base load profiles.
    settings : dict
        User-defined parameters from a settings file

    Returns
    -------
    DataFrame
        Hourly generation profiles for DG resources in each region. Not all regions
        need to be accounted for.

    Raises
    ------
    KeyError
        If the calculation method specified in settings is not 'capacity' or 'fraction_load'
    """

    year = settings["model_year"]
    dg_profiles_path = (
        Path(settings["input_folder"]) / settings["distributed_gen_profiles_fn"]
    )

    hourly_norm_profiles = pd.read_csv(dg_profiles_path)
    profile_regions = hourly_norm_profiles.columns

    dg_calc_methods = settings["distributed_gen_method"]
    dg_calc_values = settings["distributed_gen_values"]

    assert (
        year in dg_calc_values
    ), f"The years in settings parameter 'distributed_gen_values' do not match the model years."

    for region in dg_calc_values[year]:
        assert region in set(profile_regions), (
            "The profile regions in settings parameter 'distributed_gen_values' do not\n"
            f"match the regions in {settings['distributed_gen_profiles_fn']} for year {year}"
        )

    if "fraction_load" in dg_calc_methods.values():
        regional_load = make_load_curves(pg_engine, settings)

    dg_hourly_gen = pd.DataFrame(columns=dg_calc_methods.keys())

    for region, method in dg_calc_methods.items():
        region_norm_profile = hourly_norm_profiles[region]
        region_calc_value = dg_calc_values[year][region]

        if method == "capacity":
            dg_hourly_gen[region] = calc_dg_capacity_method(
                region_norm_profile, region_calc_value
            )
        elif method == "fraction_load":
            region_load = regional_load[region]
            dg_hourly_gen[region] = calc_dg_frac_load_method(
                region_norm_profile, region_calc_value, region_load, settings
            )
        else:
            raise KeyError(
                "The settings parameter 'distributed_gen_method' can only have key "
                "values of 'capapacity' or 'fraction_load' for each region.\n"
                f"The value in your settings file is {method}"
            )

    if len(dg_hourly_gen) == 8784:
        remove_feb_29(dg_hourly_gen)

    return dg_hourly_gen


def calc_dg_capacity_method(dg_profile, dg_capacity):
    """Calculate the hourly distributed generation in a single region when given
    installed capacity.

    Parameters
    ----------
    dg_profile : Series
        Hourly normalized generation profile
    dg_capacity : float
        Total installed DG capacity

    Returns
    -------
    Series
        8760 hourly generation
    """

    hourly_gen = dg_profile * dg_capacity

    return hourly_gen.values


def calc_dg_frac_load_method(dg_profile, dg_requirement, regional_load, settings):
    """Calculate the hourly distributed generation in a single region where generation
    required to be a fraction of total sales.

    Parameters
    ----------
    dg_profile : Series
        Hourly normalized generation profile
    dg_requirement : float
        The fraction of total sales that DG must constitute
    regional_load : Series
        Hourly load for a given region
    settings : dict
        User-defined parameters from a settings file

    Returns
    -------
    Series
        8760 hourly generation
    """

    annual_load = regional_load.sum()
    dg_capacity_factor = dg_profile.mean()
    distribution_loss = settings["avg_distribution_loss"]

    required_dg_gen = annual_load * dg_requirement * (1 - distribution_loss)
    dg_capacity = required_dg_gen / 8760 / dg_capacity_factor

    hourly_gen = dg_profile * dg_capacity

    return hourly_gen


# powergenome/nrelatb.py

"""
Functions to fetch and modify NREL ATB data from PUDL
"""

import copy
import collections
import logging
import operator
from pathlib import Path
from typing import Dict, List, Tuple, Union

import numpy as np
import pandas as pd
import sqlalchemy

from powergenome.params import DATA_PATHS, build_resource_clusters
from powergenome.price_adjustment import inflation_price_adjustment
from powergenome.resource_clusters import ClusterBuilder, map_nrel_atb_technology
from powergenome.util import reverse_dict_of_lists, remove_leading_zero

idx = pd.IndexSlice
logger = logging.getLogger(__name__)


def fetch_atb_costs(
    pg_engine: sqlalchemy.engine.base.Engine,
    settings: dict,
    offshore_spur_costs: pd.DataFrame = None,
) -> pd.DataFrame:
    """Get NREL ATB power plant cost data from database, filter where applicable.

    This function can also remove NREL ATB offshore spur costs if more accurate costs
    will be included elsewhere (e.g. as part of total interconnection costs).

    Parameters
    ----------
    pg_engine : sqlalchemy.Engine
        A sqlalchemy connection for use by pandas
    settings : dict
        User-defined parameters from a settings file. Needs to have keys
        `atb_data_year`, `atb_new_gen`, and `target_usd_year`. If the key
        `atb_financial_case` is not included, the default value will be "Market".
    offshore_spur_costs : pd.DataFrame
        An optional dataframe with spur costs for offshore wind resources. These costs
        are included in ATB for a fixed distance (same for all sites). PowerGenome
        interconnection costs for offshore sites include a spur cost calculated
        using actual distance from shore.

    Returns
    -------
    pd.DataFrame
        Power plant cost data with columns:
        ['technology', 'cap_recovery_years', 'cost_case', 'financial_case',
       'basis_year', 'tech_detail', 'fixed_o_m_mw', 'variable_o_m_mwh', 'capex', 'cf',
       'fuel', 'lcoe', 'wacc_real']
    """
    logger.info("Loading NREL ATB data")

    col_names = [
        "technology",
        "tech_detail",
        "cost_case",
        "parameter",
        "basis_year",
        "parameter_value",
        "dollar_year",
    ]
    atb_year = settings["atb_data_year"]
    fin_case = settings.get("atb_financial_case", "Market")

    # Fetch cost data from sqlite and create dataframe. Only get values for techs/cases
    # listed in the settings file.
    all_rows = []
    wacc_rows = []
    tech_list = []
    techs = settings["atb_new_gen"]
    mod_techs = []
    if settings.get("modified_atb_new_gen"):
        for _, m in settings.get("modified_atb_new_gen").items():
            mod_techs.append(
                [m["atb_technology"], m["atb_tech_detail"], m["atb_cost_case"], None]
            )

    cost_params = (
        "capex_mw",
        "fixed_o_m_mw",
        "variable_o_m_mwh",
        "capex_mwh",
        "fixed_o_m_mwh",
    )
    # add_pv_wacc = True
    cols = ["technology", "tech_detail", "financial_case", "cost_case", "atb_year"]
    valid_inputs = db_col_values(pg_engine, "technology_costs_nrelatb", cols)
    for tech in techs + mod_techs:
        tech, tech_detail, cost_case, _ = tech
        # if tech == "UtilityPV":
        #     add_pv_wacc = False
        tech_params = [tech, tech_detail, fin_case, cost_case, atb_year]
        for param in tech_params:
            if param not in valid_inputs:
                raise ValueError(
                    f"When getting technology costs, the parameter {param} does not have "
                    "a valid matching value in the database table."
                )
        s = f"""
        SELECT technology, tech_detail, cost_case, parameter, basis_year, parameter_value, dollar_year
        from technology_costs_nrelatb
        where
            technology == "{tech}"
            AND tech_detail == "{tech_detail}"
            AND financial_case == "{fin_case}"
            AND cost_case == "{cost_case}"
            AND atb_year == {atb_year}
            AND parameter IN ({','.join('?'*len(cost_params))})
        """
        all_rows.extend(pg_engine.execute(s, cost_params).fetchall())

        if (tech, cost_case) not in tech_list:
            # ATB2020 summary file provides a single WACC for each technology and a single
            # tech detail of "*", so need to fetch this separately from other cost params.
            # Only need to fetch once per technology.
            wacc_s = f"""
            select technology, cost_case, basis_year, parameter_value
            from technology_costs_nrelatb
            where
                technology == "{tech}"
                AND financial_case == "{fin_case}"
                AND cost_case == "{cost_case}"
                AND atb_year == {atb_year}
                AND parameter == "wacc_real"
            """
            wacc_rows.extend(pg_engine.execute(wacc_s).fetchall())

        tech_list.append((tech, cost_case))
    tech_names = [t[0] for t in tech_list]
    if "Battery" not in tech_names:
        df = pd.DataFrame(all_rows, columns=col_names)
        wacc_df = pd.DataFrame(
            wacc_rows, columns=["technology", "cost_case", "basis_year", "wacc_real"]
        )
    else:
        # ATB doesn't have a WACC for battery storage. We use UtilityPV WACC as a default
        # stand-in -- make sure we have it in case.
        s = 'SELECT DISTINCT("technology") from technology_costs_nrelatb WHERE parameter == "wacc_real"'
        atb_techs = [x[0] for x in pg_engine.execute(s).fetchall()]
        battery_wacc_standin = settings.get("atb_battery_wacc")
        battery_tech = [x for x in techs if x[0] == "Battery"][0]
        if isinstance(battery_wacc_standin, float):
            if battery_wacc_standin > 0.1:
                logger.warning(
                    f"You defined a battery WACC of {battery_wacc_standin}, which seems"
                    " very high. Check settings parameter `atb_battery_wacc`."
                )
            battery_wacc_rows = [
                (battery_tech[0], battery_tech[2], year, battery_wacc_standin)
                for year in range(2017, 2051)
            ]
            wacc_rows.extend(battery_wacc_rows)
        elif battery_wacc_standin in atb_techs:
            # if battery_wacc_standin in tech_list:
            #     pass
            # else:
            logger.info(
                f"Using {battery_wacc_standin} {fin_case} WACC for Battery storage."
            )
            for cost_case in ["Mid", "Moderate"]:
                wacc_s = f"""
                select technology, cost_case, basis_year, parameter_value
                from technology_costs_nrelatb
                where
                    technology == "{battery_wacc_standin}"
                    AND financial_case == "{fin_case}"
                    AND cost_case == "{cost_case}"
                    AND atb_year == {atb_year}
                    AND parameter == "wacc_real"

                """
                b_rows = pg_engine.execute(wacc_s).fetchall()
                battery_wacc_rows = [
                    (battery_tech[0], battery_tech[2], b_row[2], b_row[3])
                    for b_row in b_rows
                ]
                wacc_rows.extend(battery_wacc_rows)
        else:
            raise ValueError(
                f"The settings key `atb_battery_wacc` value is {battery_wacc_standin}. It "
                f"should either be a float or a string from the list {atb_techs}."
            )

        df = pd.DataFrame(all_rows, columns=col_names)
        wacc_df = pd.DataFrame(
            wacc_rows, columns=["technology", "cost_case", "basis_year", "wacc_real"]
        )

    # Transform from tidy to wide dataframe, which makes it easier to fill generator
    # rows with the correct values.
    atb_costs = (
        df.drop_duplicates()
        .set_index(
            [
                "technology",
                "tech_detail",
                "cost_case",
                "dollar_year",
                "basis_year",
                "parameter",
            ]
        )
        .unstack(level=-1)
    )
    atb_costs.columns = atb_costs.columns.droplevel(0)
    atb_costs = (
        atb_costs.reset_index()
        .merge(wacc_df, on=["technology", "cost_case", "basis_year"], how="left")
        .drop_duplicates()
    )
    atb_costs = atb_costs.fillna(0)

    usd_columns = [
        "fixed_o_m_mw",
        "fixed_o_m_mwh",
        "variable_o_m_mwh",
        "capex_mw",
        "capex_mwh",
    ]
    for col in usd_columns:
        if col not in atb_costs.columns:
            atb_costs[col] = 0

    atb_target_year = settings["target_usd_year"]
    if not atb_costs.empty:
        atb_costs[usd_columns] = atb_costs.apply(
            lambda row: inflation_price_adjustment(
                row[usd_columns],
                base_year=row["dollar_year"],
                target_year=atb_target_year,
            ),
            axis=1,
        )

    if any("PV" in tech for tech in tech_list) and atb_year == 2019:
        print("Inflating ATB 2019 PV costs from DC to AC")
        atb_costs.loc[
            atb_costs["technology"].str.contains("PV"),
            ["capex_mw", "fixed_o_m_mw", "variable_o_m_mwh"],
        ] *= settings.get("pv_ac_dc_ratio", 1.3)
    elif atb_year > 2019:
        logger.info("PV costs are already in AC units, not inflating the cost.")

    if offshore_spur_costs is not None and "OffShoreWind" in atb_costs["technology"]:
        idx_cols = ["technology", "tech_detail", "cost_case", "basis_year"]
        offshore_spur_costs = offshore_spur_costs.set_index(idx_cols)
        atb_costs = atb_costs.set_index(idx_cols)

        atb_costs.loc[idx["OffShoreWind", :, :, :], "capex_mw"] = (
            atb_costs.loc[idx["OffShoreWind", :, :, :], "capex_mw"]  # .values
            - offshore_spur_costs.loc[
                idx["OffShoreWind", :, :, :], "capex_mw"
            ]  # .values
        )
        atb_costs = atb_costs.reset_index()

    return atb_costs


def db_col_values(
    engine: sqlalchemy.engine, table: str, cols: List[str]
) -> List[Union[str, int]]:
    """Find all distinct values in one or more columns of a database table.

    This function can be used to check user inputs that are applied as filters against
    existing table values to reduce the risk of SQL injection attacks. All distinct
    values are returned in a single list for simplicity.

    Parameters
    ----------
    engine : sqlalchemy.engine
        Connection engine to the database.
    table : str
        Name of the database table.
    cols : List[str]
        Name of the column(s) to check.

    Returns
    -------
    List[Union[str, int]]
        All distinct values from the database table column(s).
    """
    valid_inputs = []
    for col in cols:
        s = f"SELECT DISTINCT {col} from {table}"
        valid_inputs.extend(pd.read_sql_query(s, engine)[col].to_list())

    return valid_inputs


def fetch_atb_offshore_spur_costs(
    pg_engine: sqlalchemy.engine.base.Engine, settings: dict
) -> pd.DataFrame:
    """Load offshore spur-line costs and convert to desired dollar-year.

    Parameters
    ----------
    pg_engine : sqlalchemy.Engine
        A sqlalchemy connection for use by pandas
    settings : dict
        User-defined parameters from a settings file. Needs to have keys `atb_data_year`
        and `target_usd_year`.

    Returns
    -------
    pd.DataFrame
        Total offshore spur line capex from ATB for each technology/tech_detail/
        basis_year/cost_case combination.
    """
    spur_costs = pd.read_sql_table("offshore_spur_costs_nrelatb", pg_engine)
    spur_costs = spur_costs.loc[spur_costs["atb_year"] == settings["atb_data_year"], :]

    atb_target_year = settings["target_usd_year"]

    spur_costs["capex_mw"] = spur_costs.apply(
        lambda row: inflation_price_adjustment(
            row["capex_mw"], base_year=row["dollar_year"], target_year=atb_target_year
        ),
        axis=1,
    )

    # ATB assumes a 30km distance for offshore spur. Normalize to per mile
    spur_costs["capex_mw_mile"] = spur_costs["capex_mw"] / 30 / 1.60934

    return spur_costs


def fetch_atb_heat_rates(
    pg_engine: sqlalchemy.engine.base.Engine, settings: dict
) -> pd.DataFrame:
    """Get heat rate projections for power plants

    Data is originally from AEO, NREL does a linear interpolation between current and
    final years.

    Parameters
    ----------
    pg_engine : sqlalchemy.Engine
        A sqlalchemy connection for use by pandas
    settings : dict
        User-defined parameters from a settings file. Needs to have key `atb_data_year`.

    Returns
    -------
    pd.DataFrame
        Power plant heat rate data by year with columns:
        ['technology', 'tech_detail', 'cost_case', 'basis_year', 'heat_rate']
    """

    heat_rates = pd.read_sql_table("technology_heat_rates_nrelatb", pg_engine)
    if settings["atb_data_year"] not in heat_rates["atb_year"].unique():
        max_atb_year = heat_rates["atb_year"].max()
        logger.warning(
            f"Your settings file has parameter `atb_year` of {settings['atb_data_year']}"
            ", which isn't in the table `technology_heat_rates_nrelatb`. Using "
            f"{max_atb_year} instead."
        )
        settings["atb_data_year"] = max_atb_year
    heat_rates = heat_rates.loc[heat_rates["atb_year"] == settings["atb_data_year"], :]

    return heat_rates


def atb_fixed_var_om_existing(
    results: pd.DataFrame,
    atb_hr_df: pd.DataFrame,
    settings: dict,
    pg_engine: sqlalchemy.engine.base.Engine,
    coal_fgd_df: pd.DataFrame,
) -> pd.DataFrame:
    """Add fixed and variable O&M for existing power plants

    ATB O&M data for new power plants are used as reference values. Fixed and variable
    O&M for each technology and heat rate are calculated. Assume that O&M scales with
    heat rate from new plants to existing generators. A separate multiplier for fixed
    O&M is specified in the settings file.

    Parameters
    ----------
    results : DataFrame
        Compiled results of clustered power plants with weighted average heat rates.
        Note that column names should include "technology", "Heat_Rate_MMBTU_per_MWh",
        and "region". Technology names should not yet be converted to snake case.
    atb_hr_df : DataFrame
        Heat rate data from NREL ATB
    settings : dict
        User-defined parameters from a settings file

    Returns
    -------
    DataFrame
        Same as incoming "results" dataframe but with new columns
        "Fixed_OM_Cost_per_MWyr" and "Var_OM_Cost_per_MWh"
    """
    logger.info("Adding fixed and variable O&M for existing plants")

    existing_year = settings["atb_existing_year"]

    techs = {}
    for eia, atb in settings["eia_atb_tech_map"].items():
        if not isinstance(atb, list):
            atb = [atb]
        techs[eia] = atb[0].split("_")

    target_usd_year = settings["target_usd_year"]
    simple_o_m = {
        "Natural Gas Steam Turbine": {
            "variable_o_m_mwh": inflation_price_adjustment(1.0, 2017, target_usd_year)
        },
        "Coal": {
            "variable_o_m_mwh": inflation_price_adjustment(1.78, 2017, target_usd_year)
        },
        "Conventional Hydroelectric": {
            "fixed_o_m_mw": inflation_price_adjustment(
                44.56 * 1000, 2017, target_usd_year
            ),
            "variable_o_m_mwh": 0,
        },
        "Geothermal": {
            "fixed_o_m_mw": inflation_price_adjustment(
                198.04 * 1000, 2017, target_usd_year
            ),
            "variable_o_m_mwh": 0,
        },
        "Pumped Hydro": {
            "fixed_o_m_mw": inflation_price_adjustment(
                (23.63 + 14.83) * 1000, 2017, target_usd_year
            ),
            "variable_o_m_mwh": 0,
        },
    }

    # Load all atb O&M values at once rather than querying the db thousands of times
    nems_o_m_techs = [
        "Combined Cycle",
        "Combustion Turbine",
        "Coal",
        "Steam Turbine",
        "Hydroelectric",
        "Geothermal",
        "Nuclear",
    ]
    atb_techs = [
        (tech, tech_detail)
        for k, (tech, tech_detail) in techs.items()
        if tech not in nems_o_m_techs
    ]
    s = f"""
        select technology, tech_detail, parameter, AVG(parameter_value) as parameter_value
        from technology_costs_nrelatb
        where
            basis_year == ?
            AND financial_case == "Market"
            AND cost_case in ("Mid", "Moderate")
            AND atb_year == ?
            AND parameter in ("variable_o_m_mwh", "fixed_o_m_mw", "fixed_o_m_mwh")
            AND
                ({' OR '.join(["(technology==? and tech_detail==?)"]*len(atb_techs))})
            GROUP BY technology, tech_detail, parameter
        """
    params = [existing_year, settings["atb_data_year"]] + [
        item for sublist in atb_techs for item in sublist
    ]
    atb_om = pd.read_sql_query(
        s,
        pg_engine,
        params=params,
        index_col=["technology", "tech_detail", "parameter"],
    )
    atb_hr_df = atb_hr_df.set_index(
        ["technology", "tech_detail", "basis_year"]
    ).sort_index()

    df_list = []
    grouped_results = results.reset_index().groupby(["technology"], as_index=False)
    for group, _df in grouped_results:
        _df = calc_om(
            _df,
            atb_hr_df,
            atb_om,
            settings,
            pg_engine,
            coal_fgd_df,
            existing_year,
            techs,
            simple_o_m,
            group,
        )

        df_list.append(_df)

    mod_results = pd.concat(df_list, ignore_index=True)

    # Fill na FOM values, first by technology and then across all techs
    if not mod_results.loc[mod_results["Fixed_OM_Cost_per_MWyr"].isna()].empty:
        df_list = []
        for tech, _df in mod_results.groupby("technology"):
            _df.loc[
                _df["Fixed_OM_Cost_per_MWyr"].isna(), "Fixed_OM_Cost_per_MWyr"
            ] = _df["Fixed_OM_Cost_per_MWyr"].mean()
            df_list.append(_df)
        mod_results = pd.concat(df_list, ignore_index=True)
        mod_results.loc[
            mod_results["Fixed_OM_Cost_per_MWyr"].isna(), "Fixed_OM_Cost_per_MWyr"
        ] = mod_results["Fixed_OM_Cost_per_MWyr"].mean()
    mod_results.loc[:, "Fixed_OM_Cost_per_MWyr"] = mod_results.loc[
        :, "Fixed_OM_Cost_per_MWyr"
    ].astype(int)
    mod_results.loc[:, "Fixed_OM_Cost_per_MWhyr"] = mod_results.loc[
        :, "Fixed_OM_Cost_per_MWhyr"
    ].astype(int)
    mod_results.loc[:, "Var_OM_Cost_per_MWh"] = mod_results.loc[
        :, "Var_OM_Cost_per_MWh"
    ]

    return mod_results


def calc_om(
    df: pd.DataFrame,
    atb_hr_df: pd.DataFrame,
    atb_om_df: pd.DataFrame,
    settings: dict,
    pg_engine: sqlalchemy.engine.base.Engine,
    coal_fgd_df: pd.DataFrame,
    existing_year: int,
    techs: Dict[str, List[str]],
    simple_o_m: Dict[str, float],
    group: str,
) -> pd.DataFrame:
    """Calculate fixed and variable O&M for a single technology.

    Parameters
    ----------
    df : pd.DataFrame
        Units of generators at a single plant
    atb_hr_df : pd.DataFrame
        Heat rate data from NREL ATB
    atb_om_df : pd.DataFrame
        O&M costs from NREL ATB
    settings : dict
        User-defined parameters from a settings file
    pg_engine : sqlalchemy.engine.base.Engine
        Connection to the PowerGenome database
    coal_fgd_df : pd.DataFrame
        Table showing if a coal generating unit has FGD control technology
    existing_year : int
        Year of ATB to use as a proxy for existing plants
    techs : Dict[str, List[str]]
        Mapping of EIA technology name (key) to the ATB technology, tech_detail (value)
    simple_o_m : Dict[str, float]
        Mapping of simple O&M costs used by some generators (not dependent on size/age)
    group : str
        The EIA technology name

    Returns
    -------
    pd.DataFrame
        Modified copy of input with fixed and variable O&M costs

    Raises
    ------
    KeyError
        _description_
    KeyError
        _description_
    """
    df["Fixed_OM_Cost_per_MWhyr"] = 0
    target_usd_year = settings["target_usd_year"]
    eia_tech = group

    try:
        atb_tech, tech_detail = techs[eia_tech]
    except KeyError:
        if eia_tech in settings.get("tech_groups", {}) or {}:
            raise KeyError(
                f"{eia_tech} is defined in 'tech_groups' but doesn't have a "
                "corresponding ATB technology in 'eia_atb_tech_map'"
            )

        else:
            raise KeyError(
                f"{eia_tech} doesn't have a corresponding ATB technology in "
                "'eia_atb_tech_map'"
            )

    try:
        new_build_hr = atb_hr_df.loc[
            (atb_tech, tech_detail, existing_year),
            "heat_rate",
        ].mean()
        if not isinstance(new_build_hr, float):
            logger.warning(
                "\n\n****************\nCAUTION!!!\n\n"
                f"The calculated new build heat rate for {atb_tech}, {tech_detail} "
                f"should be a single value but is {new_build_hr}. This could cause "
                f"issues with your variable O&M costs for {eia_tech}. Please report "
                "this as an issue on the PowerGenome repository.\n"
            )
    except (ValueError, TypeError, KeyError):
        # Not all technologies have a heat rate. If they don't, just set both values
        # to 10.34 (33% efficiency)
        df["heat_rate_mmbtu_mwh"] = 10.34
        new_build_hr = 10.34

    try:
        atb_var_om_mwh = atb_om_df.loc[
            (atb_tech, tech_detail, "variable_o_m_mwh"), "parameter_value"
        ]
    except KeyError:
        atb_var_om_mwh = 0

    nems_o_m_techs = [
        "Combined Cycle",
        "Combustion Turbine",
        "Coal",
        "Steam Turbine",
        # "Hydroelectric",
        # "Geothermal",
        "Nuclear",
    ]

    if any(t in eia_tech for t in nems_o_m_techs):
        df_list = []
        for plant_id, _df in df.groupby("plant_id_eia", as_index=False):
            # Change CC and CT O&M to EIA NEMS values, which are much higher for CCs and
            # lower for CTs than a heat rate & linear mulitpler correction to the ATB
            # values.
            # Add natural gas steam turbine O&M.
            # Also using the new values for coal plants, assuming 40-50 yr age and half
            # FGD
            # https://www.eia.gov/analysis/studies/powerplants/generationcost/pdf/full_report.pdf
            # logger.info(f"Using NEMS values for {eia_tech} fixed/variable O&M")

            if "Combined Cycle" in eia_tech:
                # https://www.eia.gov/analysis/studies/powerplants/generationcost/pdf/full_report.pdf
                assert not _df[settings["capacity_col"]].isnull().all()
                plant_capacity = _df[settings["capacity_col"]].sum()
                if plant_capacity < 500:
                    fixed = 15.62 * 1000
                    variable = 4.31
                elif 500 <= plant_capacity < 1000:
                    fixed = 9.27 * 1000
                    variable = 3.42
                else:
                    fixed = 11.68 * 1000
                    variable = 3.37

                _df["Fixed_OM_Cost_per_MWyr"] = inflation_price_adjustment(
                    fixed, 2017, target_usd_year
                )
                _df["Var_OM_Cost_per_MWh"] = inflation_price_adjustment(
                    variable, 2017, target_usd_year
                )

            if "Combustion Turbine" in eia_tech:
                # need to adjust the EIA fixed/variable costs because they have no
                # variable cost per MWh for existing CTs but they do have per MWh for
                # new build. Assume $11/MWh from new-build and 4% CF:
                # (11*8760*0.04/1000)=$3.85/kW-yr. Scale the new-build variable
                # (~$11/MWh) by relative heat rate and subtract a /kW-yr value as
                # calculated above from the FOM.
                # Based on conversation with Jesse J. on Dec 20, 2019.
                plant_capacity = _df[settings["capacity_col"]].sum()
                op, op_value = (
                    settings.get("atb_modifiers", {})
                    .get("ngct", {})
                    .get("Var_OM_Cost_per_MWh", (None, None))
                )

                if op:
                    f = operator.attrgetter(op)
                    atb_var_om_mwh = f(operator)(atb_var_om_mwh, op_value)

                variable = atb_var_om_mwh  # * (existing_hr / new_build_hr)

                if plant_capacity < 100:
                    annual_capex = 9.0 * 1000
                    fixed = annual_capex + 5.96 * 1000
                elif 100 <= plant_capacity <= 300:
                    annual_capex = 6.18 * 1000
                    fixed = annual_capex + 6.43 * 1000
                else:
                    annual_capex = 6.95 * 1000
                    fixed = annual_capex + 3.99 * 1000

                fixed = fixed - (variable * 8760 * 0.04)

                _df["Fixed_OM_Cost_per_MWyr"] = inflation_price_adjustment(
                    fixed, 2017, target_usd_year
                )
                _df["Var_OM_Cost_per_MWh"] = inflation_price_adjustment(
                    variable, 2017, target_usd_year
                )

            if "Natural Gas Steam Turbine" in eia_tech:
                # https://www.eia.gov/analysis/studies/powerplants/generationcost/pdf/full_report.pdf
                assert not _df[settings["capacity_col"]].isnull().all()
                plant_capacity = _df[settings["capacity_col"]].sum()
                if plant_capacity < 500:
                    annual_capex = 18.86 * 1000
                    fixed = annual_capex + 29.73 * 1000
                elif 500 <= plant_capacity < 1000:
                    annual_capex = 11.57 * 1000
                    fixed = annual_capex + 17.98 * 1000
                else:
                    annual_capex = 10.82 * 1000
                    fixed = annual_capex + 14.51 * 1000

                _df["Fixed_OM_Cost_per_MWyr"] = inflation_price_adjustment(
                    fixed, 2017, target_usd_year
                )
                _df["Var_OM_Cost_per_MWh"] = simple_o_m["Natural Gas Steam Turbine"][
                    "variable_o_m_mwh"
                ]

            if "Coal" in eia_tech:
                assert not _df[settings["capacity_col"]].isnull().all()
                plant_capacity = _df[settings["capacity_col"]].sum()

                age = settings["model_year"] - _df.operating_date.dt.year
                try:
                    age = age.fillna(age.mean())
                except:
                    age = age.fillna(40)
                gen_ids = _df["generator_id"].to_list()
                fgd = coal_fgd_df.query(
                    "plant_id_eia == @plant_id & generator_id in @gen_ids"
                )["fgd"].values
                if not np.any(fgd):
                    gen_ids = [remove_leading_zero(g) for g in gen_ids]
                    fgd = coal_fgd_df.query(
                        "plant_id_eia == @plant_id & generator_id in @gen_ids"
                    )["fgd"].values
                if not np.any(fgd):
                    # If FGD isn't found, use average of with/without FGD
                    fgd = np.ones_like(age) * 0.5

                    # https://www.eia.gov/analysis/studies/powerplants/generationcost/pdf/full_report.pdf
                annual_capex = (16.53 + (0.126 * age) + (5.68 * fgd)) * 1000

                if plant_capacity < 500:
                    fixed = 44.21 * 1000
                elif 500 <= plant_capacity < 1000:
                    fixed = 34.02 * 1000
                elif 1000 <= plant_capacity < 2000:
                    fixed = 28.52 * 1000
                else:
                    fixed = 33.27 * 1000

                _df["Fixed_OM_Cost_per_MWyr"] = inflation_price_adjustment(
                    fixed + annual_capex, 2017, target_usd_year
                )
                _df["Var_OM_Cost_per_MWh"] = simple_o_m["Coal"]["variable_o_m_mwh"]

            if "Nuclear" in eia_tech:
                num_units = len(_df)
                plant_capacity = _df[settings["capacity_col"]].sum()

                # Operating costs for different size/num units in 2016 INL report
                # "Economic and Market Challenges Facing the U.S. Nuclear Fleet"
                # https://gain.inl.gov/Shared%20Documents/Economics-Nuclear-Fleet.pdf,
                # table 1. Average of the two costs are used in each case.
                # The costs in that report include fuel and VOM. Assume $0.66/mmbtu
                # and $2.32/MWh plus 90% CF (ATB 2020) to get the costs below.
                # The INL report doesn't give a dollar year for costs, assume 2015.
                if num_units == 1 and plant_capacity < 900:
                    fixed = 315000
                elif num_units == 1 and plant_capacity >= 900:
                    fixed = 252000
                else:
                    fixed = 177000
                    # age = (settings["model_year"] - _df.operating_date.dt.year).values
                    # age = age.fillna(age.mean())
                    # age = age.fillna(40)
                    # EIA, 2020, "Assumptions to Annual Energy Outlook, Electricity Market Module,"
                    # Available: https://www.eia.gov/outlooks/aeo/assumptions/pdf/electricity.pdf
                    # fixed = np.ones_like(age)
                    # fixed[age < 30] *= 27 * 1000
                    # fixed[age >= 30] *= (27+37) * 1000

                _df["Fixed_OM_Cost_per_MWyr"] = inflation_price_adjustment(
                    fixed, 2015, target_usd_year
                )

                # If nuclear heat rates are NaN, set them to new build value
                _df.loc[
                    _df["heat_rate_mmbtu_mwh"].isna(), "heat_rate_mmbtu_mwh"
                ] = new_build_hr
                _df["Var_OM_Cost_per_MWh"] = atb_var_om_mwh * (
                    _df["heat_rate_mmbtu_mwh"].mean() / new_build_hr
                )

            df_list.append(_df)

        return pd.concat(df_list, ignore_index=True)

    elif "Hydroelectric" in eia_tech:
        df["Fixed_OM_Cost_per_MWyr"] = simple_o_m["Conventional Hydroelectric"][
            "fixed_o_m_mw"
        ]
        df["Var_OM_Cost_per_MWh"] = simple_o_m["Conventional Hydroelectric"][
            "variable_o_m_mwh"
        ]
    elif "Geothermal" in eia_tech:
        df["Fixed_OM_Cost_per_MWyr"] = simple_o_m["Geothermal"]["fixed_o_m_mw"]
        df["Var_OM_Cost_per_MWh"] = simple_o_m["Geothermal"]["variable_o_m_mwh"]
    elif "Pumped" in eia_tech:
        df["Fixed_OM_Cost_per_MWyr"] = simple_o_m["Pumped Hydro"]["fixed_o_m_mw"]
        df["Var_OM_Cost_per_MWh"] = simple_o_m["Pumped Hydro"]["variable_o_m_mwh"]
    else:
        try:
            atb_fixed_om_mw_yr = atb_om_df.loc[
                (atb_tech, tech_detail, "fixed_o_m_mw"), "parameter_value"
            ]
        except KeyError:
            atb_fixed_om_mw_yr = 0
        df["Fixed_OM_Cost_per_MWyr"] = atb_fixed_om_mw_yr
        df["Var_OM_Cost_per_MWh"] = atb_var_om_mwh * (
            df["heat_rate_mmbtu_mwh"] / new_build_hr
        )
    if atb_tech == "Battery":
        atb_fixed_om_mwh = atb_om_df.loc[
            (atb_tech, tech_detail, "fixed_o_m_mwh"), "parameter_value"
        ]
        df["Fixed_OM_Cost_per_MWhyr"] = atb_fixed_om_mwh

    return df


def single_generator_row(
    atb_costs_hr: pd.DataFrame,
    new_gen_type: str,
    model_year_range: Union[Tuple[int], List[int]],
) -> pd.DataFrame:
    """Create a data row with NREL ATB costs and performace for a single technology

    Parameters
    ----------
    atb_costs : pd.DataFrame
        Data from the sqlite tables of both resources costs and heat rates
    new_gen_type : str
        type of generating resource
    model_year_range : Union[Tuple[int], List[int]]
        All of the years that should be averaged over

    Returns
    -------
    pd.DataFrame
        A single row dataframe with average cost and performence values over the study
        period.
    """

    technology, tech_detail, cost_case, size_mw = new_gen_type
    numeric_cols = [
        "basis_year",
        "fixed_o_m_mw",
        "fixed_o_m_mwh",
        "variable_o_m_mwh",
        "capex_mw",
        "capex_mwh",
        "wacc_real",
        "heat_rate",
    ]
    s = atb_costs_hr.loc[
        (atb_costs_hr["technology"] == technology)
        & (atb_costs_hr["tech_detail"] == tech_detail)
        & (atb_costs_hr["cost_case"] == cost_case)
        & (atb_costs_hr["basis_year"].isin(model_year_range)),
        numeric_cols,
    ].mean()
    cols = ["technology", "cost_case", "tech_detail"] + numeric_cols
    row = pd.DataFrame([technology, cost_case, tech_detail] + s.to_list(), index=cols).T

    row["Cap_Size"] = size_mw

    return row


def investment_cost_calculator(capex, wacc, cap_rec_years):
    capex = np.asarray(capex, dtype=float)
    wacc = np.asarray(wacc, dtype=float)

    for variable in capex, wacc, cap_rec_years:
        if np.isnan(variable).any():
            raise ValueError(f"Investment costs contains nan values")

    inv_cost = capex * (
        np.exp(wacc * cap_rec_years)
        * (np.exp(wacc) - 1)
        / (np.exp(wacc * cap_rec_years) - 1)
    )

    return inv_cost


def regional_capex_multiplier(
    df: pd.DataFrame,
    region: str,
    region_map: Dict[str, str],
    tech_map: Dict[str, str],
    regional_multipliers: pd.DataFrame,
) -> pd.DataFrame:
    cost_region = region_map[region]
    tech_multiplier = regional_multipliers.loc[cost_region, :].squeeze()
    avg_multiplier = tech_multiplier.mean()

    tech_multiplier = tech_multiplier.fillna(avg_multiplier)

    tech_multiplier_map = {}
    for atb_tech, eia_tech in tech_map.items():
        if df["technology"].str.contains(atb_tech).sum() > 0:
            full_atb_tech = df.loc[
                df["technology"].str.contains(atb_tech).idxmax(), "technology"
            ]
            tech_multiplier_map[full_atb_tech] = tech_multiplier.at[eia_tech]
        if df["technology"].str.contains(atb_tech).sum() > 1:
            s = f"""
    ***************************
    There is an issue with assigning regional cost multipliers. In your settings file
    under the parameter 'cost_multiplier_technology_map`, the EIA technology '{eia_tech}'
    has an ATB technology '{atb_tech}'. This ATB name matches more than one new ATB tech
    listed in the settings parameter 'atb_new_gen'. Only the first matching tech in
    'atb_new_gen' will get a valid regional cost multiplier; the rest will have values of
    0, which will lead to annual investment costs of $0.
        """
            logger.warning(s)
    df["Inv_Cost_per_MWyr"] *= df["technology"].map(tech_multiplier_map)
    df["Inv_Cost_per_MWhyr"] *= df["technology"].map(tech_multiplier_map)
    df["regional_cost_multiplier"] = df["technology"].map(tech_multiplier_map)

    return df


def add_modified_atb_generators(
    settings: dict,
    atb_costs_hr: pd.DataFrame,
    model_year_range: Union[Tuple[int], List[int]],
) -> pd.DataFrame:
    """Create a modified version of an ATB generator.

    For each parameter (capex, heat_rate, etc) that users want modified they should
    specify a list of [<operator>, <value>]. The operator can be add, mul, truediv, or
    sub (substract). This is used to modify individual parameters of the ATB resource.

    Parameters
    ----------
    settings : dict
        User-defined parameters from a settings file
    atb_costs_hr : pd.DataFrame
        Cost and heat rate data for ATB resources
    model_year_range : Union[Tuple[int], List[int]]
        A list or range of years to average ATB values from.

    Returns
    -------
    pd.DataFrame
        Row or rows of modified ATB resources. Each row includes the colums:
        ['technology', 'cost_case', 'tech_detail', 'basis_year', 'fixed_o_m_mw',
       'fixed_o_m_mwh', 'variable_o_m_mwh', 'capex', 'capex_mwh', 'cf', 'fuel',
       'lcoe', 'o_m', 'wacc_real', 'heat_rate', 'Cap_Size'].
    """

    # copy settings so popped keys aren't removed permenantly
    _settings = copy.deepcopy(settings)

    allowed_operators = ["add", "mul", "truediv", "sub"]

    mod_tech_list = []
    for name, mod_tech in _settings["modified_atb_new_gen"].items():
        atb_technology = mod_tech.pop("atb_technology")
        atb_tech_detail = mod_tech.pop("atb_tech_detail")
        atb_cost_case = mod_tech.pop("atb_cost_case")
        size_mw = mod_tech.pop("size_mw")

        new_gen_type = (atb_technology, atb_tech_detail, atb_cost_case, size_mw)

        gen = single_generator_row(atb_costs_hr, new_gen_type, model_year_range)
        gen["technology"] = mod_tech.pop("new_technology")
        gen["tech_detail"] = mod_tech.pop("new_tech_detail", "")
        gen["cost_case"] = mod_tech.pop("new_cost_case")

        for parameter, op_list in mod_tech.items():
            if isinstance(op_list, float) | isinstance(op_list, int):
                gen[parameter] = op_list
            else:
                assert len(op_list) == 2, (
                    "Two values, an operator and a numeric value, are needed in the parameter\n"
                    f"'{parameter}' for technology '{name}' in 'modified_atb_new_gen'."
                )
                op, op_value = op_list

                assert parameter in gen.columns, (
                    f"'{parameter}' is not a valid parameter for new resources. Check '{name}'\n"
                    "in 'modified_atb_new_gen' of the settings file."
                )
                assert op in allowed_operators, (
                    f"The key {parameter} for technology {name} needs a valid operator from the list\n"
                    f"{allowed_operators}\n"
                    "in the format [<operator>, <value>] to modify the properties of an existing generator.\n"
                )

                f = operator.attrgetter(op)
                gen[parameter] = f(operator)(gen[parameter], op_value)

        mod_tech_list.append(gen)

    mod_gens = pd.concat(mod_tech_list, ignore_index=True)

    return mod_gens


def atb_new_generators(atb_costs, atb_hr, settings, cluster_builder=None):
    """Add rows for new generators in each region

    Parameters
    ----------
    atb_costs : DataFrame
        All cost parameters from the SQL table for new generators. Should include:
        ['technology', 'cost_case', 'financial_case', 'basis_year', 'tech_detail',
        'capex', 'capex_mwh', 'fixed_o_m_mw', 'fixed_o_m_mwh', 'variable_o_m_mwh',
        'wacc_real']
    atb_hr : DataFrame
        The technology, tech_detail, and heat_rate of new generators from ATB.
    settings : dict
        User-defined parameters from a settings file
    cluster_builder : ClusterBuilder
        ClusterBuilder object. Reuse to save time. None by default.

    Returns
    -------
    DataFrame
        New generating resources in every region. Contains the columns:
        ['technology', 'basis_year', 'Fixed_OM_Cost_per_MWyr',
       'Fixed_OM_Cost_per_MWhyr', 'Var_OM_Cost_per_MWh', 'capex', 'capex_mwh',
       'Inv_Cost_per_MWyr', 'Inv_Cost_per_MWhyr', 'Heat_Rate_MMBTU_per_MWh',
       'Cap_Size', 'region']
    """
    logger.info("Creating new resources for each region.")
    new_gen_types = settings["atb_new_gen"]
    model_year = settings["model_year"]
    try:
        first_planning_year = settings["model_first_planning_year"]
        model_year_range = range(first_planning_year, model_year + 1)
    except KeyError:
        model_year_range = list(range(model_year + 1))

    regions = settings["model_regions"]

    atb_costs_hr = atb_costs.merge(
        atb_hr, on=["technology", "tech_detail", "cost_case", "basis_year"], how="left"
    )

    if new_gen_types:
        new_gen_df = pd.concat(
            [
                single_generator_row(atb_costs_hr, new_gen, model_year_range)
                for new_gen in new_gen_types
            ],
            ignore_index=True,
        )
    else:
        new_gen_df = pd.DataFrame(
            columns=["region", "technology", "tech_detail", "cost_case"]
        )
    # Add user-defined technologies
    # This should probably be separate from ATB techs, and the regional cost multipliers
    # should be its own function.
    if settings.get("additional_technologies_fn"):
        if isinstance(settings.get("additional_new_gen"), list):
            # user_costs, user_hr = load_user_defined_techs(settings)
            user_tech = load_user_defined_techs(settings)
            # new_gen_df = pd.concat([new_gen_df, user_costs], ignore_index=True, sort=False)
            new_gen_df = pd.concat(
                [new_gen_df, user_tech], ignore_index=True, sort=False
            )
            # atb_hr = pd.concat([atb_hr, user_hr], ignore_index=True, sort=False)
        else:
            logger.warning(
                "A filename for additional technologies was included but no technologies"
                " were specified in the settings file."
            )

    if settings.get("modified_atb_new_gen"):
        modified_gens = add_modified_atb_generators(
            settings, atb_costs_hr, model_year_range
        )
        new_gen_df = pd.concat(
            [new_gen_df, modified_gens], ignore_index=True, sort=False
        )

    new_gen_df = new_gen_df.rename(
        columns={
            "heat_rate": "Heat_Rate_MMBTU_per_MWh",
            "fixed_o_m_mw": "Fixed_OM_Cost_per_MWyr",
            "fixed_o_m_mwh": "Fixed_OM_Cost_per_MWhyr",
            "variable_o_m_mwh": "Var_OM_Cost_per_MWh",
        }
    )

    # Adjust values for CT/CC generators to match advanced techs in NEMS rather than
    # ATB average of advanced and conventional.
    # This is now generalized for changes to ATB values for any technology type.
    for tech, _tech_modifiers in (settings.get("atb_modifiers") or {}).items():
        tech_modifiers = copy.deepcopy(_tech_modifiers)
        assert isinstance(tech_modifiers, dict), (
            "The settings parameter 'atb_modifiers' must be a nested list.\n"
            "Each top-level key is a short name of the technology, with a nested"
            " dictionary of items below it."
        )
        assert (
            "technology" in tech_modifiers
        ), "Each nested dictionary in atb_modifiers must have a 'technology' key."
        assert (
            "tech_detail" in tech_modifiers
        ), "Each nested dictionary in atb_modifiers must have a 'tech_detail' key."

        technology = tech_modifiers.pop("technology")
        tech_detail = tech_modifiers.pop("tech_detail")

        allowed_operators = ["add", "mul", "truediv", "sub"]

        for key, op_list in tech_modifiers.items():
            if isinstance(op_list, float) | isinstance(op_list, int):
                new_gen_df.loc[
                    (new_gen_df.technology == technology)
                    & (new_gen_df.tech_detail == tech_detail),
                    key,
                ] = op_list
            else:
                assert len(op_list) == 2, (
                    "Two values, an operator and a numeric value, are needed in the parameter\n"
                    f"'{key}' for technology '{tech}' in 'atb_modifiers'."
                )
                op, op_value = op_list

                assert op in allowed_operators, (
                    f"The key {key} for technology {tech} needs a valid operator from the list\n"
                    f"{allowed_operators}\n"
                    "in the format [<operator>, <value>] to modify the properties of an existing generator.\n"
                )

                f = operator.attrgetter(op)
                new_gen_df.loc[
                    (new_gen_df.technology == technology)
                    & (new_gen_df.tech_detail == tech_detail),
                    key,
                ] = f(operator)(
                    new_gen_df.loc[
                        (new_gen_df.technology == technology)
                        & (new_gen_df.tech_detail == tech_detail),
                        key,
                    ],
                    op_value,
                )

    new_gen_df["technology"] = (
        new_gen_df[["technology", "tech_detail", "cost_case"]]
        .astype(str)
        .agg("_".join, axis=1)
    )

    new_gen_df["cap_recovery_years"] = settings["atb_cap_recovery_years"]

    if new_gen_df.empty:
        results = new_gen_df.copy()
    else:
        for tech, years in (settings.get("alt_atb_cap_recovery_years") or {}).items():
            new_gen_df.loc[
                new_gen_df["technology"].str.lower().str.contains(tech.lower()),
                "cap_recovery_years",
            ] = years

        new_gen_df["Inv_Cost_per_MWyr"] = investment_cost_calculator(
            capex=new_gen_df["capex_mw"],
            wacc=new_gen_df["wacc_real"],
            cap_rec_years=new_gen_df["cap_recovery_years"],
        )

        new_gen_df["Inv_Cost_per_MWhyr"] = investment_cost_calculator(
            capex=new_gen_df["capex_mwh"],
            wacc=new_gen_df["wacc_real"],
            cap_rec_years=new_gen_df["cap_recovery_years"],
        )

        # Set no capacity limit on new resources that aren't renewables.
        new_gen_df["Max_Cap_MW"] = -1
        new_gen_df["Max_Cap_MWh"] = -1
        regional_cost_multipliers = pd.read_csv(
            DATA_PATHS["cost_multipliers"]
            / settings.get(
                "cost_multiplier_fn", "AEO_2020_regional_cost_corrections.csv"
            ),
            index_col=0,
        )
        if settings.get("user_regional_cost_multiplier_fn"):
            user_cost_multipliers = pd.read_csv(
                Path(settings["input_folder"])
                / settings["user_regional_cost_multiplier_fn"],
                index_col=0,
            )
            regional_cost_multipliers = pd.concat(
                [regional_cost_multipliers, user_cost_multipliers], axis=1
            )
        rev_mult_region_map = reverse_dict_of_lists(
            settings["cost_multiplier_region_map"]
        )
        rev_mult_tech_map = reverse_dict_of_lists(
            settings["cost_multiplier_technology_map"]
        )
        df_list = []
        for region in regions:
            _df = new_gen_df.copy()
            _df["region"] = region
            _df = regional_capex_multiplier(
                _df,
                region,
                rev_mult_region_map,
                rev_mult_tech_map,
                regional_cost_multipliers,
            )
            _df = add_renewables_clusters(_df, region, settings, cluster_builder)

            if region in (settings.get("new_gen_not_available") or {}):
                techs = settings["new_gen_not_available"][region]
                for tech in techs:
                    _df = _df.loc[~_df["technology"].str.contains(tech), :]

            df_list.append(_df)

        results = pd.concat(df_list, ignore_index=True, sort=False)

        int_cols = [
            "Fixed_OM_Cost_per_MWyr",
            "Fixed_OM_Cost_per_MWhyr",
            "Inv_Cost_per_MWyr",
            "Inv_Cost_per_MWhyr",
            "cluster",
        ]
        int_cols = [c for c in int_cols if c in results.columns]
        results = results.fillna(0)
        results[int_cols] = results[int_cols].astype(int)
        results["Var_OM_Cost_per_MWh"] = results["Var_OM_Cost_per_MWh"].astype(float)

    return results


def add_renewables_clusters(
    df: pd.DataFrame,
    region: str,
    settings: dict,
    cluster_builder: ClusterBuilder = None,
) -> pd.DataFrame:
    """
    Add renewables clusters

    Parameters
    ----------
    df
        New generation technologies.
            - `technology`: NREL ATB technology in the format
                <technology>_<tech_detail>_<cost_case>. Must be unique.
            - `region`: Model region.
    region
        Model region.
    settings
        Dictionary with the following keys:
            - `renewables_clusters`: Determines the clusters built for the region.
            - `region_aggregations`: Maps the model region to IPM regions.
    cluster_builder
        ClusterBuilder object. Reuse to save time. None by default.

    Returns
    -------
    pd.DataFrame
        Copy of the input dataframe joined to rows for renewables clusters
        on matching NREL ATB technology and model region.

    Raises
    ------
    ValueError
        NREL ATB technologies are not unique.
    ValueError
        Renewables clusters do not match NREL ATB technologies.
    ValueError
        Renewables clusters match multiple NREL ATB technologies.
    """
    if not df["technology"].is_unique:
        raise ValueError(
            f"NREL ATB technologies are not unique: {df['technology'].to_list()}"
        )
    atb_map = {
        x: map_nrel_atb_technology(x.split("_")[0], x.split("_")[1])
        for x in df["technology"]
    }
    mask = df["technology"].isin([tech for tech, match in atb_map.items() if match]) & (
        df["region"] == region
    )
    cdfs = []
    if region in (settings.get("region_aggregations", {}) or {}):
        ipm_regions = settings.get("region_aggregations", {})[region]
        ipm_regions.append(region)  # Add model region, sometimes listed in RG file
    else:
        ipm_regions = [region]
    for scenario in settings.get("renewables_clusters", []) or []:
        if scenario["region"] != region:
            continue
        # Match cluster technology to NREL ATB technologies
        technologies = [
            k
            for k, v in atb_map.items()
            if v and all([scenario.get(ki) == vi for ki, vi in v.items()])
        ]
        if not technologies:
            raise ValueError(
                f"Renewables clusters do not match NREL ATB technologies: {scenario}"
            )
        if len(technologies) > 1:
            raise ValueError(
                f"Renewables clusters match multiple NREL ATB technologies: {scenario}"
            )
        technology = technologies[0]
        # region not an argument to ClusterBuilder.get_clusters()
        scenario = scenario.copy()
        scenario.pop("region")
        if not cluster_builder:
            cluster_builder = build_resource_clusters(settings.get("RESOURCE_GROUPS"))
        clusters = (
            cluster_builder.get_clusters(
                **scenario,
                ipm_regions=ipm_regions,
                existing=False,
                utc_offset=settings.get("utc_offset", 0),
            )
            .rename(columns={"mw": "Max_Cap_MW"})
            .assign(technology=technology, region=region)
        )
        clusters["cluster"] = range(1, 1 + len(clusters))
        if scenario.get("min_capacity"):
            # Warn if total capacity less than expected
            capacity = clusters["Max_Cap_MW"].sum()
            if capacity < scenario["min_capacity"]:
                logger.warning(
                    f"Selected technology {scenario['technology']} capacity"
                    + f" in region {region}"
                    + f" less than minimum ({capacity} < {scenario['min_capacity']} MW)"
                )
        row = df[df["technology"] == technology].to_dict("records")[0]
        new_tech_name = "_".join(
            [
                str(v)
                for k, v in scenario.items()
                if k not in ["region", "technology", "max_clusters", "min_capacity"]
            ]
        )
        clusters["technology"] = clusters["technology"] + "_" + new_tech_name
        kwargs = {k: v for k, v in row.items() if k not in clusters}
        cdfs.append(clusters.assign(**kwargs))
    return pd.concat([df[~mask]] + cdfs, sort=False)


def load_user_defined_techs(settings: dict) -> pd.DataFrame:
    """Load user-defined technologies from a CSV file. Returns cost columns and heat
    rate.

    Parameters
    ----------
    settings : dict
        User-defined parameters from a settings file. It must have the key
        'additional_technologies_fn'. The value can either be a string (name of a single
        file) or a dictionary. If the value is a dictionary it should have integer keys
        corresponding to model years and corresponding string values (file name).

        settings['additional_technologies_fn'] = 'user_techs.csv'
        OR
        settings['additional_technologies_fn'] = {
            2030: 'user_techs_2030.csv',
            2045: 'user_techs_2045.csv'
        }

    Returns
    -------
    pd.DataFrame
        A dataframe of user-defined resources with cost and heat rate columns.
    """
    if isinstance(settings["additional_technologies_fn"], collections.abc.Mapping):
        fn = settings["additional_technologies_fn"][settings["model_year"]]
    else:
        fn = settings["additional_technologies_fn"]

    # Search the extra inputs folder first, then the legacy additional_techs folder
    # in repo
    if (Path(settings["input_folder"]) / fn).exists():
        user_techs = pd.read_csv(Path(settings["input_folder"]) / fn)
    else:
        logger.warning(
            "The file with your user defined technologies is not in the user input "
            "folder. Reading the file from PowerGenome/data/additional_technolgies "
            "instead. This may be depreciated in a future version, please move "
            f"{fn} to the folder {settings['input_folder']}."
        )
        user_techs = pd.read_csv(DATA_PATHS["additional_techs"] / fn)

    user_techs = user_techs.loc[
        (user_techs["technology"].isin(settings["additional_new_gen"]))
        & (user_techs["planning_year"] == settings["model_year"]),
        :,
    ]

    user_techs = user_techs.fillna(0)

    if "tech_detail" not in user_techs.columns:
        user_techs["tech_detail"] = ""
    if "cost_case" not in user_techs.columns:
        user_techs["cost_case"] = ""
    if "Cap_Size" not in user_techs.columns:
        user_techs["Cap_Size"] = 1

    if "dollar_year" in user_techs.columns:
        for idx, row in user_techs.iterrows():
            for col in [
                "capex_mw",
                "capex_mwh",
                "fixed_o_m_mw",
                "fixed_o_m_mwh",
                "variable_o_m_mwh",
            ]:
                user_techs.loc[idx, col] = inflation_price_adjustment(
                    row[col], row["dollar_year"], settings["target_usd_year"]
                )

    cols = [
        "technology",
        "tech_detail",
        "cost_case",
        "capex_mw",
        "capex_mwh",
        "fixed_o_m_mw",
        "fixed_o_m_mwh",
        "variable_o_m_mwh",
        "wacc_real",
        "heat_rate",
        "Cap_Size",
        "dollar_year",
    ]

    return user_techs[cols]


# powergenome/params.py

"""
Parameters and settings
"""
from typing import Union
import os
from pathlib import Path

from dotenv import find_dotenv, load_dotenv

from powergenome import __file__
from powergenome.resource_clusters import ClusterBuilder

# Not convinced this is the best way to set folder paths but it works!
powergenome_path = Path(__file__).parent
project_path = powergenome_path.parent

load_dotenv(dotenv_path=powergenome_path / ".env")

DATA_PATHS = {}
DATA_PATHS["results"] = project_path / "results"
DATA_PATHS["powergenome"] = project_path / "powergenome"
DATA_PATHS["data"] = project_path / "data"
DATA_PATHS["atb_storage_costs"] = DATA_PATHS["data"] / "NREL_ATB_battery_costs.csv"
DATA_PATHS["ipm_shapefiles"] = DATA_PATHS["data"] / "IPM Regions v617 04-05-17"
DATA_PATHS["tests"] = project_path / "tests"
DATA_PATHS["test_data"] = DATA_PATHS["tests"] / "data"
DATA_PATHS["settings"] = project_path / "settings"
DATA_PATHS["eia"] = DATA_PATHS["data"] / "eia"
DATA_PATHS["eia_860m"] = DATA_PATHS["eia"] / "860m"
DATA_PATHS["cost_multipliers"] = DATA_PATHS["data"] / "cost_multipliers"
DATA_PATHS["cache"] = DATA_PATHS["data"] / "cache"
DATA_PATHS["cache"].mkdir(exist_ok=True)
DATA_PATHS["additional_techs"] = DATA_PATHS["data"] / "additional_technologies"
DATA_PATHS["coal_fgd"] = DATA_PATHS["data"] / "coal_fgd" / "fgd_output.csv"
DATA_PATHS["cpi_data"] = DATA_PATHS["data"] / "cpi_data" / "cpi_data.csv"

IPM_SHAPEFILE_PATH = DATA_PATHS["ipm_shapefiles"] / "IPM_Regions_201770405.shp"
IPM_GEOJSON_PATH = DATA_PATHS["data"] / "ipm_regions_simple.geojson"

SETTINGS = {}
SETTINGS["PUDL_DB"] = os.environ.get("PUDL_DB")
SETTINGS["PG_DB"] = os.environ.get("PG_DB")
SETTINGS["EIA_API_KEY"] = os.environ.get("EIA_API_KEY")
SETTINGS["EFS_DATA"] = os.environ.get("EFS_DATA")
SETTINGS["RESOURCE_GROUPS"] = os.environ.get("RESOURCE_GROUPS")


def build_resource_clusters(group_path: Union[str, Path] = None):
    if not group_path:
        group_path = SETTINGS.get("RESOURCE_GROUPS")
    if not group_path:
        cluster_builder = ClusterBuilder([])
    else:
        cluster_builder = ClusterBuilder.from_json(
            Path(group_path, ".").glob("**/*.json")
        )
    return cluster_builder


# powergenome/price_adjustment.py

"""
Adjust price/cost from one year to another
"""

from functools import lru_cache
import logging

import requests
import json
from typing import NamedTuple, Union
from datetime import date
import pandas as pd
import numpy as np
from pathlib import Path
from powergenome.params import DATA_PATHS


class MonthlyCPI(NamedTuple):
    year: int
    period: int
    value: float


logger = logging.getLogger(__name__)


def get_cpi_data(start_year: int = 1980, end_year: int = None) -> pd.DataFrame:
    """Get monthly consumer price index data from the US BLS API, create annual averages.

    BLS has 2 versions of their API. V1 does not require an API key but it restricts the
    number of API queries per day and only returns monthly data. Since V2 requires a key
    I'm using V1, even though it requires a little extra work to go from monthly to
    annual data.

    According to BLS documentation (https://www.bls.gov/cpi/factsheets/cpi-math-calculations.pdf),
    "Annual averages are the sum of the 12 monthly data points (i.e. indexes), divided
    by 12."

    Parameters
    ----------
    start_year : int, optional
        First year of CPI data in the timeseries, by default 1980
    end_year : int, optional
        Final year of CPI data in the timeseries, by default None. If value is None,
        the current year will be used.

    Returns
    -------
    pd.DataFrame
        Annual averages of BLS CPI for all years from `start_year` to `end_year` (inclusive)
        that have 12 months of data.

    Examples
    --------
    >>> get_cpi_data(start_year=2015, end_year=2020)
        year    period  value
    0   2015    12      237.017
    1   2016    12      240.007167
    2   2017    12      245.119583
    3   2018    12      251.106833
    4   2019    12      255.657417
    5   2020    12      258.811167
    """
    if end_year is None:
        todays_date = date.today()
        end_year = todays_date.year
    headers = {"Content-type": "application/json"}

    df_list = []
    e_y = start_year + 10
    while start_year <= end_year:
        data = json.dumps(
            {
                "seriesid": ["CUUR0000SA0"],
                "startyear": str(start_year),
                "endyear": str(e_y),
            }
        )
        p = requests.post(
            "https://api.bls.gov/publicAPI/v1/timeseries/data/",
            data=data,
            headers=headers,
        )
        json_data = json.loads(p.text)

        data_list = []
        for m_data in json_data["Results"]["series"][0]["data"]:
            monthly_cpi = MonthlyCPI(
                int(m_data["year"]),
                int(m_data["period"].lstrip("M")),
                float(m_data["value"]),
            )
            data_list.append(monthly_cpi)

        m_cpi_df = pd.DataFrame(data_list)
        a_cpi_df = m_cpi_df.groupby("year", as_index=False).agg(
            {"period": "count", "value": "mean"}
        )
        a_cpi_df = a_cpi_df.query("period == 12")
        df_list.append(a_cpi_df)
        start_year = e_y
        e_y = start_year + 10

    annual_cpi = pd.concat(df_list)

    return annual_cpi


@lru_cache()
def load_cpi_data(
    reload_data: bool = False, data_path: Path = None, **kwargs
) -> pd.DataFrame:
    """Load BLS CPI data from CSV file if it exists and `reload_data` is False, otherwise
    get it from the API and write results to file.

    Parameters
    ----------
    reload_data : bool, optional
        If data should be reloaded from the BLS API, by default False
    data_path: Path, option
        Path to CSV file where CPI data is stored or will be saved, by default None. If
        None, use DATA_PATHS["cpi_data"] from the params module.
    **kwargs: optional
        Optional keyword arguments for the function `get_cpi_data`. Only used if
        `reload_data` is True or the CPI data file doesn't already exist.

    Returns
    -------
    pd.DataFrame
        Annual averages of BLS CPI. If data are reloaded then the results will be
    """
    if not data_path:
        data_path = DATA_PATHS["cpi_data"]
    if reload_data or not data_path.exists():
        data_path.parent.mkdir(exist_ok=True)
        start_year = 1980
        end_year = None
        for k, v in kwargs.items():
            if k == "start_year":
                start_year = v
            if k == "end_year":
                end_year = v
        cpi_data = get_cpi_data(start_year, end_year)
        cpi_data.to_csv(data_path, index=False, float_format="%g")
    else:
        cpi_data = pd.read_csv(data_path)

    return cpi_data


def inflation_price_adjustment(
    price: Union[int, float, pd.Series, pd.DataFrame, np.ndarray],
    base_year: int,
    target_year: int,
    **kwargs,
) -> float:
    """Convert costs from one dollar-year to another dollar-year using BLS annual CPI data.

    Parameters
    ----------
    price : Union[int, float, pd.Series, pd.DataFrame, np.ndarray]
        The cost to adjust. Can be a single float or an object compatible with broadcast
        multiplication
    base_year : int
        The original data dollar-year
    target_year : int
        The target dollar-year
    **kwargs : optional
        Optional keyword arguments for the function `load_cpi_data`.

    Returns
    -------
    Union[int, float, pd.Series, pd.DataFrame, np.ndarray]
        Cost data transformed from the base dollar-year to the target dollar_year

    Raises
    ------
    ValueError
        The target dollar year is greater than years available with 12 months of data.
    ValueError
        The base dollar year is lower than the minimum year available.

    Examples
    --------
    >>> p = 10.0
    >>> inflation_price_adjustment(p, 2000, 2010)
    1.2662921022067364
    >>> p = pd.Series([1, 10, 100])
    >>> inflation_price_adjustment(p, 2000, 2010)
    0      1.266292
    1     12.662921
    2    126.629210
    >>> p = pd.DataFrame(data=[[1, 2], [3, 4]], columns=["a", "b"])
    >>> inflation_price_adjustment(p, 2000, 2010)
        a	        b
    0	1.266292	2.532584
    1	3.798876	5.065168
    >>> inflation_price_adjustment(p, 2020, 2050)
    ValueError: CPI data are only available through 2020. Your target year is 2050
    """

    base_year = int(base_year)
    target_year = int(target_year)
    data_path = kwargs.pop("data_path", None)
    reload_data = kwargs.pop("reload_data", False)
    cpi_data = load_cpi_data(reload_data=reload_data, data_path=data_path)
    if cpi_data["year"].max() < max(target_year, base_year):
        logger.info("Updating CPI data")
        kwargs.update({"end_year": target_year})
        cpi_data = load_cpi_data(reload_data=True, data_path=data_path, **kwargs)
        if cpi_data["year"].max() < target_year:
            raise ValueError(
                f"CPI data are only available through {cpi_data['year'].max()}. Your target year is "
                f"{target_year}"
            )
    if cpi_data["year"].min() > base_year:
        logger.info("Updating CPI data")
        kwargs.update({"start_year": base_year})
        cpi_data = load_cpi_data(reload_data=True, data_path=data_path, **kwargs)
        if cpi_data["year"].min() > base_year:
            raise ValueError(
                f"CPI data only start in year {cpi_data['year'].min()}. Your base year is "
                f"{base_year}"
            )
    cpi_data = cpi_data.set_index("year")
    price = price * (
        cpi_data.loc[target_year, "value"] / cpi_data.loc[base_year, "value"]
    )

    return price


# powergenome/resource_clusters.py

import copy
import glob
import json
import os
import re
from typing import Any, Callable, Dict, Iterable, List, Mapping, Optional, Union

import numpy as np
import pandas as pd
import pyarrow
import pyarrow.parquet as pq
import scipy.cluster.hierarchy

CAPACITY = "mw"
MERGE = {
    "sums": [CAPACITY, "area"],
    "means": [
        "lcoe",
        "interconnect_annuity",
        "interconnect_capex_mw",
        "offshore_spur_miles",
        "spur_miles",
        "tx_miles",
        "site_substation_spur_miles",
        "substation_metro_tx_miles",
        "site_metro_spur_miles",
        "m_popden",
    ],
    "weight": CAPACITY,
    "uniques": ["ipm_region", "metro_id"],
}
NREL_ATB_TECHNOLOGY_MAP = {
    ("utilitypv", None): {"technology": "utilitypv"},
    ("landbasedwind", None): {"technology": "landbasedwind"},
    ("offshorewind", None): {"technology": "offshorewind"},
    ("hydropower", None): {"technology": "hydro"},
    **{
        ("offshorewind", f"otrg{x}"): {
            "technology": "offshorewind",
            "turbine_type": "fixed",
        }
        for x in range(1, 8)
    },
    **{
        ("offshorewind", f"class{x}"): {
            "technology": "offshorewind",
            "turbine_type": "fixed",
        }
        for x in range(1, 8)
    },
    **{
        ("offshorewind", f"otrg{x}"): {
            "technology": "offshorewind",
            "turbine_type": "floating",
        }
        for x in range(8, 16)
    },
    **{
        ("offshorewind", f"class{x}"): {
            "technology": "offshorewind",
            "turbine_type": "floating",
        }
        for x in range(8, 16)
    },
}
EIA_TECHNOLOGY_MAP = {
    "conventionalhydroelectric": {"technology": "hydro", "small": False},
    "smallhydroelectric": {"technology": "hydro", "small": True},
    "onshorewindturbine": {"technology": "landbasedwind"},
    "offshorewindturbine": {"technology": "offshorewind"},
    "solarphotovoltaic": {"technology": "utilitypv"},
}


def _normalize(x: Optional[str]) -> Optional[str]:
    """
    Normalize string to lowercase, no whitespace, and no underscores.

    Examples
    --------
    >>> _normalize('Offshore Wind')
    'offshorewind'
    >>> _normalize('OffshoreWind')
    'offshorewind'
    >>> _normalize('Offshore_Wind')
    'offshorewind'
    >>> _normalize(None) is None
    True
    """
    if not x:
        return x
    return re.sub(r"\s+|_", "", x.lower())


def map_nrel_atb_technology(tech: str, detail: str = None) -> Dict[str, Any]:
    """
    Map NREL ATB technology to resource groups.

    Parameters
    ----------
    tech
        Technology.
    detail
        Technology detail.

    Returns
    -------
    dict
        Key, value pairs identifying one or more resource groups.

    Examples
    --------
    >>> map_nrel_atb_technology('UtilityPV', 'LosAngeles')
    {'technology': 'utilitypv'}
    >>> map_nrel_atb_technology('LandbasedWind', 'LTRG1')
    {'technology': 'landbasedwind'}
    >>> map_nrel_atb_technology('OffShoreWind')
    {'technology': 'offshorewind'}
    >>> map_nrel_atb_technology('OffShoreWind', 'OTRG3')
    {'technology': 'offshorewind', 'turbine_type': 'fixed'}
    >>> map_nrel_atb_technology('OffShoreWind', 'OTRG7')
    {'technology': 'offshorewind', 'turbine_type': 'floating'}
    >>> map_nrel_atb_technology('Hydropower')
    {'technology': 'hydro'}
    >>> map_nrel_atb_technology('Hydropower', 'NSD4')
    {'technology': 'hydro'}
    >>> map_nrel_atb_technology('Unknown')
    {}
    """
    tech = _normalize(tech)
    detail = _normalize(detail)
    group = {}
    for k, v in NREL_ATB_TECHNOLOGY_MAP.items():
        if (tech == k[0] or not k[0]) and (detail == k[1] or not k[1]):
            group.update(v)
    return group


def map_eia_technology(tech: str) -> Dict[str, Any]:
    """
    Map EIA technology to resource groups.

    Parameters
    ----------
    tech
        Technology.

    Returns
    -------
    dict
        Key, value pairs identifying one or more resource groups.

    Examples
    --------
    >>> map_eia_technology('Solar Photovoltaic')
    {'technology': 'utilitypv'}
    >>> map_eia_technology('solar_photovoltaic')
    {'technology': 'utilitypv'}
    >>> map_eia_technology('Onshore Wind Turbine')
    {'technology': 'landbasedwind'}
    >>> map_eia_technology('Offshore Wind Turbine')
    {'technology': 'offshorewind'}
    >>> map_eia_technology('Conventional Hydroelectric')
    {'technology': 'hydro', 'small': False}
    >>> map_eia_technology('Small Hydroelectric')
    {'technology': 'hydro', 'small': True}
    >>> map_eia_technology('Unknown')
    {}
    """
    tech = _normalize(tech)
    group = {}
    for k, v in EIA_TECHNOLOGY_MAP.items():
        if tech == k or not k:
            group.update(v)
    return group


class Table:
    """
    Cached interface for tabular data.

    Supports parquet and csv formats.

    Parameters
    ----------
    path
        Path to dataset.
    df
        In-memory dataframe.

    Attributes
    ----------
    path : Union[str, os.PathLike]
        Path to the dataset.
    df : pd.DataFrame
        Cached dataframe.
    format : str
        Dataset format ('parquet' or 'csv'), or `None` if in-memory only.
    columns : list
        Dataset column names.

    Raises
    ------
    ValueError
        Missing either path or dataframe.
    ValueError
        Dataframe columns are not all strings.

    Examples
    --------
    In-memory dataframe:

    >>> df = pd.DataFrame({'id': [1, 2], 'x': [10, 20]})
    >>> table = Table(df = df)
    >>> table.format is None
    True
    >>> table.columns
    ['id', 'x']
    >>> table.read()
       id   x
    0   1  10
    1   2  20
    >>> table.read(columns=['id'])
       id
    0   1
    1   2
    >>> table.clear()
    >>> table.df is not None
    True

    File dataset (csv):

    >>> import tempfile
    >>> fp = tempfile.NamedTemporaryFile()
    >>> df.to_csv(fp.name, index=False)
    >>> table = Table(path = fp.name)
    >>> table.format
    'csv'
    >>> table.columns
    ['id', 'x']
    >>> table.read(cache=False)
       id   x
    0   1  10
    1   2  20
    >>> table.df is None
    True
    >>> table.read(columns=['id'], cache=True)
       id
    0   1
    1   2
    >>> table.df is not None
    True
    >>> table.clear()
    >>> table.df is None
    True
    >>> fp.close()
    """

    def __init__(
        self, path: Union[str, os.PathLike] = None, df: pd.DataFrame = None
    ) -> None:
        self.path = path
        self.df = df
        if df is not None:
            if any(not isinstance(x, str) for x in df.columns):
                raise ValueError("Dataframe columns are not all strings")
        self.format = None
        self._dataset = None
        self._columns = None
        if path is not None:
            try:
                self._dataset = pq.ParquetDataset(path)
                self._columns = self._dataset.schema.names
                self.format = "parquet"
            except pyarrow.lib.ArrowInvalid:
                # Assume CSV file
                self.format = "csv"
        if path is None and df is None:
            raise ValueError("Mising either path to tabular data or a pandas DataFrame")

    @property
    def columns(self) -> list:
        if self.df is not None:
            return list(self.df.columns)
        if self._columns is None:
            if self.format == "csv":
                self._columns = pd.read_csv(self.path, nrows=0).columns
        return list(self._columns)

    def read(self, columns: Iterable = None, cache: bool = None) -> pd.DataFrame:
        """
        Read data from memory or from disk.

        Parameters
        ----------
        columns
            Names of column to read. If `None`, all columns are read.
        cache
            Whether to cache the full dataset in memory. If `None`,
            the dataset is cached if `columns` is `None`, and not otherwise.

        Returns
        -------
        pd.DataFrame
            Data as a dataframe.
        """
        if self.df is not None:
            return self.df[columns] if columns is not None else self.df
        if cache is None:
            cache = columns is None
        read_columns = None if cache else columns
        if self.format == "csv":
            df = pd.read_csv(self.path, usecols=read_columns, dtype={"metro_id": str})
        elif self.format == "parquet":
            df = self._dataset.read(columns=read_columns).to_pandas()
        if cache:
            self.df = df
        return df[columns] if columns is not None else df

    def clear(self) -> None:
        """
        Clear the dataset cache.

        Only applies if :attr:`path` is set so that the dataset can be reread from file.
        """
        if self.path is not None:
            self.df = None


class ResourceGroup:
    """
    Group of resources sharing common attributes.

    Parameters
    ----------
    group
        Group metadata.

        - `technology` : str
          Resource type ('utilitypv', 'landbasedwind', or 'offshorewind').
        - `existing` : bool
          Whether resources are new (`False`, default) or existing (`True`).
        - `tree` : str, optional
          The name of the resource metadata attribute by
          which to differentiate between multiple precomputed hierarchical trees.
          Defaults to `None` (resource group does not represent hierarchical trees).
        - `metadata` : str, optional
          Relative path to resource metadata dataset (optional if `metadata` is `None`).
        - `profiles` : str, optional
          Relative path to resource profiles dataset.
        - ... and any additional (optional) keys.

    metadata
        Resource metadata, with one resource per row.

        - `id`: int
          Resource identifier, unique within the group.
        - `ipm_region` : str
          IPM region to which the resource delivers power.
        - `mw` : float
          Maximum resource capacity in MW.
        - `lcoe` : float, optional
          Levelized cost of energy, used to guide the selection
          (from lowest to highest) and clustering (by nearest) of resources.
          If missing, selection and clustering is by largest and nearest `mw`.

        Resources representing hierarchical trees (see `group.tree`)
        require additional attributes.

        - `parent_id` : int
          Identifier of the resource formed by clustering this resource with the one
          other resource with the same `parent_id`.
          Only resources with `level` of 1 have no `parent_id`.
        - `level` : int
          Level of tree where the resource first appears, from `m`
          (the number of resources at the base of the tree), to 1.
        - `[group.tree]` : Any
          Each unique value of this grouping attribute represents a precomputed
          hierarchical tree. When clustering resources, every tree is traversed to its
          crown before the singleton resources from the trees are clustered together.

        The following resource attributes (all float) are propagaged as:

        - weighted means (weighted by `mw`):

            - `lcoe`
            - `interconnect_annuity`
            - `tx_miles`
            - `spur_miles`
            - `offshore_spur_miles`
            - `site_substation_spur_miles`
            - `substation_metro_tx_miles`
            - `site_metro_spur_miles`

        - sums:

            - `mw`
            - `area`

        - uniques:

            - `ipm_region`
            - `metro_id`

    profiles
        Variable resource capacity profiles with normalized capacity factors
        (from 0 to 1) for every hour of the year (either 8760 or 8784 for a leap year).
        Each profile must be a column whose name matches the resource `metadata.id`.
    path
        Directory relative to which the file paths `group.metadata` and `group.profiles`
        should be read.

    Attributes
    ----------
    group : Dict[str, Any]
    metadata : Table
        Cached interface to resource metadata.
    profiles : Optional[Table]
        Cached interface to resource profiles.

    Examples
    --------
    >>> group = {'technology': 'utilitypv'}
    >>> metadata = pd.DataFrame({'id': [0, 1], 'ipm_region': ['A', 'A'], 'mw': [1, 2]})
    >>> profiles = pd.DataFrame({'0': np.full(8784, 0.1), '1': np.full(8784, 0.4)})
    >>> rg = ResourceGroup(group, metadata, profiles)
    >>> rg.test_metadata()
    >>> rg.test_profiles()
    >>> rg.get_clusters(max_clusters=1)
           ipm_region  mw                                            profile
    (1, 0)          A   3  [0.3, 0.3, 0.3, 0.3, 0.3, 0.3, 0.3, 0.3, 0.3, ...
    """

    def __init__(
        self,
        group: Dict[str, Any],
        metadata: pd.DataFrame = None,
        profiles: pd.DataFrame = None,
        path: str = ".",
    ) -> None:
        self.group = {"existing": False, "tree": None, **group.copy()}
        for key in ["metadata", "profiles"]:
            if self.group.get(key):
                # Convert relative paths (relative to group file) to absolute paths
                self.group[key] = os.path.abspath(os.path.join(path, self.group[key]))
        required = ["technology"]
        if metadata is None:
            required.append("metadata")
        missing = [key for key in required if not self.group.get(key)]
        if missing:
            raise ValueError(
                f"Group metadata missing required keys {missing}: {self.group}"
            )
        self.metadata = Table(df=metadata, path=self.group.get("metadata"))
        self.profiles = None
        if profiles is not None or self.group.get("profiles"):
            self.profiles = Table(df=profiles, path=self.group.get("profiles"))

    @classmethod
    def from_json(cls, path: Union[str, os.PathLike]) -> "ResourceGroup":
        """
        Build from JSON file.

        Parameters
        ----------
        path
            Path to JSON file.
        """
        with open(path, mode="r") as fp:
            group = json.load(fp)
        return cls(group, path=os.path.dirname(path))

    def test_metadata(self) -> None:
        """
        Test that `:attr:metadata` is valid.

        Raises
        ------
        ValueError
            Resource metadata missing required keys.
        """
        columns = self.metadata.columns
        required = ["ipm_region", "id", "mw"]
        if self.group.get("tree"):
            required.extend(["parent_id", "level", self.group["tree"]])
        missing = [key for key in required if key not in columns]
        if missing:
            raise ValueError(f"Resource metadata missing required keys {missing}")

    def test_profiles(self) -> None:
        """
        Test that `:attr:profiles` is valid.

        Raises
        ------
        ValueError
            Resource profiles column names do not match resource identifiers.
        ValueError
            Resource profiles are not either 8760 or 8784 elements.
        """
        if self.profiles is None:
            return None
        # Cast identifiers to string to match profile columns
        ids = self.metadata.read(columns=["id"])["id"].astype(str)
        columns = self.profiles.columns
        if not set(columns) == set(ids):
            raise ValueError(
                f"Resource profiles column names do not match resource identifiers"
            )
        df = self.profiles.read(columns=columns[0])
        if len(df) not in [8760, 8784]:
            raise ValueError(f"Resource profiles are not either 8760 or 8784 elements")

    def get_clusters(
        self,
        ipm_regions: Iterable[str] = None,
        min_capacity: float = None,
        max_clusters: int = None,
        max_lcoe: float = None,
        cap_multiplier: float = None,
        profiles: bool = True,
        utc_offset: int = 0,
        sub_region=None,
    ) -> pd.DataFrame:
        """
        Compute resource clusters.

        Parameters
        ----------
        ipm_regions
            IPM regions in which to select resources.
            If `None`, all IPM regions are selected.
        min_capacity
            Minimum total capacity (MW). Resources are selected,
            from lowest to highest levelized cost of energy (lcoe),
            or from highest to lowest capacity if lcoe not available,
            until the minimum capacity is just exceeded.
            If `None`, all resources are selected for clustering.
        max_clusters
            Maximum number of resource clusters to compute.
            If `None`, no clustering is performed; resources are returned unchanged.
        max_lcoe
            Select only the resources with a levelized cost of electricity (lcoe)
            below this maximum. Takes precedence over `min_capacity`.
        cap_multiplier
            Multiplier applied to resource capacity before selection by `min_capacity`.
        profiles
            Whether to include cluster profiles, if available, in column `profile`.

        Returns
        -------
        pd.DataFrame
            Clustered resources whose indices are tuples of the resource identifiers
            from which they were constructed.

        Raises
        ------
        ValueError
            No resources found or selected.
        """
        df = self.metadata.read().set_index("id")
        if ipm_regions is not None:
            # Filter by IPM region
            df = df[df["ipm_region"].isin(ipm_regions)]
        if sub_region is not None:
            df = df.loc[df[self.group["sub_region"]] == sub_region, :]
        if cap_multiplier is not None:
            # Apply capacity multiplier
            df[CAPACITY] *= cap_multiplier

        if df.empty:
            drop_keys = ["tree", "metadata", "site_cluster", "profiles"]
            group_info = {k: v for k, v in self.group.items() if k not in drop_keys}
            raise KeyError(
                f"""
                No resources for the group
                {group_info}
                were found in the model region containing IPM Regions
                {ipm_regions}
            """
            )
        # Sort resources by lcoe (ascending) or capacity (descending)
        by = "lcoe" if "lcoe" in df else CAPACITY
        df = df.sort_values(by, ascending=by == "lcoe")
        # Select base resources
        tree = self.group["tree"]
        if tree:
            idx = df.index
            df = df.set_index(tree)
            max_level = df.index.map(df.groupby(tree)["level"].max())
            df = df.reset_index()
            df.index = idx
            base = (df["level"] == max_level).values
            mask = base.copy()
        else:
            mask = np.ones(len(df), dtype=bool)
        if min_capacity:
            # Select resources until min_capacity reached
            temp = (df.loc[mask, CAPACITY].cumsum() < min_capacity).values
            temp[temp.argmin()] = True
            mask[mask] = temp
        if max_lcoe and "lcoe" in df:
            # Select clusters with LCOE below the cutoff
            mask[mask] = df.loc[mask, "lcoe"] <= max_lcoe
        if not mask.any():
            raise ValueError(f"No resources found or selected")
        if tree:
            # Only keep trees with one ore more base resources
            if isinstance(tree, list):
                df["tree"] = ""
                for t in tree:
                    df["tree"] += df[t].astype(str)
            else:
                df["tree"] = df[tree]
            selected = (
                pd.Series(mask, index=df.index)
                .groupby(df["tree"])
                .transform(lambda x: x.sum() > 0)
            )
            # Add non-base resources to selected trees
            mask |= selected & ~base
        # Apply mask
        df = df[mask]
        # Prepare merge
        merge = copy.deepcopy(MERGE)
        # Prepare profiles
        if profiles and self.profiles is not None:
            df["profile"] = list(
                np.roll(
                    self.profiles.read(columns=df.index.astype(str)).values.T,
                    utc_offset,
                )
            )
            merge["means"].append("profile")
        # Compute clusters
        if tree:
            return cluster_trees(df, by=by, tree=tree, max_rows=max_clusters, **merge)
        return cluster_rows(df, by=df[[by]], max_rows=max_clusters, **merge)


class ClusterBuilder:
    """
    Builds clusters of resources.

    Parameters
    ----------
    groups
        Groups of resources. See :class:`ResourceGroup`.

    Attributes
    ----------
    groups : Iterable[ResourceGroup]

    Examples
    --------
    Prepare the resource groups.

    >>> groups = []
    >>> group = {'technology': 'utilitypv'}
    >>> metadata = pd.DataFrame({'id': [0, 1], 'ipm_region': ['A', 'A'], 'mw': [1, 2]})
    >>> profiles = pd.DataFrame({'0': np.full(8784, 0.1), '1': np.full(8784, 0.4)})
    >>> groups.append(ResourceGroup(group, metadata, profiles))
    >>> group = {'technology': 'utilitypv', 'existing': True}
    >>> metadata = pd.DataFrame({'id': [0, 1], 'ipm_region': ['B', 'B'], 'mw': [1, 2]})
    >>> profiles = pd.DataFrame({'0': np.full(8784, 0.1), '1': np.full(8784, 0.4)})
    >>> groups.append(ResourceGroup(group, metadata, profiles))
    >>> builder = ClusterBuilder(groups)

    Compute resource clusters.

    >>> builder.get_clusters(ipm_regions=['A'], max_clusters=1,
    ...     technology='utilitypv', existing=False)
          ids ipm_region  mw  ...         profile technology  existing
    0  (1, 0)          A   3  [0.3, 0.3, 0.3, ...  utilitypv     False
    >>> builder.get_clusters(ipm_regions=['B'], min_capacity=2,
    ...     technology='utilitypv', existing=True)
        ids ipm_region  mw  ...         profile technology  existing
    0  (1,)          B   2  [0.4, 0.4, 0.4, ...  utilitypv      True

    Errors arise if search criteria is either ambiguous or results in an empty result.

    >>> builder.get_clusters(ipm_regions=['A'], technology='utilitypv')
    Traceback (most recent call last):
      ...
    ValueError: Parameters match multiple resource groups: [{...}, {...}]
    >>> builder.get_clusters(ipm_regions=['B'], technology='utilitypv', existing=False)
    Traceback (most recent call last):
      ...
    ValueError: No resources found or selected
    """

    def __init__(self, groups: Iterable[ResourceGroup]) -> None:
        self.groups = groups

    @classmethod
    def from_json(cls, paths: Iterable[Union[str, os.PathLike]]) -> "ClusterBuilder":
        """
        Load resources from resource group JSON files.

        Parameters
        ----------
        paths
            Paths to resource group JSON files.

        Raises
        ------
        ValueError
            No resource groups specified.
        """
        paths = list(paths)
        if not paths:
            raise ValueError(f"No resource groups specified {paths}")
        return cls([ResourceGroup.from_json(path) for path in paths])

    def find_groups(self, **kwargs: Any) -> List[ResourceGroup]:
        """
        Return the resource groups matching the specified arguments.

        Parameters
        ----------
        **kwargs
            Parameters to match against resource group metadata.
        """
        return [
            rg
            for rg in self.groups
            if all(
                k in rg.group and rg.group[k] == v
                for k, v in kwargs.items()
                if k not in ["sub_region", "extra_uniques"]
            )
        ]

    def get_clusters(
        self,
        ipm_regions: Iterable[str] = None,
        min_capacity: float = None,
        max_clusters: int = None,
        max_lcoe: float = None,
        cap_multiplier: float = None,
        utc_offset: int = 0,
        **kwargs: Any,
    ) -> pd.DataFrame:
        """
        Compute resource clusters.

        See :meth:`ResourceGroup.get_clusters` for parameter descriptions.

        The following fields are added:

        - `ids` (tuple): Original resource identifiers.
        - **kwargs: Parameters used to uniquely identify the group.

        Parameters
        ----------
        ipm_regions
        min_capacity
        max_clusters
        max_lcoe
        cap_multiplier
        **kwargs
            Parameters to :meth:`find_groups` for selecting the resource group.

        Raises
        ------
        ValueError
            Parameters do not match any resource groups.
        ValueError
            Parameters match multiple resource groups.
        """
        groups = self.find_groups(**kwargs)
        if not groups:
            raise ValueError(f"Parameters do not match any resource groups: {kwargs}")
        if len(groups) > 1:
            meta = [rg.group for rg in groups]
            raise ValueError(f"Parameters match multiple resource groups: {meta}")
        return (
            groups[0]
            .get_clusters(
                ipm_regions=ipm_regions,
                min_capacity=min_capacity,
                max_clusters=max_clusters,
                max_lcoe=max_lcoe,
                cap_multiplier=cap_multiplier,
                utc_offset=utc_offset,
                sub_region=kwargs.get("sub_region"),
            )
            .assign(**kwargs)
            .rename_axis("ids")
            .reset_index()
        )


def _tuple(x: Any) -> tuple:
    """
    Cast object to tuple.

    Examples
    --------
    >>> _tuple(1)
    (1,)
    >>> _tuple([1])
    (1,)
    >>> _tuple('string')
    ('string',)
    """
    if np.iterable(x) and not isinstance(x, str):
        return tuple(x)
    return (x,)


def merge_row_pair(
    a: Mapping,
    b: Mapping,
    sums: Iterable = None,
    means: Iterable = None,
    weight: Any = None,
    uniques: Iterable = None,
) -> dict:
    """
    Merge two mappings into one.

    Parameters
    ----------
    a
        First mapping (e.g. :class:`dict`, :class:`pd.Series`).
    b
        Second mapping.
    means
        Keys of values to average.
    weight
        Key of values to use as weights for weighted averages.
        If `None`, averages are not weighted.
    uniques
        Keys of values for which to return the value if equal, and `None` if not.

    Returns
    -------
    dict
        Merged row as a dictionary.

    Examples
    --------
    >>> df = pd.DataFrame({'mw': [1, 2], 'area': [10, 20], 'lcoe': [0.1, 0.4]})
    >>> a, b = df.to_dict('records')
    >>> merge_row_pair(a, b, sums=['area', 'mw'], means=['lcoe'], weight='mw')
    {'area': 30, 'mw': 3, 'lcoe': 0.3}
    >>> merge_row_pair(a, b, sums=['area', 'mw'], means=['lcoe'])
    {'area': 30, 'mw': 3, 'lcoe': 0.25}
    >>> b['mw'] = 1
    >>> merge_row_pair(a, b, uniques=['mw', 'area'])
    {'mw': 1, 'area': None}
    """
    merge = {}
    if sums:
        for key in sums:
            merge[key] = a[key] + b[key]
    if means:
        if weight:
            total = a[weight] + b[weight]
            aw = a[weight] / total
            bw = b[weight] / total
        else:
            aw = 0.5
            bw = 0.5
        for key in means:
            merge[key] = a[key] * aw + b[key] * bw
    if uniques:
        for key in uniques:
            merge[key] = a[key] if a[key] == b[key] else None
    return merge


def cluster_rows(
    df: pd.DataFrame, by: Iterable[Iterable], max_rows: int = None, **kwargs: Any
) -> pd.DataFrame:
    """
    Merge rows in dataframe by hierarchical clustering.

    Uses the Ward variance minimization algorithm to incrementally merge rows.
    See :func:`scipy.cluster.hierarchy.linkage`.

    Parameters
    ----------
    df
        Rows to merge (m, ...).
    by
        2-dimensional array of observation vectors (m, ...) from which to compute
        distances between each row pair.
    max_rows
        Number of rows at which to stop merging rows.
        If `None`, no clustering is performed.
    **kwargs
        Optional parameters to :func:`merge_row_pair`.

    Returns
    -------
    pd.DataFrame
        Merged rows as a dataframe.
        Their indices are tuples of the original row indices from which they were built.
        If original indices were already iterables, they are merged
        (e.g. (1, 2) and (3, ) becomes (1, 2, 3)).

    Raises
    ------
    ValueError
        Max number of rows must be greater than zero.

    Examples
    --------
    With the default (range) row index:

    >>> df = pd.DataFrame({'mw': [1, 2, 3], 'area': [4, 5, 6], 'lcoe': [0.1, 0.4, 0.2]})
    >>> kwargs = {'sums': ['mw', 'area'], 'means': ['lcoe'], 'weight': 'mw'}
    >>> cluster_rows(df, by=df[['lcoe']], **kwargs)
          mw  area  lcoe
    (0,)   1     4   0.1
    (1,)   2     5   0.4
    (2,)   3     6   0.2
    >>> cluster_rows(df, by=df[['lcoe']], max_rows=2, **kwargs)
            mw  area   lcoe
    (1,)     2     5  0.400
    (0, 2)   4    10  0.175

    With a custom row index:

    >>> df.index = ['a', 'b', 'c']
    >>> cluster_rows(df, by=df[['lcoe']], max_rows=2, **kwargs)
            mw  area   lcoe
    (b,)     2     5  0.400
    (a, c)   4    10  0.175

    With an iterable row index:

    >>> df.index = [(1, 2), (4, ), (3, )]
    >>> cluster_rows(df, by=df[['lcoe']], max_rows=2, **kwargs)
               mw  area   lcoe
    (4,)        2     5  0.400
    (1, 2, 3)   4    10  0.175
    """
    nrows = len(df)
    if max_rows is None:
        max_rows = len(df)
    elif max_rows < 1:
        raise ValueError("Max number of rows must be greater than zero")
    drows = nrows - max_rows
    index = [_tuple(x) for x in df.index] + [None] * drows
    merge = prepare_merge(kwargs, df)
    df = df[get_merge_columns(merge, df)].reset_index(drop=True)
    if drows < 1:
        df.index = index
        return df
    # Convert dataframe rows to dictionaries
    rows = df.to_dict("records")
    # Preallocate new rows
    rows += [None] * drows
    # Preallocate new rows
    Z = scipy.cluster.hierarchy.ward(by)
    n = nrows + drows
    mask = np.ones(n, dtype=bool)
    for i, link in enumerate(Z[:drows, 0:2].astype(int)):
        mask[link] = False
        pid = nrows + i
        rows[pid] = merge_row_pair(rows[link[0]], rows[link[1]], **merge)
        index[pid] = index[link[0]] + index[link[1]]
    clusters = pd.DataFrame([x for x, m in zip(rows, mask) if m])
    # Preserve original column order
    clusters = clusters[[x for x in df.columns if x in clusters]]
    clusters.index = [x for x, m in zip(index, mask) if m]
    return clusters


def build_tree(
    df: pd.DataFrame, by: Iterable[Iterable], max_level: int = None, **kwargs: Any
) -> pd.DataFrame:
    """
    Build a hierarchical tree of rows in a dataframe.

    Uses the Ward variance minimization algorithm to incrementally merge rows.
    See :func:`scipy.cluster.hierarchy.linkage`.

    Parameters
    ----------
    df
        Rows to merge (m, ...).
        Should not have columns `id`, `parent_id`, and `level`, as these are appended to
        the result dataframe.
    by
        2-dimensional array of observation vectors (m, ...) from which to compute
        distances between each row pair.
    max_level
        Maximum level of tree to return,
        from m (the number of rows in `df`, if `None`) to 1.
    **kwargs
        Optional parameters to :func:`merge_row_pair`.

    Returns
    -------
    pd.DataFrame
        Hierarchical tree as a dataframe.
        Row indices are tuples of the original row indices from which they were built.
        If original indices were already iterables, they are merged
        (e.g. (1, 2) and (3, ) becomes (1, 2, 3)).
        The following columns are added:

        - `id` (int): New row identifier (0, ..., 0 + n).
        - `parent_id` (Int64): New row identifer of parent row.
        - `level` (int): Tree level of row (max_level, ..., 1).

    Raises
    ------
    ValueError
        Max level of tree must be greater than zero.

    Examples
    --------
    >>> df = pd.DataFrame({'mw': [1, 2, 3], 'area': [4, 5, 6], 'lcoe': [0.1, 0.4, 0.2]})
    >>> kwargs = {'sums': ['area', 'mw'], 'means': ['lcoe'], 'weight': 'mw'}
    >>> build_tree(df, by=df[['lcoe']], **kwargs)
               mw  area   lcoe  id  parent_id  level
    (0,)        1     4  0.100   0          3      3
    (1,)        2     5  0.400   1          4      3
    (2,)        3     6  0.200   2          3      3
    (0, 2)      4    10  0.175   3          4      2
    (1, 0, 2)   6    15  0.250   4        NaN      1
    >>> build_tree(df, by=df[['lcoe']], max_level=2, **kwargs)
               mw  area   lcoe  id  parent_id  level
    (1,)        2     5  0.400   0          2      2
    (0, 2)      4    10  0.175   1          2      2
    (1, 0, 2)   6    15  0.250   2        NaN      1
    >>> build_tree(df, by=df[['lcoe']], max_level=1, **kwargs)
               mw  area  lcoe  id  parent_id  level
    (1, 0, 2)   6    15  0.25   0        NaN      1
    """
    nrows = len(df)
    if max_level is None:
        max_level = nrows
    else:
        max_level = min(max_level, nrows)
        if max_level < 1:
            raise ValueError("Max level of tree must be greater than zero")
    drows = nrows - 1
    index = [_tuple(x) for x in df.index] + [None] * drows
    df = df.reset_index(drop=True)
    merge = prepare_merge(kwargs, df)
    columns = get_merge_columns(merge, df)
    df = df[columns]
    if drows < 1:
        df.index = index
        return df
    # Convert dataframe rows to dictionaries
    rows = df.to_dict("records")
    # Preallocate new rows
    rows += [None] * drows
    Z = scipy.cluster.hierarchy.linkage(by, method="ward")
    n = nrows + drows
    mask = np.ones(n, dtype=bool)
    level = np.concatenate((np.full(nrows, nrows), np.arange(drows, 0, -1)))
    parent_id = np.zeros(n)
    drop = nrows - max_level
    for i, link in enumerate(Z[:, 0:2].astype(int)):
        if i < drop:
            mask[link] = False
        pid = nrows + i
        parent_id[link] = pid
        rows[pid] = merge_row_pair(rows[link[0]], rows[link[1]], **merge)
        index[pid] = index[link[0]] + index[link[1]]
    tree = pd.DataFrame([x for x, m in zip(rows, mask) if m])
    # Restore original column order
    tree = tree[columns]
    # Normalize ids to 0, ..., n
    old_ids = np.where(mask)[0]
    new_ids = np.arange(len(old_ids))
    new_parent_ids = pd.Series(np.searchsorted(old_ids, parent_id[mask]), dtype="Int64")
    new_parent_ids.iloc[-1] = np.nan
    # Bump lower levels to max_level
    level = level[mask]
    if max_level < nrows:
        stop = level.size - np.searchsorted(level[::-1], max_level, side="right")
        level[:stop] = max_level
    tree = tree.assign(id=new_ids, parent_id=new_parent_ids, level=level)
    tree.index = [x for x, m in zip(index, mask) if m]
    return tree


def cluster_trees(
    df: pd.DataFrame, by: str, tree: str = None, max_rows: int = None, **kwargs: Any
) -> pd.DataFrame:
    """
    Merge rows in a dataframe following precomputed hierarchical trees.

    Parameters
    ----------
    df
        Rows to merge.
        Must have columns `parent_id` (matching values in index), `level`, and
        the columns named in **by** and **tree**.
    by
        Name of column to use for determining merge order.
        Children with the smallest pairwise distance on this column are merged first.
    tree
        Name of column to use for differentiating between hierarchical trees.
        If `None`, assumes rows represent a single tree.
    max_rows
        Number of rows at which to stop merging rows.
        If smaller than the number of trees, :func:`cluster_rows` is used to merge
        tree heads.
        If `None`, no merging is performed and only the base rows are returned.
    **kwargs
        Optional parameters to :func:`merge_row_pair`.

    Returns
    -------
    pd.DataFrame
        Merged rows as a dataframe.
        Their indices are tuples of the original row indices from which they were built.
        If original indices were already iterables, they are merged
        (e.g. (1, 2) and (3, ) becomes (1, 2, 3)).

    Raises
    ------
    ValueError
        Max number of rows must be greater than zero.
    ValueError
        Missing required fields.
    ValueError
        `by` column not included in row merge arguments (`kwargs`).

    Examples
    --------
    >>> df = pd.DataFrame({
    ...     'level': [3, 3, 3, 2, 1],
    ...     'parent_id': pd.Series([3, 3, 4, 4, float('nan')], dtype='Int64'),
    ...     'mw': [0.1, 0.1, 0.1, 0.2, 0.3],
    ...     'area': [1, 1, 1, 2, 3]
    ... }, index=[0, 1, 2, 3, 4])
    >>> cluster_trees(df, by='mw', sums=['mw', 'area'], max_rows=2)
             mw  area
    (2,)    0.1     1
    (0, 1)  0.2     2
    >>> cluster_trees(df, by='mw', sums=['mw'], max_rows=1)
                mw
    (2, 0, 1)  0.3
    >>> cluster_trees(df, by='mw', sums=['mw'])
           mw
    (0,)  0.1
    (1,)  0.1
    (2,)  0.1
    """
    required = ["parent_id", "level", by]
    if tree:
        if isinstance(tree, list):
            required.extend(tree)
        else:
            required.append(tree)
    missing = [key for key in required if key not in df]
    if missing:
        raise ValueError(f"Missing required fields {missing}")
    if tree:
        idx = df.index
        df = df.set_index(tree)
        mask = df["level"] == df.index.map(df.groupby(tree)["level"].max())
        mask.index = idx
        df = df.reset_index()
        df.index = idx
    else:
        mask = df["level"] == df["level"].max()
    nrows = mask.sum()
    if max_rows is None:
        max_rows = nrows
    elif max_rows < 1:
        raise ValueError("Max number of rows must be greater than zero")
    merge = prepare_merge(kwargs, df)
    columns = get_merge_columns(merge, df)
    if by not in columns:
        raise ValueError(f"{by} not included in row merge arguments")
    drows = nrows - max_rows
    if drows < 1:
        df = df.loc[mask, columns].copy()
        df.index = [_tuple(x) for x in df.index]
        return df
    df = df[set(columns + required)].assign(
        _id=df.index, _ids=[_tuple(x) for x in df.index], _mask=mask
    )
    diff = lambda x: abs(x.max() - x.min())
    while drows > 0:
        # Sort parents by ascending distance of children
        # NOTE: Inefficient to recompute for all parents every time
        parents = (
            df[df["_mask"]]
            .groupby("parent_id", sort=False)
            .agg(ids=("_id", list), n=("_id", "count"), distance=(by, diff))
            .sort_values(["n", "distance"], ascending=[False, True])
        )
        if parents.empty:
            break
        if parents["n"].iloc[0] == 2:
            # Choose complete parent with lowest distance of children
            pid = parents.index[0]
            ids = parents["ids"].iloc[0]
            children = df.loc[ids].to_dict("records")
            # Compute parent
            parent = {
                # Initial attributes
                # Can access series because all columns integer
                **df.loc[pid, ["_id", "parent_id", "level"]],
                # Merged children attributes
                # NOTE: Needed only if a child is incomplete
                **merge_row_pair(children[0], children[1], **merge),
                # Indices of all past children
                "_ids": df.loc[ids[0], "_ids"] + df.loc[ids[1], "_ids"],
                "_mask": True,
            }
            # Add parent
            df.loc[pid] = pd.Series(parent, dtype=object)
            # Drop children
            df.loc[ids, "_mask"] = False
            # Decrement rows
            drows -= 1
        else:
            # Promote child with deepest parent
            parent_id = df.loc[parents.index, "level"].idxmax()
            child_id = parents.loc[parent_id, "ids"][0]
            # Update child
            tree_columns = ["_id", "parent_id", "level"]
            df.loc[child_id, tree_columns] = df.loc[parent_id, tree_columns]
            # Update index
            df.rename(index={child_id: parent_id, parent_id: np.nan}, inplace=True)
    # Apply mask
    df = df[df["_mask"]]
    # Drop temporary columns
    df.index = df["_ids"].values
    df = df.drop(columns=["_id", "_ids", "_mask"])
    if len(df) > max_rows:
        df = cluster_rows(df, by=df[[by]], max_rows=max_rows, **kwargs)
    return df[columns]


def group_rows(
    df: pd.DataFrame, ids: Iterable[Iterable]
) -> pd.core.groupby.DataFrameGroupBy:
    """
    Group dataframe rows by index.

    Parameters
    ----------
    df
        Dataframe to group.
    ids
        Groups of rows indices.

    Returns
    -------
    pd.core.groupby.DataFrameGroupBy
        Rows of `df` grouped by their membership in each index group.

    Examples
    --------
    >>> df = pd.DataFrame({'x': [2, 1, 3]}, index=[2, 1, 3])
    >>> group_rows(df, [(1, ), (2, 3), (1, 2, 3)]).sum()
       x
    0  1
    1  5
    2  6
    """
    groups = np.repeat(np.arange(len(ids)), [len(x) for x in ids])
    index = np.concatenate(ids)
    return df.loc[index].groupby(groups, sort=False)


def prune_tree(df: pd.DataFrame, level: int) -> pd.DataFrame:
    """
    Prune base levels of hierarchical tree.

    Parameters
    ----------
    df
        Dataframe representing a hierarchical tree.
        Must have columns `id`, `parent_id` and `level`.
    level
        Level at which to prune tree.

    Returns
    -------
    pd.DataFrame
        Pruned hierarchical tree.
        Column `id` (and `parent_id`) is reset to (0, ..., nrows - 1).

    Examples
    --------
    >>> parent_id = pd.Series([3, 3, 4, 4, None], dtype='Int64')
    >>> df = pd.DataFrame({
    ...     'id': [0, 1, 2, 3, 4],
    ...     'parent_id': parent_id,
    ...     'level': [3, 3, 3, 2, 1]
    ... })
    >>> prune_tree(df, level=2)
       id  parent_id  level
    2   0          2      2
    3   1          2      2
    4   2        NaN      1
    """
    levels = df["level"].max()
    if level > levels:
        return df
    # Drop direct children of all parents up to and including max_level
    pids = df["id"][(df["level"] >= level) & (df["level"] < levels)]
    df = df[~df["parent_id"].isin(pids)].copy()
    # Bump level of remaining children
    df.loc[df["level"] > level, "level"] = level
    # Normalize ids to 0, ..., n
    mask = ~df["parent_id"].isna()
    df.loc[mask, "parent_id"] = np.searchsorted(df["id"], df["parent_id"][mask])
    df["id"] = np.arange(len(df))
    return df


def prepare_merge(merge: dict, df: pd.DataFrame) -> dict:
    """
    Prepare merge for a target dataframe.

    Parameters
    ----------
    merge
        Parameters to :func:`merge_row_pair`.
    df
        Dataframe to prepare merge for.

    Raises
    ------
    ValueError
        Column names duplicated in merge.
    ValueError
        Weights not present in dataframe.
    ValueError
        Weights not included in merge.

    Examples
    --------
    >>> df = pd.DataFrame(columns=['mw', 'lcoe'])
    >>> merge = {'sums': ['mw', 'area'], 'means': ['lcoe'], 'weight': 'mw'}
    >>> prepare_merge(merge, df)
    {'sums': ['mw'], 'means': ['lcoe'], 'weight': 'mw'}
    """
    reduced = {}
    for key in "sums", "means", "uniques":
        if merge.get(key):
            reduced[key] = [x for x in merge[key] if x in df]
    columns = get_merge_columns(reduced)
    if reduced.get("means") and merge.get("weight"):
        weight = merge["weight"]
        if weight not in df:
            raise ValueError(f"Weights {weight} not present in dataframe")
        if weight not in columns:
            raise ValueError(f"Weights {weight} not included in merge")
        reduced["weight"] = weight
    return reduced


def get_merge_columns(merge: dict, df: pd.DataFrame = None) -> list:
    """
    Get columns included in merge.

    Parameters
    ----------
    merge
        Parameters to :func:`merge_row_pair`.
    df
        Dataframe.
        If provided, only matching column names are returned, in order of appearance.

    Raises
    ------
    ValueError
        Column names duplicated in merge.

    Examples
    --------
    >>> merge = {'sums': ['mw'], 'means': ['lcoe'], 'uniques': None, 'weight': 'lcoe'}
    >>> get_merge_columns(merge)
    ['mw', 'lcoe']
    >>> get_merge_columns(merge, pd.DataFrame(columns=['lcoe', 'mw']))
    ['lcoe', 'mw']
    >>> get_merge_columns({'sums': ['mw'], 'means': ['mw']})
    Traceback (most recent call last):
      ...
    ValueError: Column names duplicated in merge
    """
    columns = (
        (merge.get("sums") or [])
        + (merge.get("means") or [])
        + (merge.get("uniques") or [])
    )
    if len(columns) > len(set(columns)):
        visited = set()
        dup = [x for x in columns if x in visited or (visited.add(x) or False)]
        raise ValueError(f"Column names {dup} duplicated in merge")
    if df is not None:
        return [x for x in df if x in columns]
    return columns


# powergenome/run_powergenome_multiple_outputs_cli.py

import argparse
import copy
import logging
import shutil
import sys
from datetime import datetime as dt
from pathlib import Path

import pandas as pd

import powergenome
from powergenome.fuels import fuel_cost_table
from powergenome.generators import (
    GeneratorClusters,
    add_fuel_labels,
    add_genx_model_tags,
)
from powergenome.GenX import (
    add_cap_res_network,
    check_resource_tags,
    create_policy_req,
    create_regional_cap_res,
    fix_min_power_values,
    hydro_energy_to_power,
    min_cap_req,
    max_cap_req,
    reduce_time_domain,
    add_misc_gen_values,
    network_line_loss,
    network_max_reinforcement,
    network_reinforcement_cost,
    round_col_values,
    set_int_cols,
    calculate_partial_CES_values,
)
from powergenome.load_profiles import make_final_load_curves
from powergenome.transmission import (
    agg_transmission_constraints,
    transmission_line_distance,
)
from powergenome.nrelatb import atb_fixed_var_om_existing
from powergenome.external_data import (
    insert_user_tx_costs,
    load_user_tx_costs,
    make_generator_variability,
)
from powergenome.util import (
    build_scenario_settings,
    check_settings,
    init_pudl_connection,
    load_settings,
    remove_fuel_scenario_name,
    remove_fuel_gen_scenario_name,
    update_dictionary,
    write_case_settings_file,
    write_results_file,
    load_ipm_shapefile,
)

if not sys.warnoptions:
    import warnings

    warnings.simplefilter("ignore")


def parse_command_line(argv):
    """
    Parse command line arguments. See the -h option.

    :param argv: arguments on the command line must include caller file name.
    """
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "-sf",
        "--settings_file",
        dest="settings_file",
        type=str,
        default="example_settings.yml",
        help="Specify a YAML settings file.",
    )
    parser.add_argument(
        "-rf",
        "--results_folder",
        dest="results_folder",
        type=str,
        default=dt.now().strftime("%Y-%m-%d %H.%M.%S"),
        help="Specify the results subfolder to write output",
    )
    parser.add_argument(
        "--no-current-gens",
        dest="current_gens",
        action="store_false",
        help="Don't load and cluster current generators.",
    )
    parser.add_argument(
        "--no-gens",
        dest="gens",
        action="store_false",
        help="Use flag to not calculate generator clusters.",
    )
    parser.add_argument(
        "--no-load",
        dest="load",
        action="store_false",
        help="Calculate hourly load. If False, file will not be written.",
    )
    parser.add_argument(
        "--no-transmission",
        dest="transmission",
        action="store_false",
        help="Calculate transmission constraints. If False, file will not be written.",
    )
    parser.add_argument(
        "-f",
        "--no-fuel",
        dest="fuel",
        action="store_false",
        help=(
            "Create fuel table. If False, file will not be written."
            " Can not be created without the generators."
        ),
    )
    parser.add_argument(
        "-s",
        "--sort-gens",
        dest="sort_gens",
        action="store_true",
        help=(
            "Sort generators alphabetically within region. Existing resources will "
            "still be separate from new resources."
        ),
    )
    arguments = parser.parse_args(argv[1:])
    return arguments


def main(**kwargs):
    args = parse_command_line(sys.argv)
    args.__dict__.update(kwargs)
    cwd = Path.cwd()

    out_folder = cwd / args.results_folder
    out_folder.mkdir(exist_ok=True)

    # Create a logger to output any messages we might have...
    logger = logging.getLogger(powergenome.__name__)
    logger.setLevel(logging.INFO)
    handler = logging.StreamHandler()
    formatter = logging.Formatter(
        # More extensive test-like formatter...
        "%(asctime)s [%(levelname)8s] %(name)s:%(lineno)s %(message)s",
        # This is the datetime format string.
        "%Y-%m-%d %H:%M:%S",
    )
    handler.setFormatter(formatter)
    logger.addHandler(handler)

    filehandler = logging.FileHandler(out_folder / "log.txt")
    filehandler.setFormatter(formatter)
    logger.addHandler(filehandler)

    logger.info("Reading settings file")
    settings = load_settings(path=args.settings_file)

    # Copy the settings file to results folder
    if Path(args.settings_file).is_file():
        shutil.copy(args.settings_file, out_folder)
    else:
        shutil.copytree(
            args.settings_file, out_folder / "pg_settings", dirs_exist_ok=True
        )

    logger.info("Initiating PUDL connections")

    pudl_engine, pudl_out, pg_engine = init_pudl_connection(
        freq="AS",
        start_year=min(settings.get("data_years")),
        end_year=max(settings.get("data_years")),
    )

    check_settings(settings, pg_engine)

    # Make sure everything in model_regions is either an aggregate region
    # or an IPM region. Will need to change this once we start using non-IPM
    # regions.
    ipm_regions = pd.read_sql_table("regions_entity_epaipm", pg_engine)[
        "region_id_epaipm"
    ]
    all_valid_regions = ipm_regions.tolist() + list(
        settings.get("region_aggregations", {}) or {}
    )
    good_regions = [region in all_valid_regions for region in settings["model_regions"]]

    if not all(good_regions):
        logger.warning(
            "One or more model regions is not valid. Check to make sure all regions "
            "are either in IPM or region_aggregations in the settings YAML file."
        )

    # Sort zones in the settings to make sure they are correctly sorted everywhere.
    settings["model_regions"] = sorted(settings["model_regions"])
    zones = settings["model_regions"]
    logger.info(f"Sorted zones are {', '.join(zones)}")
    zone_num_map = {
        zone: f"{number + 1}" for zone, number in zip(zones, range(len(zones)))
    }

    input_folder = Path(args.settings_file).parent / settings["input_folder"]
    settings["input_folder"] = input_folder

    scenario_definitions = pd.read_csv(
        input_folder / settings["scenario_definitions_fn"]
    )

    assert set(scenario_definitions["year"]) == set(
        settings["model_year"]
    ), "The years included the secenario definitions file must match the settings parameter `model_year`"
    assert len(settings["model_year"]) == len(
        settings["model_first_planning_year"]
    ), "The number of years in the settings parameter 'model_year' must be the same as 'model_first_planning_year'"

    # Build a dictionary of settings for every planning year and case_id
    scenario_settings = build_scenario_settings(settings, scenario_definitions)

    i = 0
    model_regions_gdf = None
    for year in scenario_settings:
        for case_id, _settings in scenario_settings[year].items():
            case_folder = (
                out_folder / f"{year}" / f"{case_id}_{year}_{_settings['case_name']}"
            )
            _settings["extra_outputs"] = case_folder / "extra_outputs"
            _settings["extra_outputs"].mkdir(parents=True, exist_ok=True)
            if i == 0:
                if args.gens:
                    gc = GeneratorClusters(
                        pudl_engine=pudl_engine,
                        pudl_out=pudl_out,
                        pg_engine=pg_engine,
                        settings=_settings,
                        current_gens=args.current_gens,
                        sort_gens=args.sort_gens,
                    )
                    gen_clusters = gc.create_all_generators()
                    if args.fuel and args.gens:
                        fuels = fuel_cost_table(
                            fuel_costs=gc.fuel_prices,
                            generators=gc.all_resources,
                            settings=_settings,
                        )
                        fuels.index.name = "Time_Index"
                        write_results_file(
                            df=remove_fuel_scenario_name(fuels, _settings),
                            folder=case_folder,
                            file_name="Fuels_data.csv",
                            include_index=True,
                        )

                    gen_clusters["Zone"] = gen_clusters["region"].map(zone_num_map)
                    gen_clusters = add_misc_gen_values(gen_clusters, _settings)
                    gen_clusters = hydro_energy_to_power(
                        gen_clusters,
                        _settings.get("hydro_factor"),
                        _settings.get("regional_hydro_factor", {}),
                    )

                    # Save existing resources that aren't demand response for use in
                    # other cases
                    existing_gens = gc.existing_resources.copy()

                    logger.info(
                        f"\nFinished first round with year {year} scenario {case_id}\n"
                    )
                    gen_variability = make_generator_variability(gen_clusters)
                    gen_variability.index.name = "Time_Index"
                    gen_variability.columns = gen_clusters["Resource"]
                    gens = fix_min_power_values(gen_clusters, gen_variability)
                    for col in _settings["generator_columns"]:
                        if col not in gens.columns:
                            gens[col] = 0
                    cols = [c for c in _settings["generator_columns"] if c in gens]

                    write_results_file(
                        df=remove_fuel_gen_scenario_name(
                            gens[cols].fillna(0), _settings
                        )
                        .pipe(set_int_cols)
                        .pipe(round_col_values)
                        .pipe(check_resource_tags),
                        folder=case_folder,
                        file_name="Generators_data.csv",
                        include_index=False,
                    )

                    i += 1
                if args.transmission:
                    if args.gens is False:
                        model_regions_gdf = load_ipm_shapefile(_settings)

            else:
                logger.info(f"\nStarting year {year} scenario {case_id}\n")
                if args.gens:
                    gc.settings = _settings

                    gen_clusters = gc.create_all_generators()
                    gen_clusters = add_misc_gen_values(gen_clusters, _settings)
                    gen_clusters = hydro_energy_to_power(
                        gen_clusters,
                        _settings.get("hydro_factor"),
                        _settings.get("regional_hydro_factor"),
                    )
                    gen_clusters = set_int_cols(gen_clusters)
                    gen_clusters["Zone"] = gen_clusters["region"].map(zone_num_map)

                    fuels = fuel_cost_table(
                        fuel_costs=gc.fuel_prices,
                        generators=gc.all_resources,
                        settings=_settings,
                    )
                    gen_variability = make_generator_variability(gen_clusters)
                    gen_variability.index.name = "Time_Index"
                    gen_variability.columns = gen_clusters["Resource"]
                    gens = fix_min_power_values(gen_clusters, gen_variability)
                    cols = [c for c in _settings["generator_columns"] if c in gens]
                    write_results_file(
                        df=remove_fuel_gen_scenario_name(
                            gens[cols].fillna(0), _settings
                        )
                        .pipe(set_int_cols)
                        .pipe(round_col_values)
                        .pipe(check_resource_tags),
                        folder=case_folder,
                        file_name="Generators_data.csv",
                        include_index=False,
                    )

            if args.load:
                load = make_final_load_curves(pg_engine=pg_engine, settings=_settings)
                load.columns = "Load_MW_z" + load.columns.map(zone_num_map)

                (
                    reduced_resource_profile,
                    reduced_load_profile,
                    time_series_mapping,
                    representative_point,
                ) = reduce_time_domain(gen_variability, load, _settings)
                reduced_resource_profile.index.name = "Time_Index"
                write_results_file(
                    df=reduced_load_profile,
                    folder=case_folder,
                    file_name="Load_data.csv",
                    include_index=False,
                )

                write_results_file(
                    df=reduced_resource_profile,
                    folder=case_folder,
                    file_name="Generators_variability.csv",
                    include_index=True,
                    float_format="%.3f",
                )
                if time_series_mapping is not None:
                    write_results_file(
                        df=time_series_mapping,
                        folder=case_folder,
                        file_name="Period_map.csv",
                        include_index=False,
                    )
                if representative_point is not None:
                    write_results_file(
                        df=representative_point,
                        folder=case_folder,
                        file_name="Representative_Period.csv",
                        include_index=False,
                    )

            if args.transmission:
                model_regions_gdf = gc.model_regions_gdf
                if _settings.get("user_transmission_costs"):
                    user_tx_costs = load_user_tx_costs(
                        _settings["input_folder"]
                        / _settings["user_transmission_costs"],
                        _settings["model_regions"],
                    )
                    transmission = agg_transmission_constraints(
                        pg_engine=pg_engine, settings=_settings
                    ).pipe(insert_user_tx_costs, user_costs=user_tx_costs)
                else:
                    transmission = (
                        agg_transmission_constraints(
                            pg_engine=pg_engine, settings=_settings
                        )
                        .pipe(
                            transmission_line_distance,
                            ipm_shapefile=model_regions_gdf,
                            settings=_settings,
                            units="mile",
                        )
                        .pipe(network_line_loss, settings=_settings)
                        .pipe(network_reinforcement_cost, settings=_settings)
                    )
                transmission = (
                    transmission.pipe(network_max_reinforcement, settings=_settings)
                    .pipe(set_int_cols)
                    .pipe(round_col_values)
                    .pipe(add_cap_res_network, settings=_settings)
                )
                zones = settings["model_regions"]
                network_zones = [f"z{n+1}" for n in range(len(zones))]
                nz_df = pd.Series(data=network_zones, name="Network_zones")
                network = pd.concat([pd.DataFrame(nz_df), transmission], axis=1)

                if _settings.get("emission_policies_fn"):
                    energy_share_req = create_policy_req(_settings, col_str_match="ESR")
                    co2_cap = create_policy_req(_settings, col_str_match="CO_2")
                else:
                    energy_share_req = None
                    co2_cap = None
                min_cap = min_cap_req(_settings)
                max_cap = max_cap_req(_settings)

                cap_res = create_regional_cap_res(_settings)

                write_results_file(
                    df=network,
                    folder=case_folder,
                    file_name="Network.csv",
                    include_index=False,
                )
                if energy_share_req is not None:
                    write_results_file(
                        df=energy_share_req,
                        folder=case_folder,
                        file_name="Energy_share_requirement.csv",
                        include_index=False,
                    )
                if cap_res is not None:
                    write_results_file(
                        df=cap_res,
                        folder=case_folder,
                        file_name="Capacity_reserve_margin.csv",
                        include_index=True,
                    )
                if co2_cap is not None:
                    co2_cap = co2_cap.set_index("Region_description")
                    co2_cap.index.name = None
                    write_results_file(
                        df=co2_cap,
                        folder=case_folder,
                        file_name="CO2_cap.csv",
                        include_index=True,
                    )
                if min_cap is not None:
                    write_results_file(
                        df=min_cap,
                        folder=case_folder,
                        file_name="Minimum_capacity_requirement.csv",
                        include_index=False,
                    )
                if max_cap is not None:
                    write_results_file(
                        df=max_cap,
                        folder=case_folder,
                        file_name="Maximum_capacity_limit.csv",
                        include_index=False,
                    )

            if args.fuel and args.gens:
                fuels = fuel_cost_table(
                    fuel_costs=gc.fuel_prices,
                    generators=gc.all_resources,
                    settings=_settings,
                )

                fuels.index.name = "Time_Index"
                write_results_file(
                    df=remove_fuel_scenario_name(fuels, _settings)
                    .pipe(set_int_cols)
                    .pipe(round_col_values),
                    folder=case_folder,
                    file_name="Fuels_data.csv",
                    include_index=True,
                )
            if _settings.get("reserves_fn"):
                shutil.copy(
                    _settings["input_folder"] / _settings["reserves_fn"],
                    case_folder / "Inputs",
                )

            if _settings.get("genx_settings_fn"):
                shutil.copy(cwd / _settings["genx_settings_fn"], case_folder / "Inputs")

            if _settings.get("genx_settings_folder"):
                genx_settings_folder = case_folder / "Settings"
                genx_settings_folder.mkdir(exist_ok=True)
                for f in (cwd / _settings.get("genx_settings_folder")).glob("*.yml"):
                    shutil.copy(f, genx_settings_folder)

            write_case_settings_file(
                settings=_settings,
                folder=case_folder,
                file_name="powergenome_case_settings.yml",
            )


if __name__ == "__main__":
    main()


# powergenome/time_reduction.py

"Functions to cluster or otherwise reduce the number of hours in generation and load profiles"

from sklearn.cluster import KMeans
from sklearn.preprocessing import minmax_scale
import numpy as np
import datetime
import pandas as pd


def kmeans_time_clustering(
    resource_profiles,
    load_profiles,
    days_in_group,
    num_clusters,
    include_peak_day=True,
    load_weight=1,
    variable_resources_only=True,
):
    """Reduce the number of hours in load and resource variability timeseries using
    kmeans clustering.

    This script is adapted from work originally created by Dharik Mallapragada. For more
    information see:
    - Mallapragada, D. S., Papageorgiou, D. J., Venkatesh, A., Lara, C. L., & Grossmann,
    I. E. (2018). Impact of model resolution on scenario outcomes for electricity sector
    system expansion. Energy, 163, 1231–1244.
    https://doi.org/10.1016/j.energy.2018.08.015
    - Mallapragada, D. S., Sepulveda, N. A., & Jenkins, J. D. (2020). Long-run system
    value of battery energy storage in future grids with increasing wind and solar
    generation. Applied Energy, 275, 115390.
    https://doi.org/10.1016/j.apenergy.2020.115390


    Parameters
    ----------
    resource_profiles : DataFrame
        Hourly generation profiles for all resources. Each column is a resource with
        a unique name, each row is a consecutive hour.
    load_profiles : DataFrame
        Hourly demand profiles of load. Each column is a region with a unique name. each
        row is a consecutive hour.
    days_in_group : int
        The number of 24 hour periods included in each group/cluster
    num_clusters : int
        The number of clusters to include in the output
    include_peak_day : bool, optional
        If the days with system peak demand should be included in outputs, by default
        True
    load_weight : int, optional
        A weighting factor for load profiles during clustering, by default 1
    variable_resources_only : bool, optional
        If clustering should only consider resources with variable (non-zero standard
        deviation) profiles, by default True

    Returns
    -------
    (dict, list, list)
        This function returns multiple items. The dict has keys ['load_profiles',
        'resource_profiles', 'ClusterWeights', 'AnnualGenScaleFactor', 'RMSE', and
        'AnnualProfile']

        The first list has strings with the order of periods selected e.g. ['p42','p26',
        'p3', 'p13', 'p32', 'p8'].

        The second list has integer weights of each cluster.
    """
    resource_col_names = resource_profiles.columns
    if variable_resources_only:
        input_std = resource_profiles.describe().loc["std", :]
        var_col_names = [col for col in input_std.index if input_std[col] > 0]
        resource_profiles = resource_profiles.loc[:, var_col_names]
    # Initialize dataframes to store final and intermediate data in

    input_data = pd.concat(
        [
            load_profiles.reset_index(drop=True),
            resource_profiles.reset_index(drop=True),
        ],
        axis=1,
    )
    input_data = input_data.reset_index(drop=True)
    original_col_names = input_data.columns.tolist()
    # CAUTION: Load Column lables should be named with the phrase "Load_"
    load_col_names = load_profiles.columns

    # Columns to be reported in output files
    new_col_names = input_data.columns.tolist() + ["GrpWeight"]
    # Dataframe storing final outputs
    final_output_data = pd.DataFrame(columns=new_col_names)

    # Dataframe storing normalized inputs
    norm_tseries = pd.DataFrame(columns=original_col_names)

    hours_per_year = len(input_data)

    # Normalized all load and renewables data 0 and LoadWeight, All Renewables b/w 0
    # and 1
    norm_tseries = pd.DataFrame(
        data=minmax_scale(input_data), columns=input_data.columns
    )
    norm_tseries.loc[:, load_col_names] *= load_weight

    # Identify hour with maximum system wide load
    hr_maxSysLoad = input_data.loc[:, load_col_names].sum(axis=1).idxmax()
    ################################ pre-processing data to create concatenated column
    # of load, pv and wind data

    # Number of such samples in a year - by avoiding division by float we are excluding
    # a few days in each sample set
    # Hence annual generation for each zone will not exactly match up with raw data
    # num_data_points = round(hours_per_year / 24 / days_in_group)
    num_data_points = int(hours_per_year / 24 // days_in_group)

    DataPointAsColumns = [f"p{j}" for j in range(1, num_data_points + 1)]

    # Create a dictionary storing groups of time periods to average over for each hour
    HourlyGroupings = {
        i: [j for j in range(days_in_group * 24 * (i - 1), days_in_group * 24 * i)]
        for i in range(1, num_data_points + 1)
    }

    #  Create a new dataframe storing aggregated load and renewables time series
    ModifiedDataNormalized = pd.DataFrame(columns=DataPointAsColumns)
    # Original data organized in concatenated column
    ModifiedData = pd.DataFrame(columns=DataPointAsColumns)

    # Creating the dataframe with concatenated columns
    for j in range(num_data_points):
        if j == 1:  # Store  variable names for the concatenated column
            ConcatenatedRowNames = norm_tseries.loc[HourlyGroupings[j + 1], :].melt(
                id_vars=None
            )["variable"]

        ModifiedDataNormalized[DataPointAsColumns[j]] = norm_tseries.loc[
            HourlyGroupings[j + 1], :
        ].melt(id_vars=None)["value"]
        ModifiedData[DataPointAsColumns[j]] = input_data.loc[
            HourlyGroupings[j + 1], :
        ].melt(id_vars=None)["value"]

    # Eliminate grouping including the hour with largest system laod (GW) - this group
    # will be manually included in the outputs
    if include_peak_day:
        # IP()
        GroupingwithPeakLoad = ["p" + str(int(hr_maxSysLoad / 24 / days_in_group + 1))]
        ClusteringInputDF = ModifiedDataNormalized.drop(GroupingwithPeakLoad, axis=1)
    else:
        ClusteringInputDF = ModifiedDataNormalized

    ################################## k-means clustering process
    # create Kmeans clustering model and specify the number of clusters gathered
    # number of replications =100, squared euclidean distance

    if include_peak_day:  # If peak day in cluster, generate one less cluster
        num_clusters = num_clusters - 1

    # K-means clutering with 100 trials with randomly selected starting values
    model = KMeans(
        n_clusters=num_clusters, n_init=100, init="k-means++", random_state=42
    )
    model.fit(ClusteringInputDF.values.transpose())

    # Store clustered data
    # Create an empty list storing weight of each cluster
    EachClusterWeight = [None] * num_clusters

    # Create an empty list storing name of each data point
    EachClusterRepPoint = [None] * num_clusters

    # creating a dataframe for storing the mapping between representative time period and the entire year
    time_series_mapping = pd.DataFrame(columns=["Period_Index", "Rep_Period_Index"])

    for k in range(num_clusters):
        # Number of points in kth cluster (i.e. label=0)
        EachClusterWeight[k] = len(model.labels_[model.labels_ == k])

        # Compute Euclidean distance of each point from centroid of cluster k
        dist = {
            ClusteringInputDF.loc[:, model.labels_ == k].columns[j]: np.linalg.norm(
                ClusteringInputDF.loc[:, model.labels_ == k].values.transpose()[j]
                - model.cluster_centers_[k]
            )
            for j in range(EachClusterWeight[k])
        }

        # Select column name closest with the smallest euclidean distance to the mean
        EachClusterRepPoint[k] = min(dist, key=lambda k: dist[k])

        # Creating a list that matches each week to a representative week
        for j in range(EachClusterWeight[k]):
            time_series_mapping = time_series_mapping.append(
                pd.DataFrame(
                    {
                        "Period_Index": int(
                            ClusteringInputDF.loc[:, model.labels_ == k].columns[j][1:]
                        ),
                        "Rep_Period_Index": k + 1,
                    },
                    index=[0],
                ),
                ignore_index=True,
            )
    if include_peak_day:
        # appending the week representing peak load
        time_series_mapping = time_series_mapping.append(
            pd.DataFrame(
                {
                    "Period_Index": int(GroupingwithPeakLoad[0][1:]),
                    "Rep_Period_Index": k + 2,
                },
                index=[0],
            ),
            ignore_index=True,
        )

    # same CSV file that will be used in GenX
    time_series_mapping = time_series_mapping.sort_values(by=["Period_Index"])
    time_series_mapping = time_series_mapping.reset_index(drop=True)

    # extract month corresponding to each time slot
    time_series_mapping["Month"] = 0
    for Period_Index in time_series_mapping["Period_Index"]:
        dayOfYear = days_in_group * Period_Index
        d = datetime.datetime.strptime("{} {}".format(dayOfYear, 2011), "%j %Y")
        time_series_mapping["Month"][Period_Index - 1] = d.month

    # Storing selected groupings in a new data frame with appropriate dimensions
    # (E.g. load in GW)
    ClusterOutputDataTemp = ModifiedData[EachClusterRepPoint]

    # Selecting rows corresponding to Load in excluded subperiods and exclude them from
    # scale factor calculation
    NRowsLoad = len(load_col_names)
    # Excluding grouping with peak hr from scale factor calculation
    if include_peak_day:
        Actualdata = ModifiedData.loc[0 : 24 * days_in_group * NRowsLoad - 1, :].drop(
            GroupingwithPeakLoad, axis=1
        )
    else:
        Actualdata = ModifiedData.loc[0 : 24 * days_in_group * NRowsLoad - 1, :]

    # Scale factor to adjust total generation in original data set to be equal to scaled
    # up total generation in sampled data set
    SampleweeksAnnualTWh = sum(
        [
            ClusterOutputDataTemp.loc[
                0 : 24 * days_in_group * NRowsLoad - 1, EachClusterRepPoint[j]
            ].sum()
            * EachClusterWeight[j]
            for j in range(num_clusters)
        ]
    )
    ScaleFactor = (
        Actualdata.loc[0 : 24 * days_in_group * NRowsLoad - 1, :].sum().sum()
        / SampleweeksAnnualTWh
    )

    # Updated load values in GW
    ClusterOutputDataTemp.loc[0 : 24 * days_in_group * NRowsLoad - 1, :] = (
        ScaleFactor
        * ClusterOutputDataTemp.loc[0 : 24 * days_in_group * NRowsLoad - 1, :]
    )

    # Add the grouping with the peak hour back into the cluster if that is excluded in
    # the clustering
    if include_peak_day:
        EachClusterRepPoint = EachClusterRepPoint + GroupingwithPeakLoad
        EachClusterWeight = EachClusterWeight + [1]
        ClusterOutputData = pd.concat(
            [ClusterOutputDataTemp, ModifiedData[GroupingwithPeakLoad]],
            axis=1,
            sort=False,
        )
    else:
        ClusterOutputData = ClusterOutputDataTemp

    # Store weights for each selected hour  Number of days *24, for each week
    ClusteredWeights = pd.DataFrame(
        EachClusterWeight * np.ones([days_in_group * 24, len(EachClusterWeight)]),
        columns=EachClusterRepPoint,
    )

    # Storing weights in final output data column
    final_output_data["GrpWeight"] = ClusteredWeights.melt(id_vars=None)["value"]

    # Regenerating data organized by time series (columns) and representative time
    # periods (hours)
    for i in range(len(new_col_names) - 1):
        final_output_data[new_col_names[i]] = ClusterOutputData.loc[
            ConcatenatedRowNames == new_col_names[i], :
        ].melt(id_vars=None)["value"]

    # Calculating error metrics and Annual profile
    FullLengthOutputs = final_output_data
    for j in range(len(EachClusterWeight)):
        # Selecting rows of the FinalOutputData dataframe to append
        df_try = final_output_data.truncate(
            before=days_in_group * 24 * j, after=days_in_group * 24 * (j + 1) - 1
        )
        #        print(EachClusterWeight[j])
        if (
            EachClusterWeight[j] > 1
        ):  # Need to duplicate entries only weight is greater than 1
            FullLengthOutputs = FullLengthOutputs.append(
                [df_try] * (EachClusterWeight[j] - 1), ignore_index=True
            )

    #  Root mean square error between the duration curves of each time series
    # Only conisder the points consider in the k-means clustering - ignoring any days
    # dropped off from original data set  due to rounding
    RMSE = {
        col: np.linalg.norm(
            np.sort(input_data.truncate(after=len(FullLengthOutputs) - 1)[col].values)
            - np.sort(FullLengthOutputs[col].values)
        )
        for col in original_col_names
    }

    load_df = final_output_data.loc[:, load_col_names]
    # if variable_only:
    resource_df = pd.DataFrame(
        columns=resource_col_names, index=final_output_data.index
    )
    for col in resource_col_names:
        try:
            resource_df[col] = final_output_data.loc[:, col].values
        except KeyError:
            pass
    resource_df = resource_df.fillna(value=1)

    # load_df["Sub_Weights"] = np.nan
    # load_df.loc[: len(EachClusterWeight) - 1, "Sub_Weights"] = (
    #     np.array(EachClusterWeight) * NumGrpDays * 24
    # )
    # load_df.to_csv("load_time_reduced.csv", index=False)
    # renewable_df = FinalOutputData.loc[
    #     :, [col for col in FinalOutputData.columns if "Load_" not in col]
    # ]
    # renewable_df = renewable_df.drop(columns=["GrpWeight"])
    # renewable_df.insert(loc=0, column="Resource", value=renewable_df.index + 1)
    # renewable_df.to_csv("renewables_time_reduced.csv", index=False)
    rep_period_map = {i + 1: int(p[1:]) for i, p in enumerate(EachClusterRepPoint)}
    time_series_mapping["Rep_Period"] = time_series_mapping["Rep_Period_Index"].map(
        rep_period_map
    )
    EachClusterRepPoint = pd.DataFrame(EachClusterRepPoint, columns=["slot"])
    return (
        {
            "load_profiles": load_df,  # Scaled Output Load and Renewables profiles for the sampled representative groupings
            "resource_profiles": resource_df,
            "ClusterWeights": EachClusterWeight,  # Weight of each for the representative groupings
            "AnnualGenScaleFactor": ScaleFactor,  # Scale factor used to adjust load output to match annual generation of original data
            "RMSE": RMSE,  # Root mean square error between full year data and modeled full year data (duration curves)
            "AnnualProfile": FullLengthOutputs,
            "time_series_mapping": time_series_mapping,
        },
        EachClusterRepPoint,
        EachClusterWeight,
    )  # Modeled duration curves GW


# powergenome/transmission.py

"""
Transmission constraints between regions and line distance
"""

import itertools
import logging
from math import asin, cos, radians, sin, sqrt
from pathlib import Path

import pandas as pd
import sqlalchemy as sa

from powergenome.util import map_agg_region_names, reverse_dict_of_lists, find_centroid

logger = logging.getLogger(__name__)


def agg_transmission_constraints(
    pg_engine: sa.engine.base.Engine,
    settings: dict,
    pg_table: str = "transmission_single_epaipm",
    settings_agg_key: str = "region_aggregations",
) -> pd.DataFrame:
    """Aggregate transmission constraints/capacity between model regions

    Model regions can consist of one or more individual regions. When two or more regions
    are in a model region, the transmission capacity between individual regions is
    combined.

    Values in a user transmission table will override the database transmission table.

    Parameters
    ----------
    pg_engine : sa.engine.base.Engine
        Engine to conect with a database
    settings : dict
        Dictionary of settings parameters. Must include "model_regions". Optional parameters
        include "tx_value_col" (name of the data column in `pg_table`, default value is
        "firm_ttc_mw"), "user_transmission_constraints_fn" (name of user data file, must
        be combined with the parameter "input_folder"), and the value of `settings_agg_key`
        if any regions are aggregated.
    pg_table : str, optional
        Name of the database table with transmission capacity, by default "transmission_single_epaipm"
    settings_agg_key : str, optional
        Name of the settings parameter where regions are aggregated, by default "region_aggregations"

    Returns
    -------
    pd.DataFrame
        Network lines connecting regions and the min/max flow of each line (if different
        based on direction)

    Raises
    ------
    KeyError
        The specified data column is not in the database transmission table
    KeyError
        The database transmission table has duplicate lines in the same direction
    KeyError
        The specified data column is not in the user supplied transmission table
    KeyError
        The user transmission table has duplicate lines in the same direction
    """
    tx_value_col = settings.get("tx_value_col")
    if not tx_value_col:
        logger.warning(
            "No transmission value column (e.g. firm vs non-firm) was specified in the "
            "settings. The column 'firm_ttc_mw' will be used as a default. This is a change "
            "from previous versions of PG, where 'nonfirm_ttc_mw' was used. Firm transmission "
            "capacity is lower or equal to non-firm capacity."
        )
        tx_value_col = "firm_ttc_mw"
    zones = settings["model_regions"]
    zone_num_map = {
        zone: f"z{number + 1}" for zone, number in zip(zones, range(len(zones)))
    }

    combos = list(itertools.combinations(zones, 2))
    reverse_combos = [(combo[-1], combo[0]) for combo in combos]

    logger.info("Loading transmission constraints from PUDL")
    transmission_constraints_table = pd.read_sql_table(pg_table, con=pg_engine)

    if tx_value_col not in transmission_constraints_table.columns:
        raise KeyError(
            f"There is no column {tx_value_col} in the transmission capacity table '{pg_table}'"
        )
    if transmission_constraints_table.duplicated(
        subset=["region_from", "region_to"]
    ).any():
        dup_lines = transmission_constraints_table.loc[
            transmission_constraints_table.duplicated(
                subset=["region_from", "region_to"]
            ),
            ["region_from", "region_to"],
        ]

        raise KeyError(
            "The transmission table has duplicate lines. This table should only have unique lines.\n",
            dup_lines,
        )
    if settings.get("user_transmission_constraints_fn"):
        user_tx_constraints = pd.read_csv(
            Path(settings["input_folder"])
            / settings["user_transmission_constraints_fn"]
        )
        if tx_value_col not in user_tx_constraints.columns:
            raise KeyError(
                f"There is no column {tx_value_col} in the user supplied transmission capacity table"
            )
        if user_tx_constraints.duplicated(subset=["region_from", "region_to"]).any():
            dup_lines = user_tx_constraints.loc[
                user_tx_constraints.duplicated(subset=["region_from", "region_to"]),
                ["region_from", "region_to"],
            ]

            raise KeyError(
                "The user transmission table has duplicate lines. This table should only have unique lines.\n",
                dup_lines,
            )

        # user constraints are needed bidirectionaly
        transmission_constraints_table = pd.concat(
            [
                transmission_constraints_table,
                user_tx_constraints,
                user_tx_constraints.rename(
                    columns={"region_from": "region_to", "region_to": "region_from"}
                ),
            ]
        )

        if transmission_constraints_table.duplicated(
            subset=["region_from", "region_to"]
        ).any():
            logger.warning(
                "The user transmission capacity table duplicates values from the "
                "database. Database values will be discarded in favor of user values."
            )

        transmission_constraints_table = transmission_constraints_table.drop_duplicates(
            keep="last"
        )

    # Settings has a dictionary of lists for regional aggregations. Need
    # to reverse this to use in a map method.
    region_agg_map = reverse_dict_of_lists(settings.get(settings_agg_key))

    # IPM regions to keep. Regions not in this list will be dropped from the
    # dataframe
    keep_regions = [
        x
        for x in settings["model_regions"] + list(region_agg_map)
        if x not in region_agg_map.values()
    ]

    # Create new column "model_region_from"  and "model_region_to" with labels that
    # we're using for aggregated regions
    transmission_constraints_table = transmission_constraints_table.loc[
        (transmission_constraints_table.region_from.isin(keep_regions))
        & (transmission_constraints_table.region_to.isin(keep_regions)),
        :,
    ].drop(columns="id", errors="ignore")

    logger.info("Map and aggregate region names for transmission constraints")
    for col in ["region_from", "region_to"]:
        model_col = "model_" + col

        transmission_constraints_table = map_agg_region_names(
            df=transmission_constraints_table,
            region_agg_map=region_agg_map,
            original_col_name=col,
            new_col_name=model_col,
        )

    keep_cols = ["model_region_from", "model_region_to", tx_value_col]
    drop_cols = [
        c for c in transmission_constraints_table.columns if c not in keep_cols
    ]
    transmission_constraints_table.drop(columns=drop_cols, inplace=True)
    transmission_constraints_table = transmission_constraints_table.groupby(
        ["model_region_from", "model_region_to"]
    ).sum()

    # Build the final output dataframe
    logger.info(
        "Build a new transmission constraints dataframe with a single line between "
        "regions"
    )
    tc_joined = pd.DataFrame(
        columns=["Network_Lines"] + zones + ["Line_Max_Flow_MW", "Line_Min_Flow_MW"],
        index=transmission_constraints_table.reindex(combos).dropna().index,
        data=0,
    )

    if tc_joined.empty:
        logger.info(f"No transmission lines exist between model regions {combos}")
        tc_joined["transmission_path_name"] = None
        tc_joined.rename(columns=zone_num_map, inplace=True)
        return tc_joined.reset_index(drop=True)

    tc_joined["Network_Lines"] = range(1, len(tc_joined) + 1)
    tc_joined["Line_Max_Flow_MW"] = transmission_constraints_table.reindex(
        combos
    ).dropna()

    reverse_tc = transmission_constraints_table.reindex(reverse_combos).dropna() * -1
    reverse_tc.index = tc_joined.index
    tc_joined["Line_Min_Flow_MW"] = reverse_tc

    for idx, _ in tc_joined.iterrows():
        tc_joined.loc[idx, idx[0]] = 1
        tc_joined.loc[idx, idx[-1]] = -1

    tc_joined.rename(columns=zone_num_map, inplace=True)
    tc_joined = tc_joined.reset_index()
    tc_joined["transmission_path_name"] = (
        tc_joined["model_region_from"] + "_to_" + tc_joined["model_region_to"]
    )
    # tc_joined = tc_joined.set_index("Transmission Path Name")
    tc_joined.drop(columns=["model_region_from", "model_region_to"], inplace=True)

    return tc_joined


def haversine(lon1, lat1, lon2, lat2, units="mile"):
    """
    Calculate the great circle distance between two points
    on the earth (specified in decimal degrees)

    https://gis.stackexchange.com/questions/166820/geopandas-return-lat-and-long-of-a-centroid-point
    """
    # convert decimal degrees to radians
    lon1, lat1, lon2, lat2 = map(radians, [lon1, lat1, lon2, lat2])

    # haversine formula
    dlon = lon2 - lon1
    dlat = lat2 - lat1
    a = sin(dlat / 2) ** 2 + cos(lat1) * cos(lat2) * sin(dlon / 2) ** 2
    c = 2 * asin(sqrt(a))

    if units == "mile":
        r = 3956  # Radius of earth in miles. Use 6371 for kilometers, 3956 for miles
    elif units == "km":
        r = 6371
    else:
        raise ValueError(f"Units are {units}, but should be 'mile' or 'km'")

    return c * r


def getXY(pt):
    "Return the X and Y parts of a coordinate point"
    return (pt.x, pt.y)


def single_line_distance(line_name, region_centroids, units):
    """Calculate the transmission line distance between centroids of two regions.

    Parameters
    ----------
    line_name : str
        Two region names in the format <start>_to_<end>
    region_centroids : geoseries
        Centroid points of each region with region names as the index
    units : str
        Name of the distance units to use. Options are 'mile' or 'km'.

    Returns
    -------
    float
        The distance
    """

    start, end = line_name.split("_to_")
    start_lon, start_lat = getXY(region_centroids[start])
    end_lon, end_lat = getXY(region_centroids[end])
    distance = haversine(start_lon, start_lat, end_lon, end_lat, units=units)

    return distance


def transmission_line_distance(
    trans_constraints_df, ipm_shapefile, settings, units="mile"
):
    logger.info("Calculating transmission line distance")
    ipm_shapefile["geometry"] = ipm_shapefile.buffer(0.01)
    model_polygons = ipm_shapefile.dissolve(by="model_region")
    model_polygons = model_polygons.to_crs(epsg=4326)
    region_centroids = find_centroid(model_polygons)

    distances = [
        single_line_distance(line_name, region_centroids, units=units)
        for line_name in trans_constraints_df["transmission_path_name"]
    ]
    trans_constraints_df[f"distance_{units}"] = distances
    trans_constraints_df[f"distance_{units}"] = trans_constraints_df[
        f"distance_{units}"
    ]

    return trans_constraints_df


# powergenome/util.py

import collections
from copy import deepcopy
import functools
import itertools
import logging
import re
import subprocess
from typing import Dict, List, Tuple, Union
from collections.abc import Iterable

import pandas as pd
import geopandas as gpd
import pudl
import requests
import sqlalchemy as sa

from flatten_dict import flatten
import yaml
from ruamel.yaml import YAML
from pathlib import Path
from frozendict import frozendict

from powergenome.params import IPM_GEOJSON_PATH, SETTINGS

logger = logging.getLogger(__name__)


def load_settings(path: Union[str, Path]) -> dict:
    """Load a YAML file or a dictionary of YAML files with settings parameters

    Parameters
    ----------
    path : Union[str, Path]
        Name of the settings file or folder

    Returns
    -------
    dict
        All parameters listed in the YAML file(s)
    """

    path = Path(path)
    if path.is_file():
        with open(path, "r") as f:
            #     settings = yaml.safe_load(f)
            yaml = YAML(typ="safe")
            settings = yaml.load(f)
    elif path.is_dir():
        settings = {}
        for sf in path.glob("*.yml"):
            yaml = YAML(typ="safe")
            s = yaml.load(sf)
            if s:
                settings.update(s)

    if settings.get("input_folder"):
        settings["input_folder"] = path.parent / settings["input_folder"]
    return fix_param_names(settings)


def fix_param_names(settings: dict) -> dict:
    fix_params = {
        "historical_load_region_maps": "historical_load_region_map",
        "demand_response_resources": "flexible_demand_resources",
    }
    for k, v in fix_params.items():
        if k in settings:
            settings[v] = settings[k]
            s = f"""
            The settings parameter named {k} has been changed to {v}. Please correct it in
            your settings file.

            """
            logger.warning(s)
    return settings


def findkeys(node: Union[dict, list], kv: str):
    """
    Return all values in a dictionary from a matching key
    https://stackoverflow.com/a/19871956
    """
    if isinstance(node, list):
        for i in node:
            for x in findkeys(i, kv):
                yield x
    elif isinstance(node, dict):
        if kv in node:
            yield node[kv]
        for j in node.values():
            for x in findkeys(j, kv):
                yield x


def check_atb_scenario(settings: dict, pg_engine: sa.engine.base.Engine):
    """Check the

    Parameters
    ----------
    settings : dict
        Parameters and values from the YAML settings file.
    pg_engine : sa.engine.base.Engine
        Connection to the PG sqlite database.

    Raises
    ------
    KeyError
        Raises an error if an ATB technology scenario in the settings file doesn't match
        the list of available values for that year of ATB data.
    """
    atb_year = settings.get("atb_data_year")

    s = f"""
    SELECT DISTINCT cost_case
    FROM technology_costs_nrelatb
    WHERE
        atb_year == {atb_year}
    """

    atb_cases = [c[0] for c in pg_engine.execute(s).fetchall()]

    techs = []
    for l in findkeys(settings, "atb_new_gen"):
        techs.extend(l)

    cases = [tech[2] for tech in techs]

    for l in findkeys(settings, "atb_cost_case"):
        cases.append(l)

    bad_case_names = []
    for case in cases:
        if case not in atb_cases:
            bad_case_names.append(case)
    if bad_case_names:
        bad_names = list(set(bad_case_names))
        raise KeyError(
            f"There is an error with the ATB tech scenario key in your settings file."
            f" You are using ATB data from {atb_year}, which has cost cases of:\n\n "
            f"{atb_cases}\n\n"
            "Under either 'atb_new_gen' or 'modified_atb_new_gen' you have cost cases "
            f"of:\n\n{bad_names}\n\n "
            "Try searching your settings file for these "
            "values and replacing them with valid cost cases for your ATB year."
        )


def check_settings(settings: dict, pg_engine: sa.engine) -> None:
    """Check for user errors in the settings file.

    The YAML settings file is loaded as a dictionary object. It has many different parts
    that need to have consistent values. This function checks a few (but not all!) of
    the parameters for common errors or misspelled words.

    Parameters
    ----------
    settings : dict
        Parameters and values from the YAML settings file.
    pg_engine : sa.engine
        Connection to the PG sqlite database.
    """
    if settings.get("atb_data_year"):
        check_atb_scenario(settings, pg_engine)
    ipm_region_list = pd.read_sql_table("regions_entity_epaipm", pg_engine)[
        "region_id_epaipm"
    ].to_list()

    cost_mult_regions = list(
        itertools.chain.from_iterable(
            settings.get("cost_multiplier_region_map", {}).values()
        )
    )

    aeo_fuel_regions = list(
        itertools.chain.from_iterable(settings.get("aeo_fuel_region_map", {}).values())
    )

    atb_techs = settings.get("atb_new_gen", []) or []
    atb_mod_techs = settings.get("modified_atb_new_gen", {}) or {}
    add_new_techs = settings.get("additional_new_gen", []) or []
    cost_mult_techs = []
    for k, v in settings.get("cost_multiplier_technology_map", {}).items():
        for t in v:
            cost_mult_techs.append(t)

    # Make sure atb techs are spelled correctly and are in the cost_multiplier_technology_map
    for tech in atb_techs:
        tech, tech_detail, cost_case, _ = tech

        s = f"""
        SELECT technology, tech_detail
        from technology_costs_nrelatb
        where
            technology == "{tech}"
            AND tech_detail == "{tech_detail}"
        """
        if len(pg_engine.execute(s).fetchall()) == 0:
            s = f"""
    *****************************
    The technology {tech} - {tech_detail} listed in your settings file under 'atb_new_gen'
    does not match any NREL ATB technologies. Check your settings file to ensure it is
    spelled correctly"
    *****************************
    """
            logger.warning(s)

        if f"{tech}_{tech_detail}" not in cost_mult_techs:
            s = f"""
    *****************************
    The ATB technology "{tech}_{tech_detail}" listed in your settings file under 'atb_new_gen'
    is not fully specified in the 'cost_multiplier_technology_map' settings parameter.
    Part of the <tech>_<tech_detail> string might be included, but it is best practice to
    include the full name in this format. Check your settings file.
        """
            logger.warning((s))

    for mod_tech in atb_mod_techs.values():
        mt_name = f"{mod_tech['new_technology']}_{mod_tech['new_tech_detail']}"
        if mt_name not in cost_mult_techs:
            s = f"""
    *****************************
    The modified ATB technology "{mt_name}" listed in your settings file under
    'modified_atb_new_gen' is not fully specified in the 'cost_multiplier_technology_map'
    settings parameter. Part of the <new_technology>_<new_tech_detail> string might be
    included, but it is best practice to include the full name in this format. Check
    your settings file.
        """
            logger.warning((s))

    for add_tech in add_new_techs:
        if add_tech not in cost_mult_techs:
            s = f"""
    *****************************
    The additional user-specified technology "{add_tech}" listed in your settings file under
    'additional_new_gen' is not fully specified in the 'cost_multiplier_technology_map'
    settings parameter. Part of the name string might be included, but it is best practice
    to include the full name in this format. Check your settings file.
        """
            logger.warning((s))

    for agg_region, ipm_regions in (settings.get("region_aggregations") or {}).items():
        for ipm_region in ipm_regions:
            if ipm_region not in ipm_region_list:
                s = f"""
    *****************************
    There is no IPM region {ipm_region}, which is listed in {agg_region}"
    *****************************
    """
                logger.warning(s)

    for model_region in settings["model_regions"]:
        if model_region not in cost_mult_regions:
            s = f"""
    *****************************
    The model region {model_region} is not included in the settings parameter `cost_multiplier_region_map`"
    *****************************
            """
            logger.warning(s)

        if model_region not in aeo_fuel_regions:
            s = f"""
    *****************************
    The model region {model_region} is not included in the settings parameter `aeo_fuel_region_map`"
    *****************************
            """
            logger.warning(s)

    gen_col_count = collections.Counter(settings.get("generator_columns", []))
    duplicate_cols = [c for c, num in gen_col_count.items() if num > 1]
    if duplicate_cols:
        raise KeyError(
            f"The settings parameter 'generator_columns' has duplicates of {duplicate_cols}."
            " Remove the duplicates and try again."
        )

    if settings.get("eia_aeo_year") or settings.get("fuel_eia_aeo_year"):
        fuel_aeo_year = settings.get("fuel_eia_aeo_year") or settings.get(
            "eia_aeo_year"
        )
        for k, v in settings.get("eia_series_scenario_names", {}).items():
            if "REF" in v and str(fuel_aeo_year) not in v:
                logger.warning(
                    "The settings EIA fuel scenario (eia_series_scenario_names) key "
                    f"{k} has a value of {v}, which does not match the aeo data year "
                    f"{fuel_aeo_year}. It has been changed to REF{fuel_aeo_year}."
                )
                settings["eia_series_scenario_names"][k] = f"REF{fuel_aeo_year}"

    if settings.get("eia_aeo_year") or settings.get("load_eia_aeo_year"):
        load_aeo_year = settings.get("load_eia_aeo_year") or settings.get(
            "eia_aeo_year"
        )
        growth_scenario = settings.get("growth_scenario", "")
        if "REF" in growth_scenario and str(load_aeo_year) not in growth_scenario:
            logger.warning(
                "The settings EIA demand growth scenario (growth_scenario) key "
                f"value is {growth_scenario}, which does not match the aeo data year "
                f"{load_aeo_year}. It has been changed to REF{load_aeo_year}."
            )
            settings["growth_scenario"] = f"REF{load_aeo_year}"


def init_pudl_connection(
    freq: str = "AS",
    start_year: int = None,
    end_year: int = None,
    pudl_db: str = None,
    pg_db: str = None,
) -> Tuple[sa.engine.base.Engine, pudl.output.pudltabl.PudlTabl]:
    """Initiate a connection object to the sqlite PUDL database and create a pudl
    object that can quickly access parts of the database.

    Parameters
    ----------
    freq : str, optional
        The time frequency that data should be averaged over in the `pudl_out` object,
        by default "YS" (annual data).

    Returns
    -------
    sa.Engine, pudl.pudltabl
        A sqlalchemy engine for connecting to the PUDL database, and a pudl PudlTabl
        object for quickly accessing parts of the database. `pudl_out` is used
        to access unit heat rates.
    """
    if not pudl_db:
        pudl_db = SETTINGS["PUDL_DB"]
    if not pg_db:
        if SETTINGS.get("PG_DB"):
            pg_db = SETTINGS["PG_DB"]
        else:
            logger.warning(
                "No path to a `PG_DB` database was provided or found in the .env file. Using "
                "the `PUDL_DB` path instead."
            )
            pg_db = SETTINGS["PUDL_DB"]
    pudl_engine = sa.create_engine(pudl_db)
    if start_year is not None:
        start_year = pd.to_datetime(start_year, format="%Y")
    if end_year is not None:
        end_year = pd.to_datetime(end_year, format="%Y")
    """
    pudl_out = pudl.output.pudltabl.PudlTabl(
        freq=freq, pudl_engine=pudl_engine, start_date=start_year, end_date=end_year
        #freq=freq, pudl_engine=pudl_engine, start_date=start_year, end_date=end_year, ds=""
    )
    """
    pudl_out = pudl.output.pudltabl.PudlTabl(
        freq=freq,
        pudl_engine=pudl_engine,
        start_date=start_year,
        end_date=end_year,
        ds=pudl.workspace.datastore.Datastore(),
    )
    pg_engine = sa.create_engine(pg_db)
    # if SETTINGS.get("PG_DB"):
    #     pg_engine = sa.create_engine(SETTINGS["PG_DB"])
    # else:
    #     logger.warning(
    #         "No path to a `PG_DB` database was found in the .env file. Using the "
    #         "`PUDL_DB` path instead."
    #     )
    #     pg_engine = sa.create_engine(SETTINGS["PUDL_DB"])

    return pudl_engine, pudl_out, pg_engine


def reverse_dict_of_lists(d: Dict[str, list]) -> Dict[str, List[str]]:
    """Reverse the mapping in a dictionary of lists so each list item maps to the key

    Parameters
    ----------
    d : Dict[str, List[str]]
        A dictionary with string keys and lists of strings.

    Returns
    -------
    Dict[str, str]
        A reverse mapped dictionary where the item of each list becomes a key and the
        original keys are mapped as values.
    """
    if isinstance(d, collections.abc.Mapping):
        rev = {v: k for k in d for v in d[k]}
    else:
        rev = dict()
    return rev


def map_agg_region_names(
    df: pd.DataFrame,
    region_agg_map: Dict[str, List[str]],
    original_col_name: str,
    new_col_name: str,
) -> pd.DataFrame:
    """Add a column that maps original region names to aggregated model region names.

    A dataframe with un-aggregated region names (e.g. EPA IPM regions) will have a new
    column added. Aggregated model region names will be used in the new column. If a
    model region is not part of an aggregation it will be left as-is in the new column.

    Parameters
    ----------
    df : pd.DataFrame
        Original dataframe with column 'original_col_name'
    region_agg_map : Dict[str, List[str]]
        Mapping of model region names (keys) to a list of aggregated base regions
    original_col_name : str
        Name of the original column with region names.
    new_col_name : str
        Name for the column with mapped model region values.

    Returns
    -------
    pd.DataFrame
        A modified version of the original dataframe with the new column "new_col_name"
        that has values of model regions.
    """

    df[new_col_name] = df.loc[:, original_col_name]

    df.loc[df[original_col_name].isin(region_agg_map.keys()), new_col_name] = df.loc[
        df[original_col_name].isin(region_agg_map.keys()), original_col_name
    ].map(region_agg_map)

    return df


def snake_case_col(col: pd.Series) -> pd.Series:
    "Remove special characters and convert to snake case"
    clean = (
        col.str.lower()
        .str.replace(r"[^0-9a-zA-Z\-]+", " ", regex=True)
        .str.replace("-", "")
        .str.strip()
        .str.replace(" ", "_")
    )
    return clean


def snake_case_str(s: str) -> str:
    "Remove special characters and convert to snake case"
    clean = (
        re.sub(r"[^0-9a-zA-Z\-]+", " ", s)
        .lower()
        .replace("-", "")
        .strip()
        .replace(" ", "_")
    )
    return clean


def get_git_hash():
    try:
        git_head_hash = (
            subprocess.check_output(["git", "rev-parse", "HEAD"])
            .strip()
            .decode("ascii")
        )
    except FileNotFoundError:
        git_head_hash = "Git hash unknown"

    return git_head_hash


def download_save(url: str, save_path: Union[str, Path]):
    """
    Download a file that isn't zipped and save it to a given path

    Parameters
    ----------
    url : str
        Valid url to download the zip file
    save_path : str or path object
        Destination to save the file

    """

    r = requests.get(url)
    save_path.parent.mkdir(parents=True, exist_ok=True)
    save_path.write_bytes(r.content)


def update_dictionary(d: dict, u: dict) -> dict:
    """
    Update keys in an existing dictionary (d) with values from u

    https://stackoverflow.com/a/32357112
    """
    for k, v in u.items():
        if isinstance(d, collections.abc.Mapping):
            if isinstance(v, collections.abc.Mapping):
                r = update_dictionary(d.get(k, {}), v)
                d[k] = r
            else:
                d[k] = u[k]
        else:
            d = {k: u[k]}
    return d


def remove_fuel_scenario_name(df, settings):
    _df = df.copy()
    scenarios = settings["eia_series_scenario_names"].keys()
    for s in scenarios:
        _df.columns = _df.columns.str.replace(f"_{s}", "")

    return _df


def remove_fuel_gen_scenario_name(df, settings):
    _df = df.copy()
    scenarios = settings["eia_series_scenario_names"].keys()
    for s in scenarios:
        _df["Fuel"] = _df["Fuel"].str.replace(f"_{s}", "")

    return _df


def write_results_file(
    df: pd.DataFrame,
    folder: Path,
    file_name: str,
    include_index: bool = False,
    float_format: str = None,
):
    """Write a finalized dataframe to one of the results csv files.

    Parameters
    ----------
    df : DataFrame
        Data for a single results file
    folder : Path-like
        A Path object representing the folder for a single case/scenario
    file_name : str
        Name of the file.
    include_index : bool, optional
        If pandas should include the index when writing to csv, by default False
    float_format: str
        Parameter passed to pandas .to_csv
    """
    sub_folder = folder / "Inputs"
    sub_folder.mkdir(exist_ok=True, parents=True)

    path_out = sub_folder / file_name
    df.to_csv(path_out, index=include_index, float_format=float_format)


def write_case_settings_file(settings, folder, file_name):
    """Write a finalized dictionary to YAML file.

    Parameters
    ----------
    settings : dict
        A dictionary with settings
    folder : Path-like
        A Path object representing the folder for a single case/scenario
    file_name : str
        Name of the file.
    """
    folder.mkdir(exist_ok=True, parents=True)
    path_out = folder / file_name

    # yaml = YAML(typ="unsafe")
    _settings = deepcopy(settings)
    # for key, value in _settings.items():
    #     if isinstance(value, Path):
    #         _settings[key] = str(value)
    # yaml.register_class(Path)
    # stream = file(path_out, 'w')
    with open(path_out, "w") as f:
        yaml.dump(_settings, f)


def find_centroid(gdf):
    """Find the centroid of polygons, even when in a geographic CRS

    If the crs is geographic (uses lat/lon) then it is converted to a projection before
    calculating the centroid.

    The projected CRS used here is:

    <Projected CRS: EPSG:2163>
    Name: US National Atlas Equal Area
    Axis Info [cartesian]:
    - X[east]: Easting (metre)
    - Y[north]: Northing (metre)
    Area of Use:
    - name: USA
    - bounds: (167.65, 15.56, -65.69, 74.71)
    Coordinate Operation:
    - name: US National Atlas Equal Area
    - method: Lambert Azimuthal Equal Area (Spherical)
    Datum: Not specified (based on Clarke 1866 Authalic Sphere)
    - Ellipsoid: Clarke 1866 Authalic Sphere
    - Prime Meridian: Greenwich

    Parameters
    ----------
    gdf : GeoDataFrame
        A gdf with a geometry column.

    Returns
    -------
    GeoSeries
        A GeoSeries of centroid Points.
    """

    crs = gdf.crs

    if crs.is_geographic:
        _gdf = gdf.to_crs("EPSG:2163")
        centroid = _gdf.centroid
        centroid = centroid.to_crs(crs)
    else:
        centroid = gdf.centroid

    return centroid


def regions_to_keep(
    model_regions: List[str], region_aggregations: dict = {}
) -> Tuple[list, dict]:
    """Create a list of all IPM regions that are used in the model, either as single
    regions or as part of a user-defined model region. Also includes the aggregate
    regions defined by user.

    Parameters
    ----------
    settings : dict
        User-defined parameters from a settings YAML file with keys "model_regions" and
        "region_aggregations".

    Returns
    -------
    list
        All of the IPM regions and user defined model regions.
    """
    # Settings has a dictionary of lists for regional aggregations.
    region_agg_map = reverse_dict_of_lists(region_aggregations)

    # IPM regions to keep - single in model_regions plus those aggregated by the user
    keep_regions = [
        x
        for x in model_regions + list(region_agg_map)
        if x not in region_agg_map.values()
    ]
    return keep_regions, region_agg_map


def build_case_id_name_map(settings: dict) -> dict:
    """Make a dictionary mapping of case IDs and case names from a CSV file

    Parameters
    ----------
    settings : dict
        Settings parameters. Must include `input_folder` and `case_id_description_fn`

    Returns
    -------
    dict
        Mapping of case id to case name
    """
    case_id_name_df = pd.read_csv(
        Path(settings["input_folder"]) / settings["case_id_description_fn"],
        index_col=0,
        squeeze=True,
    )
    case_id_name_df = case_id_name_df.str.replace(" ", "_")
    case_id_name_map = case_id_name_df.to_dict()

    return case_id_name_map


def build_scenario_settings(
    settings: dict, scenario_definitions: pd.DataFrame
) -> Dict[int, Dict[Union[int, str], dict]]:
    """Build a nested dictionary of settings for each planning year/scenario

    Parameters
    ----------
    settings : dict
        The full settings file, including the "settings_management" section with
        alternate values for each scenario
    scenario_definitions : pd.DataFrame
        Values from the csv file defined in the settings file "scenario_definitions_fn"
        parameter. This df has columns corresponding to categories in the
        "settings_management" section of the settings file, with row values defining
        specific case/scenario names.

    Returns
    -------
    dict
        A nested dictionary. The first set of keys are the planning years, the second
        set of keys are the case ID values associated with each case.
    """
    if settings.get("model_periods"):
        model_planning_period_dict = {
            year: (start_year, year) for (start_year, year) in settings["model_periods"]
        }
    elif isinstance(settings.get("model_year"), list) and isinstance(
        settings.get("model_first_planning_year"), list
    ):
        model_planning_period_dict = {
            year: (start_year, year)
            for year, start_year in zip(
                settings["model_year"], settings["model_first_planning_year"]
            )
        }
    else:
        raise KeyError(
            "To build a dictionary of scenario settings your settings file should include "
            "either the key 'model_periods' (a list of 2-element lists) or the keys "
            "'model_year' and 'model_first_planning_year' (each a list of years)."
        )

    case_id_name_map = build_case_id_name_map(settings)

    scenario_settings = {}
    for year in scenario_definitions["year"].unique():
        scenario_settings[year] = {}
        planning_year_settings_management = settings["settings_management"][year]

        # Create a dictionary with keys of things that change (e.g. ccs_capex) and
        # values of nested dictionaries that give case_id: scenario name
        planning_year_scenario_definitions_dict = (
            scenario_definitions.loc[scenario_definitions.year == year]
            .set_index("case_id")
            .to_dict()
        )
        planning_year_scenario_definitions_dict.pop("year")

        for case_id in scenario_definitions.query("year==@year")["case_id"].unique():
            _settings = deepcopy(settings)

            if "all_cases" in planning_year_settings_management:
                new_parameter = planning_year_settings_management["all_cases"]
                _settings = update_dictionary(_settings, new_parameter)

            # Add the scenario definition values to the settings files
            # e.g.
            # case_id	year	demand_response	growth	tx_expansion	ng_price
            # p1	    2030	moderate	            normal	high	reference
            case_scenario_definitions = scenario_definitions.loc[
                (scenario_definitions.case_id == case_id)
                & (scenario_definitions.year == year),
                :,
            ]
            for col in scenario_definitions.columns:
                _settings[col] = case_scenario_definitions.squeeze().at[col]

            modified_settings = []
            for (
                category,
                case_value_dict,
            ) in planning_year_scenario_definitions_dict.items():
                # key is the category e.g. ccs_capex, case_value_dict is p1: mid
                try:
                    case_value = case_value_dict[case_id]
                    new_parameter = (
                        planning_year_settings_management[category][case_value] or {}
                    )

                    try:
                        settings_keys = list(flatten(new_parameter).keys())
                    except AttributeError:
                        settings_keys = {}

                    for key in settings_keys:
                        assert (
                            key not in modified_settings
                        ), f"The settings key {key} is modified twice in case id {case_id}"

                        modified_settings.append(key)

                    if new_parameter is not None:
                        _settings = update_dictionary(_settings, new_parameter)
                    # print(_settings[list(new_parameter.keys())[0]])

                except KeyError:
                    pass

            _settings["model_first_planning_year"] = model_planning_period_dict[year][0]
            _settings["model_year"] = model_planning_period_dict[year][1]
            _settings["case_name"] = case_id_name_map[case_id]
            scenario_settings[year][case_id] = _settings

    return scenario_settings


def remove_feb_29(df: pd.DataFrame) -> pd.DataFrame:
    """Remove Feb 29 from a wide format leap-year dataseries

    Parameters
    ----------
    df : pd.DataFrame
        A wide format dataframe with 8784 columns

    Returns
    -------
    pd.DataFrame
        The same dataframe but without the 24 hours in Feb 29 and only 8760 rows.
    """
    idx_start = df.index.min()
    idx_name = df.index.name
    df["datetime"] = pd.date_range(start="2012-01-01", freq="H", periods=8784)

    df = df.loc[~((df.datetime.dt.month == 2) & (df.datetime.dt.day == 29)), :]
    df.index = range(idx_start, idx_start + 8760)
    df.index.name = idx_name

    return df.drop(columns=["datetime"])


def load_ipm_shapefile(settings: dict, path: Union[str, Path] = IPM_GEOJSON_PATH):
    """
    Load the shapefile of IPM regions

    Parameters
    ----------
    settings : dict
        User-defined parameters from a settings YAML file. This is where any region
        aggregations would be defined.
    path : Union[str, Path]
        Path, loction, or URL of the IPM shapefile/geojson to load. Default value is
        a simplified geojson stored in the PowerGenome data folder.

    Returns
    -------
    geodataframe
        Regions to use in the study with the matching geometry for each.
    """
    keep_regions, region_agg_map = regions_to_keep(
        settings["model_regions"], settings.get("region_aggregations", {}) or {}
    )

    ipm_regions = gpd.read_file(path)
    ipm_regions = ipm_regions.rename(columns={"IPM_Region": "region"})

    if settings.get("user_region_geodata_fn"):
        logger.info("Appending user regions to IPM Regions")
        user_regions = gpd.read_file(
            Path(settings["input_folder"]) / settings["user_region_geodata_fn"]
        )
        if "region" not in user_regions.columns:
            raise KeyError(
                "The user supplied region geodata file does not include the "
                "property 'region' for any of the region polygons! User region "
                "geodata can not be appropriately mapped to model regions."
            )
        user_regions = user_regions.to_crs(ipm_regions.crs)
        ipm_regions = ipm_regions.append(user_regions)

    model_regions_gdf = ipm_regions.loc[ipm_regions["region"].isin(keep_regions)]
    model_regions_gdf = map_agg_region_names(
        model_regions_gdf, region_agg_map, "region", "model_region"
    ).reset_index(drop=True)

    return model_regions_gdf


def deep_freeze(thing):
    """
    https://stackoverflow.com/a/66729248/3393071
    """
    from collections.abc import Collection, Mapping, Hashable
    from frozendict import frozendict

    if thing is None or isinstance(thing, str):
        return thing
    elif isinstance(thing, Mapping):
        return frozendict({k: deep_freeze(v) for k, v in thing.items()})
    elif isinstance(thing, Collection):
        return tuple(deep_freeze(i) for i in thing)
    elif not isinstance(thing, Hashable):
        raise TypeError(f"unfreezable type: '{type(thing)}'")
    else:
        return thing


def deep_freeze_args(func):
    """
    https://stackoverflow.com/a/66729248/3393071
    """
    import functools

    @functools.wraps(func)
    def wrapped(*args, **kwargs):
        return func(*deep_freeze(args), **deep_freeze(kwargs))

    return wrapped


def find_region_col(cols: Union[pd.Index, List[str]], context: str = None) -> str:
    """Find the column name that identifies regions.

    DataFrame, geospatial objects, etc might have different names for the region column.
    To retain some flexibility, only require that the region column has the string
    "region" in it (case insensitive).

    Raise an error if more than one column contains the string "region". If `context` is
    provided, include it in the error message for users.

    Parameters
    ----------
    cols : Iterable[str]
        DataFrame columns or other iterable sequence.
    context : str, optional
        Information about the sequence of names that can help a user understand what
        type of object might have multiple names containing "region", by default None

    Returns
    -------
    str
        Name of the column that identifies regions.

    Raises
    ------
    ValueError
        More than one column contains the string "region".
    ValueError
        No column contains the string "region".
    """

    region_col = [c for c in cols if "region" in c.lower()]
    if len(region_col) > 1:
        s = (
            "When attempting to identify the appropriate region columns, more than one "
            f"column in this dataframe includes the string 'region' ({region_col})."
        )
        if context:
            s += f"\n\nContext: {context}"

        raise ValueError(s)
    elif len(region_col) == 0:
        s = (
            "No columns contain the required string 'region'. The DataFrame columns "
            f"are ({cols})."
        )
        if context:
            s += f"\n\nContext: {context}"

        raise ValueError(s)
    else:
        return region_col[0]


def remove_leading_zero(id: Union[str, int]) -> Union[str, int]:
    """Remove leading zero from IDs that are otherwise integers.

    There is a discrepency between some generator IDs in PUDL and 860m where they are
    listed with a leading zero in one and an integer in the other. To better match,
    strip zeros from IDs that would be an integer without them.

    Parameters
    ----------
    id : Union[str, int]
        An integer or string identifier

    Returns
    -------
    Union[str, int]
        Either the original ID (if integer or non-numeric string) or an integer version
        of the ID with leading zeros removed
    """
    if isinstance(id, int):
        return id
    elif id.isnumeric():
        id = id.strip("0")
    return id


# powergenome/co2_pipeline_cost.py
"Load data from file for pipeline/transport/injection costs"

from pathlib import Path
from typing import Dict, List

import pandas as pd
import logging
from powergenome.price_adjustment import inflation_price_adjustment

from powergenome.util import snake_case_col, snake_case_str
from powergenome.params import DATA_PATHS

logger = logging.getLogger(__name__)


def merge_co2_pipeline_costs(
    df: pd.DataFrame,
    co2_data_path: Path,
    co2_pipeline_filters: List[dict],
    region_aggregations: Dict[str, List[str]] = None,
    fuel_emission_factors: Dict[str, float] = None,
    target_usd_year: int = None,
) -> pd.DataFrame:
    """Merge columns with CO2 pipeline costs to a dataframe.

    Data come from an external file. Filters from the settings are used to select the
    data that will be merged.

    Parameters
    ----------
    df : pd.DataFrame
        Generators data. Should include the rows "technology" and "region", along with a
        heat rate column that has "heat_rate" and "per_mwh" if fuel emission factors
        will be used to convert costs from per tonne to per MWh.
    co2_data_path : Path
        Path to the file with CO2 pipeline cost data
    co2_pipeline_filters : List[dict]
        Filters to apply to the CO2 cost data file
    region_aggregations : Dict[str, List[str]], optional
        Aggregations of multiple base regions to model regions, by default None
    fuel_emission_factors : Dict[str, float], optional
        CO2 emissions per unit of energy for different fuel types, by default None
    target_usd_year : int, optional
        The target year for cost inflation adjustment, by default None

    Returns
    -------
    pd.DataFrame
        The input dataframe (one row per resource) plus columns with cost data for CO2
        pipeline constuction/operation and CO2 disposal
    """
    co2_df = pd.read_csv(co2_data_path)
    co2_df = co2_df.loc[co2_df["parameter"] != "capacity_mw", :]
    if target_usd_year:
        for dollar_year in co2_df["dollar_year"].unique():
            co2_df.loc[
                co2_df["dollar_year"] == dollar_year, "parameter_value"
            ] = inflation_price_adjustment(
                co2_df.loc[co2_df["dollar_year"] == dollar_year, "parameter_value"],
                dollar_year,
                target_usd_year,
            )
    for k, v in (region_aggregations or {}).items():
        co2_df.loc[co2_df["region"].isin(v), "region"] = k

    # When regions are aggregated, use the lowest-cost ipm region
    co2_df = co2_df.sort_values("parameter_value").drop_duplicates(
        [
            "region",
            "technology",
            "tech_detail",
            "with_backbone",
            "percentile",
            "parameter",
        ],
        keep="first",
    )

    # co2_df["resource"] = (co2_df["technology"] + "_" + co2_df["tech_detail"]).lower()
    filter_techs = [
        filt["technology"] + "_" + filt["tech_detail"] for filt in co2_pipeline_filters
    ]
    co2_resources = {}
    for co2_r in filter_techs:
        for model_r in df["technology"].unique():
            if snake_case_str(co2_r) in snake_case_str(model_r):
                co2_resources[co2_r] = model_r

    if not co2_resources:
        logger.warning(
            "You have specified technologies for CO2 pipeline costs in the parameter "
            "'co2_pipeline_filters' but none of the technologies are being used in "
            "your system."
        )
        return df

    df_list = []
    for filt in co2_pipeline_filters:
        # tech = filt["technology"].lower() + "_" + filt["tech_detail"].lower()
        query_str = f"technology == '{filt['technology']}'"
        for k, v in filt.items():
            if k != "technology":
                if isinstance(v, str):
                    query_str += f" and {k} == '{v}'"
                else:
                    query_str += f" and {k} == {v}"

        co2_costs = co2_df.query(query_str)
        co2_costs.loc[:, "technology"] = (
            co2_costs["technology"] + "_" + co2_costs["tech_detail"]
        ).map(co2_resources)
        co2_costs_wide = co2_costs.pivot(
            index=["region", "technology"],
            columns="parameter",
            values="parameter_value",
        ).reset_index()
        df_list.append(co2_costs_wide)

    df_co2_costs = pd.merge(df, pd.concat(df_list), how="left")

    drop_idx = df_co2_costs.loc[
        (df_co2_costs["technology"].isin(co2_resources.values()))
        & (df_co2_costs[co2_costs_wide.columns].isna().any(axis=1))
    ].index

    if not drop_idx.empty:
        for tech, _df in df_co2_costs.loc[drop_idx, :].groupby("technology"):
            regs = _df["region"].to_list()
            f"The CCS resource {tech} is being removed from regions {regs} because cost "
            "data was not included in your CO2 pipeline cost file."
    df_co2_costs = df_co2_costs.drop(index=drop_idx)
    df_co2_costs.loc[:, co2_costs_wide.columns] = df_co2_costs[
        co2_costs_wide.columns
    ].fillna(0)

    mass_cols = [c for c in co2_costs_wide.columns if "tonne" in c]
    if mass_cols and fuel_emission_factors:
        df_co2_costs = mass_to_energy_costs(
            df_co2_costs, mass_cols, fuel_emission_factors
        )

    return df_co2_costs


def mass_to_energy_costs(
    df: pd.DataFrame,
    mass_cols: List[str],
    fuel_emission_factors: Dict[str, float],
    heat_rate_col: str = None,
) -> pd.DataFrame:
    """Convert costs per tonne of CO2 to a MWh basis

    Parameters
    ----------
    df : pd.DataFrame
        Generators dataframe with a fuel column, a heat rate column, and the columns
        provided in the list `mass_cols`
    mass_cols : List[str]
        Names of columns with cost per tonne of CO2 disposal
    fuel_emission_factors : Dict[str, float]
        Emission rate of tonne per unit of energy for fuels used by CCS generators
    heat_rate_col : str, optional
        Name of the column with generator heat rates, by default None

    Returns
    -------
    pd.DataFrame
        Modified copy of the input, with the new column "co2_cost_mwh"
    """
    if not heat_rate_col:
        heat_rate_cols = [
            c for c in df.columns if "heat_rate" in c.lower() and "per_mwh" in c.lower()
        ]
        if len(heat_rate_cols) > 1:
            logger.warning(
                "More than one column with 'heat_rate' and 'per_mwh' was found when "
                "converting CO2 transport costs from per tonne to per MWh. Values are "
                "not being converted."
            )
            return df

        elif not heat_rate_cols:
            logger.warning(
                "Unable to identify a heat rate column when converting CO2 transport costs "
                "from per tonne to per MWh. A heat rate column should contain 'heat_rate' and "
                "'per_mwh' (capitalization does not matter). Values are not being converted."
            )
            return df
        else:
            heat_rate_col = heat_rate_cols[0]

    fuel_col = [c for c in df.columns if c.lower() == "fuel"]
    fuel_col = fuel_col[0]
    for fuel, ef in fuel_emission_factors.items():
        df.loc[df[fuel_col].str.contains(fuel, case=False), "tonne_co2_mwh"] = (
            df.loc[df[fuel_col].str.contains(fuel, case=False), heat_rate_col] * ef
        )

    df["co2_cost_mwh"] = df[mass_cols].sum(axis=1) * df["tonne_co2_mwh"]
    df.loc[:, "co2_cost_mwh"] = df["co2_cost_mwh"].fillna(0)

    return df
