step3.py

import time

import pandas as pd
import pyomo.environ as pyo

from helpers import (
    solve_pyomo_model,
    build_generator_fleet,
    build_wind_availability,
    build_demand_data,
)
from plots_step3 import plot_base_vs_stressed_nodal_prices


def build_step3_inputs(data_folder, hour):
    """
    Build generator, demand, bus, and line data for one selected hour.

    Parameters
    ----------
    data_folder : str
        Path to the input data folder.
    hour : int
        Selected hour for the market-clearing problem.

    Returns
    -------
    generators : pandas.DataFrame
        Generator data with hourly available capacity ``p_max``.
    demands : pandas.DataFrame
        Demand data for the selected hour.
    buses : pandas.DataFrame
        Bus table.
    lines : pandas.DataFrame
        Line table with susceptance and line limits.
    info : dict
        Summary information for the selected hour.
    """
    generators = build_generator_fleet(data_folder).copy()
    wind_availability = build_wind_availability(data_folder).copy()
    demands = build_demand_data(data_folder, hour=hour).copy()

    # Read and clean transmission line data.
    lines = pd.read_csv(f"{data_folder}/rts24_lines.csv").rename(
        columns={
            "From": "from_node",
            "To": "to_node",
            "Reactance_pu": "x",
            "Capacity_MVA": "f_max",
        }
    )

    lines["line_id"] = [f"L{i + 1}" for i in range(len(lines))]
    lines[["from_node", "to_node"]] = lines[["from_node", "to_node"]].astype(int)
    lines[["x", "f_max"]] = lines[["x", "f_max"]].astype(float)
    lines["b"] = 1.0 / lines["x"]

    # Build the full bus list from generators, demand, and lines.
    bus_ids = sorted(
        set(generators["node"])
        .union(set(demands["node"]))
        .union(set(lines["from_node"]))
        .union(set(lines["to_node"]))
    )
    buses = pd.DataFrame({"node": bus_ids})

    # Start from nominal generator capacity.
    generators["p_max"] = generators["p_nom"].astype(float)

    # Adjust wind capacity using the selected-hour wind capacity factors.
    wind_cf_hour = (
        wind_availability.loc[wind_availability["hour"] == hour]
        .set_index("gen_id")["cf"]
    )
    is_wind = generators["type"].astype(str).str.lower() == "wind"
    generators.loc[is_wind, "p_max"] = (
        generators.loc[is_wind, "gen_id"].map(wind_cf_hour).fillna(0.0)
        * generators.loc[is_wind, "p_nom"]
    )

    info = {
        "hour": hour,
        "system_demand": demands["q_max"].sum(),
        "number_of_buses": len(buses),
        "number_of_lines": len(lines),
        "average_wind_cf": float(wind_cf_hour.mean()) if not wind_cf_hour.empty else 0.0,
    }

    return generators, demands, buses, lines, info


def apply_line_changes(lines, line_changes):
    """
    Apply line-capacity changes to a line table.

    Parameters
    ----------
    lines : pandas.DataFrame
        Original line table.
    line_changes : dict or None
        Mapping ``line_id -> value``.
        If ``value <= 1.0``, it is interpreted as a multiplicative factor.
        If ``value > 1.0``, it is interpreted as a new absolute capacity.

    Returns
    -------
    lines_mod : pandas.DataFrame
        Modified line table.
    """
    lines_mod = lines.copy()

    if line_changes:
        for line_id, value in line_changes.items():
            mask = lines_mod["line_id"] == line_id
            if value <= 1.0:
                lines_mod.loc[mask, "f_max"] = lines_mod.loc[mask, "f_max"] * value
            else:
                lines_mod.loc[mask, "f_max"] = float(value)

    return lines_mod


def extract_model_stats(model, solve_time):
    """
    Extract model statistics for reporting.

    Parameters
    ----------
    model : pyomo.environ.ConcreteModel
        Solved Pyomo model.
    solve_time : float
        Computational time in seconds.

    Returns
    -------
    stats : dict
        Dictionary with number of variables, number of constraints, and solve time.
    """
    return {
        "n_variables": sum(
            1 for _ in model.component_data_objects(pyo.Var, active=True)
        ),
        "n_constraints": sum(
            1 for _ in model.component_data_objects(pyo.Constraint, active=True)
        ),
        "solve_time": solve_time,
    }


# ============================================================
# NODAL MODEL
# ============================================================


def build_step3_nodal_model(generators, demands, buses, lines):
    """
    Build the nodal welfare-maximizing market model with DC power flow.

    Parameters
    ----------
    generators : pandas.DataFrame
        Generator input data.
    demands : pandas.DataFrame
        Demand input data.
    buses : pandas.DataFrame
        Bus table.
    lines : pandas.DataFrame
        Line table.

    Returns
    -------
    model : pyomo.environ.ConcreteModel
        Nodal optimization model.
    """
    model = pyo.ConcreteModel()

    generators = generators.copy()
    demands = demands.copy()
    buses = buses.copy()
    lines = lines.copy()

    generators["gen_id"] = generators["gen_id"].astype(str)
    demands["load_id"] = demands["load_id"].astype(str)
    lines["line_id"] = lines["line_id"].astype(str)

    model.G = pyo.Set(initialize=generators["gen_id"].tolist())
    model.D = pyo.Set(initialize=demands["load_id"].tolist())
    model.N = pyo.Set(initialize=buses["node"].tolist())
    model.L = pyo.Set(initialize=lines["line_id"].tolist())

    # Generator, demand, and line parameters.
    model.Pmax = pyo.Param(
        model.G, initialize=generators.set_index("gen_id")["p_max"].to_dict()
    )
    model.MCost = pyo.Param(
        model.G, initialize=generators.set_index("gen_id")["mcost"].to_dict()
    )
    model.GenNode = pyo.Param(
        model.G, initialize=generators.set_index("gen_id")["node"].to_dict()
    )

    model.Qmax = pyo.Param(
        model.D, initialize=demands.set_index("load_id")["q_max"].to_dict()
    )
    model.BidPrice = pyo.Param(
        model.D, initialize=demands.set_index("load_id")["bid_price"].to_dict()
    )
    model.LoadNode = pyo.Param(
        model.D, initialize=demands.set_index("load_id")["node"].to_dict()
    )

    model.LineFrom = pyo.Param(
        model.L, initialize=lines.set_index("line_id")["from_node"].to_dict()
    )
    model.LineTo = pyo.Param(
        model.L, initialize=lines.set_index("line_id")["to_node"].to_dict()
    )
    model.LineB = pyo.Param(
        model.L, initialize=lines.set_index("line_id")["b"].to_dict()
    )
    model.LineFmax = pyo.Param(
        model.L, initialize=lines.set_index("line_id")["f_max"].to_dict()
    )

    # Decision variables.
    model.p = pyo.Var(model.G, domain=pyo.NonNegativeReals)
    model.q = pyo.Var(model.D, domain=pyo.NonNegativeReals)
    model.theta = pyo.Var(model.N, domain=pyo.Reals)
    model.f = pyo.Var(model.L, domain=pyo.Reals)

    # Maximize welfare.
    model.Objective = pyo.Objective(
        expr=sum(model.BidPrice[d] * model.q[d] for d in model.D)
        - sum(model.MCost[g] * model.p[g] for g in model.G),
        sense=pyo.maximize,
    )

    model.GeneratorCapacity = pyo.Constraint(
        model.G, rule=lambda m, g: m.p[g] <= m.Pmax[g]
    )
    model.DemandLimit = pyo.Constraint(
        model.D, rule=lambda m, d: m.q[d] <= m.Qmax[d]
    )

    # DC line flow equations and line limits.
    def line_flow_rule(m, l):
        return m.f[l] == m.LineB[l] * (
            m.theta[m.LineFrom[l]] - m.theta[m.LineTo[l]]
        )

    model.LineFlow = pyo.Constraint(model.L, rule=line_flow_rule)
    model.LineUpper = pyo.Constraint(
        model.L, rule=lambda m, l: m.f[l] <= m.LineFmax[l]
    )
    model.LineLower = pyo.Constraint(
        model.L, rule=lambda m, l: m.f[l] >= -m.LineFmax[l]
    )

    # Nodal balance at every bus.
    def nodal_balance_rule(m, n):
        generation = sum(m.p[g] for g in m.G if m.GenNode[g] == n)
        demand = sum(m.q[d] for d in m.D if m.LoadNode[d] == n)
        inflow = sum(m.f[l] for l in m.L if m.LineTo[l] == n)
        outflow = sum(m.f[l] for l in m.L if m.LineFrom[l] == n)
        return generation + inflow == demand + outflow

    model.NodalBalance = pyo.Constraint(model.N, rule=nodal_balance_rule)

    # Fix one reference angle.
    slack_bus = int(buses["node"].min())
    model.ReferenceAngle = pyo.Constraint(expr=model.theta[slack_bus] == 0.0)

    model.dual = pyo.Suffix(direction=pyo.Suffix.IMPORT)

    return model


def solve_step3_nodal(generators, demands, buses, lines, solver_name):
    """
    Solve the nodal Step 3 model.

    Parameters
    ----------
    generators : pandas.DataFrame
        Generator input data.
    demands : pandas.DataFrame
        Demand input data.
    buses : pandas.DataFrame
        Bus table.
    lines : pandas.DataFrame
        Line table.
    solver_name : str
        Solver name.

    Returns
    -------
    model : pyomo.environ.ConcreteModel
        Solved nodal model.
    solve_time : float
        Computational time in seconds.
    """
    model = build_step3_nodal_model(generators, demands, buses, lines)

    start_time = time.perf_counter()
    solve_pyomo_model(model, solver_name=solver_name)
    solve_time = time.perf_counter() - start_time

    return model, solve_time


def extract_nodal_results(model, generators, demands, buses, lines, solve_time):
    """
    Extract detailed nodal results and summary metrics.

    Parameters
    ----------
    model : pyomo.environ.ConcreteModel
        Solved nodal model.
    generators : pandas.DataFrame
        Generator input data.
    demands : pandas.DataFrame
        Demand input data.
    buses : pandas.DataFrame
        Bus table.
    lines : pandas.DataFrame
        Line table.
    solve_time : float
        Computational time in seconds.

    Returns
    -------
    results : dict
        Detailed nodal result tables and summary metrics.
    """
    generation = generators.copy()
    generation["gen_id"] = generation["gen_id"].astype(str)
    generation["dispatch"] = generation["gen_id"].map(lambda g: pyo.value(model.p[g]))
    generation["unused_capacity"] = generation["p_max"] - generation["dispatch"]

    demand = demands.copy()
    demand["load_id"] = demand["load_id"].astype(str)
    demand["served_demand"] = demand["load_id"].map(lambda d: pyo.value(model.q[d]))
    demand["curtailed_demand"] = demand["q_max"] - demand["served_demand"]

    buses_out = buses.copy()
    buses_out["theta"] = buses_out["node"].map(lambda n: pyo.value(model.theta[n]))
    buses_out["nodal_price"] = buses_out["node"].map(
        lambda n: -model.dual.get(model.NodalBalance[n], 0.0)
    )

    nodal_demand = (
        demand.groupby("node", as_index=False)[
            ["q_max", "served_demand", "curtailed_demand"]
        ]
        .sum()
        .rename(columns={"q_max": "demand"})
    )
    buses_out = buses_out.merge(nodal_demand, on="node", how="left").fillna(
        {"demand": 0.0, "served_demand": 0.0, "curtailed_demand": 0.0}
    )

    price_by_node = buses_out.set_index("node")["nodal_price"].to_dict()

    lines_out = lines.copy()
    lines_out["line_id"] = lines_out["line_id"].astype(str)
    lines_out["flow"] = lines_out["line_id"].map(lambda l: pyo.value(model.f[l]))
    lines_out["loading_abs"] = lines_out["flow"].abs()
    lines_out["loading_pct"] = 100 * lines_out["loading_abs"] / lines_out["f_max"]
    lines_out["is_congested"] = (
        lines_out["loading_abs"] >= (lines_out["f_max"] - 1e-6)
    )

    # Add price-based economic quantities.
    generation["nodal_price"] = generation["node"].map(price_by_node)
    generation["revenue"] = generation["dispatch"] * generation["nodal_price"]
    generation["variable_cost"] = generation["dispatch"] * generation["mcost"]
    generation["profit"] = generation["revenue"] - generation["variable_cost"]

    demand["nodal_price"] = demand["node"].map(price_by_node)
    demand["utility"] = demand["served_demand"] * (
        demand["bid_price"] - demand["nodal_price"]
    )

    stats = extract_model_stats(model, solve_time)

    return {
        "generation": generation,
        "demand": demand,
        "buses": buses_out,
        "lines": lines_out,
        "summary": {
            "total_generation": generation["dispatch"].sum(),
            "total_demand": demand["q_max"].sum(),
            "total_served_demand": demand["served_demand"].sum(),
            "total_curtailed_demand": demand["curtailed_demand"].sum(),
            "total_welfare": pyo.value(model.Objective),
            "min_nodal_price": buses_out["nodal_price"].min(),
            "max_nodal_price": buses_out["nodal_price"].max(),
            "number_of_congested_lines": int(lines_out["is_congested"].sum()),
            "total_generator_profit": generation["profit"].sum(),
            "total_demand_utility": demand["utility"].sum(),
            "n_variables": stats["n_variables"],
            "n_constraints": stats["n_constraints"],
            "solve_time": stats["solve_time"],
        },
    }


# ============================================================
# ZONAL MODEL
# ============================================================


def build_zone_mapping():
    """
    Build the three-zone partition used for the zonal approximation.

    Returns
    -------
    zone_map : dict
        Mapping from node number to zone label.
    """
    zone_map = {}

    for node in [16, 17, 18, 19, 20, 21, 22, 23]:
        zone_map[node] = "Z1"
    for node in [1, 3, 4, 14, 15, 24]:
        zone_map[node] = "Z2"
    for node in [2, 5, 6, 7, 8, 9, 10, 11, 12, 13]:
        zone_map[node] = "Z3"

    return zone_map


def build_interzonal_atc(lines, zone_map):
    """
    Build interzonal ATC values from the underlying network.

    Parameters
    ----------
    lines : pandas.DataFrame
        Line table.
    zone_map : dict
        Mapping from node to zone.

    Returns
    -------
    atc : pandas.DataFrame
        Interzonal ATC table.
    """
    atc_rows = []

    for _, row in lines.iterrows():
        zone_from = zone_map[row["from_node"]]
        zone_to = zone_map[row["to_node"]]

        if zone_from != zone_to:
            z1, z2 = sorted([zone_from, zone_to])
            atc_rows.append(
                {
                    "zone_from": z1,
                    "zone_to": z2,
                    "f_max": row["f_max"],
                }
            )

    atc = (
        pd.DataFrame(atc_rows)
        .groupby(["zone_from", "zone_to"], as_index=False)["f_max"]
        .sum()
    )
    atc["link_id"] = [f"{row.zone_from}_{row.zone_to}" for row in atc.itertuples()]

    return atc


def build_step3_zonal_model(generators, demands, zone_map, atc):
    """
    Build the zonal welfare-maximizing market model with ATC constraints.

    Parameters
    ----------
    generators : pandas.DataFrame
        Generator input data.
    demands : pandas.DataFrame
        Demand input data.
    zone_map : dict
        Mapping from node to zone.
    atc : pandas.DataFrame
        Interzonal ATC table.

    Returns
    -------
    model : pyomo.environ.ConcreteModel
        Zonal optimization model.
    """
    model = pyo.ConcreteModel()

    generators = generators.copy()
    demands = demands.copy()
    atc = atc.copy()

    generators["gen_id"] = generators["gen_id"].astype(str)
    generators["zone"] = generators["node"].map(zone_map)
    demands["load_id"] = demands["load_id"].astype(str)
    demands["zone"] = demands["node"].map(zone_map)
    atc["link_id"] = atc["link_id"].astype(str)

    zones = sorted(set(zone_map.values()))

    model.G = pyo.Set(initialize=generators["gen_id"].tolist())
    model.D = pyo.Set(initialize=demands["load_id"].tolist())
    model.Z = pyo.Set(initialize=zones)
    model.K = pyo.Set(initialize=atc["link_id"].tolist())

    model.Pmax = pyo.Param(
        model.G, initialize=generators.set_index("gen_id")["p_max"].to_dict()
    )
    model.MCost = pyo.Param(
        model.G, initialize=generators.set_index("gen_id")["mcost"].to_dict()
    )
    model.GenZone = pyo.Param(
        model.G,
        initialize=generators.set_index("gen_id")["zone"].to_dict(),
        within=pyo.Any,
    )

    model.Qmax = pyo.Param(
        model.D, initialize=demands.set_index("load_id")["q_max"].to_dict()
    )
    model.BidPrice = pyo.Param(
        model.D, initialize=demands.set_index("load_id")["bid_price"].to_dict()
    )
    model.LoadZone = pyo.Param(
        model.D,
        initialize=demands.set_index("load_id")["zone"].to_dict(),
        within=pyo.Any,
    )

    model.LinkFrom = pyo.Param(
        model.K,
        initialize=atc.set_index("link_id")["zone_from"].to_dict(),
        within=pyo.Any,
    )
    model.LinkTo = pyo.Param(
        model.K,
        initialize=atc.set_index("link_id")["zone_to"].to_dict(),
        within=pyo.Any,
    )
    model.LinkFmax = pyo.Param(
        model.K, initialize=atc.set_index("link_id")["f_max"].to_dict()
    )

    model.p = pyo.Var(model.G, domain=pyo.NonNegativeReals)
    model.q = pyo.Var(model.D, domain=pyo.NonNegativeReals)
    model.f = pyo.Var(model.K, domain=pyo.Reals)

    model.Objective = pyo.Objective(
        expr=sum(model.BidPrice[d] * model.q[d] for d in model.D)
        - sum(model.MCost[g] * model.p[g] for g in model.G),
        sense=pyo.maximize,
    )

    model.GeneratorCapacity = pyo.Constraint(
        model.G, rule=lambda m, g: m.p[g] <= m.Pmax[g]
    )
    model.DemandLimit = pyo.Constraint(
        model.D, rule=lambda m, d: m.q[d] <= m.Qmax[d]
    )
    model.LinkUpper = pyo.Constraint(
        model.K, rule=lambda m, k: m.f[k] <= m.LinkFmax[k]
    )
    model.LinkLower = pyo.Constraint(
        model.K, rule=lambda m, k: m.f[k] >= -m.LinkFmax[k]
    )

    # Zonal balance.
    def zonal_balance_rule(m, z):
        generation = sum(m.p[g] for g in m.G if m.GenZone[g] == z)
        demand = sum(m.q[d] for d in m.D if m.LoadZone[d] == z)
        inflow = sum(m.f[k] for k in m.K if m.LinkTo[k] == z)
        outflow = sum(m.f[k] for k in m.K if m.LinkFrom[k] == z)
        return generation + inflow == demand + outflow

    model.ZonalBalance = pyo.Constraint(model.Z, rule=zonal_balance_rule)

    model.dual = pyo.Suffix(direction=pyo.Suffix.IMPORT)

    return model


def solve_step3_zonal(generators, demands, zone_map, atc, solver_name):
    """
    Solve the zonal Step 3 model.

    Parameters
    ----------
    generators : pandas.DataFrame
        Generator input data.
    demands : pandas.DataFrame
        Demand input data.
    zone_map : dict
        Mapping from node to zone.
    atc : pandas.DataFrame
        Interzonal ATC table.
    solver_name : str
        Solver name.

    Returns
    -------
    model : pyomo.environ.ConcreteModel
        Solved zonal model.
    solve_time : float
        Computational time in seconds.
    """
    model = build_step3_zonal_model(generators, demands, zone_map, atc)

    start_time = time.perf_counter()
    solve_pyomo_model(model, solver_name=solver_name)
    solve_time = time.perf_counter() - start_time

    return model, solve_time


def extract_zonal_results(model, generators, demands, zone_map, atc, solve_time):
    """
    Extract detailed zonal results and summary metrics.

    Parameters
    ----------
    model : pyomo.environ.ConcreteModel
        Solved zonal model.
    generators : pandas.DataFrame
        Generator input data.
    demands : pandas.DataFrame
        Demand input data.
    zone_map : dict
        Mapping from node to zone.
    atc : pandas.DataFrame
        Interzonal ATC table.
    solve_time : float
        Computational time in seconds.

    Returns
    -------
    results : dict
        Detailed zonal result tables and summary metrics.
    """
    generators = generators.copy()
    demands = demands.copy()
    atc = atc.copy()

    generators["gen_id"] = generators["gen_id"].astype(str)
    generators["zone"] = generators["node"].map(zone_map)
    demands["load_id"] = demands["load_id"].astype(str)
    demands["zone"] = demands["node"].map(zone_map)
    atc["link_id"] = atc["link_id"].astype(str)

    generation = generators.copy()
    generation["dispatch"] = generation["gen_id"].map(lambda g: pyo.value(model.p[g]))
    generation["unused_capacity"] = generation["p_max"] - generation["dispatch"]

    demand = demands.copy()
    demand["served_demand"] = demand["load_id"].map(lambda d: pyo.value(model.q[d]))
    demand["curtailed_demand"] = demand["q_max"] - demand["served_demand"]

    zones = sorted(set(zone_map.values()))
    zones_out = pd.DataFrame({"zone": zones})
    zones_out["zonal_price"] = zones_out["zone"].map(
        lambda z: -model.dual.get(model.ZonalBalance[z], 0.0)
    )

    zones_out = zones_out.merge(
        demand.groupby("zone", as_index=False)["q_max"]
        .sum()
        .rename(columns={"q_max": "demand"}),
        on="zone",
        how="left",
    )
    zones_out = zones_out.merge(
        demand.groupby("zone", as_index=False)["served_demand"]
        .sum()
        .rename(columns={"served_demand": "total_served_demand"}),
        on="zone",
        how="left",
    )
    zones_out[["demand", "total_served_demand"]] = zones_out[
        ["demand", "total_served_demand"]
    ].fillna(0.0)

    generation_by_zone = generation.groupby("zone", as_index=False)["dispatch"].sum()
    zones_out = zones_out.merge(
        generation_by_zone.rename(columns={"dispatch": "total_generation"}),
        on="zone",
        how="left",
    )
    zones_out["total_generation"] = zones_out["total_generation"].fillna(0.0)
    zones_out["net_export"] = (
        zones_out["total_generation"] - zones_out["total_served_demand"]
    )

    zonal_price_map = zones_out.set_index("zone")["zonal_price"].to_dict()

    generation["zonal_price"] = generation["zone"].map(zonal_price_map)
    generation["revenue"] = generation["dispatch"] * generation["zonal_price"]
    generation["variable_cost"] = generation["dispatch"] * generation["mcost"]
    generation["profit"] = generation["revenue"] - generation["variable_cost"]

    demand["zonal_price"] = demand["zone"].map(zonal_price_map)
    demand["utility"] = demand["served_demand"] * (
        demand["bid_price"] - demand["zonal_price"]
    )

    links_out = atc.copy()
    links_out["flow"] = links_out["link_id"].map(lambda k: pyo.value(model.f[k]))
    links_out["loading_abs"] = links_out["flow"].abs()
    links_out["loading_pct"] = 100 * links_out["loading_abs"] / links_out["f_max"]
    links_out["is_congested"] = (
        links_out["loading_abs"] >= (links_out["f_max"] - 1e-6)
    )

    stats = extract_model_stats(model, solve_time)

    return {
        "generation": generation,
        "demand": demand,
        "zones": zones_out,
        "links": links_out,
        "summary": {
            "total_generation": generation["dispatch"].sum(),
            "total_demand": demand["q_max"].sum(),
            "total_served_demand": demand["served_demand"].sum(),
            "total_curtailed_demand": demand["curtailed_demand"].sum(),
            "total_welfare": pyo.value(model.Objective),
            "min_zonal_price": zones_out["zonal_price"].min(),
            "max_zonal_price": zones_out["zonal_price"].max(),
            "number_of_congested_links": int(links_out["is_congested"].sum()),
            "total_generator_profit": generation["profit"].sum(),
            "total_demand_utility": demand["utility"].sum(),
            "n_variables": stats["n_variables"],
            "n_constraints": stats["n_constraints"],
            "solve_time": stats["solve_time"],
        },
    }


# ============================================================
# RUNNERS
# ============================================================


def run_nodal_case(data_folder, hour, solver_name, line_changes):
    """
    Run one nodal case.

    Parameters
    ----------
    data_folder : str
        Path to the input data.
    hour : int
        Selected hour.
    solver_name : str
        Solver name.
    line_changes : dict or None
        Line-capacity changes.

    Returns
    -------
    results : dict
        Nodal results.
    info : dict
        Input summary information.
    generators : pandas.DataFrame
        Generator table used.
    demands : pandas.DataFrame
        Demand table used.
    buses : pandas.DataFrame
        Bus table used.
    lines : pandas.DataFrame
        Line table used.
    """
    generators, demands, buses, lines, info = build_step3_inputs(data_folder, hour)
    lines = apply_line_changes(lines, line_changes)

    model, solve_time = solve_step3_nodal(
        generators,
        demands,
        buses,
        lines,
        solver_name,
    )
    results = extract_nodal_results(
        model,
        generators,
        demands,
        buses,
        lines,
        solve_time,
    )

    return results, info, generators, demands, buses, lines


def run_sensitivity_analysis(data_folder, hour, line_changes_by_case, solver_name):
    """
    Run nodal sensitivity cases with user-defined line-capacity changes.

    Parameters
    ----------
    data_folder : str
        Path to the input data.
    hour : int
        Selected hour.
    line_changes_by_case : dict
        Mapping from case name to line-capacity changes.
    solver_name : str
        Solver name.

    Returns
    -------
    summary_df : pandas.DataFrame
        Summary table across nodal sensitivity cases.
    case_results : dict
        Detailed nodal results by case.
    """
    rows = []
    case_results = {}

    for case_name, line_changes in line_changes_by_case.items():
        results, _, _, _, _, lines_used = run_nodal_case(
            data_folder,
            hour,
            solver_name,
            line_changes,
        )

        changed_lines = lines_used.loc[
            lines_used["line_id"].isin(line_changes.keys()),
            ["line_id", "f_max"],
        ]
        changed_lines_desc = (
            ", ".join(f"{row.line_id}={row.f_max:.1f}" for row in changed_lines.itertuples())
            if not changed_lines.empty
            else "none"
        )

        rows.append(
            {
                "case": case_name,
                "changed_lines": changed_lines_desc,
                "total_welfare": results["summary"]["total_welfare"],
                "min_nodal_price": results["summary"]["min_nodal_price"],
                "max_nodal_price": results["summary"]["max_nodal_price"],
                "price_spread": results["summary"]["max_nodal_price"]
                - results["summary"]["min_nodal_price"],
                "number_of_congested_lines": results["summary"]["number_of_congested_lines"],
                "n_variables": results["summary"]["n_variables"],
                "n_constraints": results["summary"]["n_constraints"],
                "solve_time": results["summary"]["solve_time"],
            }
        )
        case_results[case_name] = results

    return pd.DataFrame(rows), case_results


def run_zonal_case(
    data_folder,
    hour,
    solver_name,
    zone_map,
    line_changes,
    atc_factor,
    atc_changes,
):
    """
    Run one zonal case.

    Parameters
    ----------
    data_folder : str
        Path to the input data.
    hour : int
        Selected hour.
    solver_name : str
        Solver name.
    zone_map : dict or None
        Mapping from node to zone.
    line_changes : dict or None
        Line-capacity changes applied before building ATCs.
    atc_factor : float
        Uniform factor applied to all ATCs.
    atc_changes : dict or None
        Link-specific ATC scaling factors.

    Returns
    -------
    results : dict
        Zonal results.
    info : dict
        Input summary information.
    generators : pandas.DataFrame
        Generator table used.
    demands : pandas.DataFrame
        Demand table used.
    zone_map : dict
        Zone mapping used.
    lines : pandas.DataFrame
        Line table used.
    atc : pandas.DataFrame
        Interzonal ATC table used.
    """
    generators, demands, _, lines, info = build_step3_inputs(data_folder, hour)
    lines = apply_line_changes(lines, line_changes)

    zone_map = build_zone_mapping() if zone_map is None else zone_map
    atc = build_interzonal_atc(lines, zone_map).copy()
    atc["f_max"] = atc["f_max"] * atc_factor

    if atc_changes:
        for link_id, factor in atc_changes.items():
            mask = atc["link_id"] == link_id
            atc.loc[mask, "f_max"] = atc.loc[mask, "f_max"] * factor

    model, solve_time = solve_step3_zonal(
        generators,
        demands,
        zone_map,
        atc,
        solver_name,
    )
    results = extract_zonal_results(
        model,
        generators,
        demands,
        zone_map,
        atc,
        solve_time,
    )

    return results, info, generators, demands, zone_map, lines, atc


def run_zonal_atc_sensitivity(
    data_folder,
    hour,
    zone_map,
    line_changes,
    atc_factors,
    target_link_id,
    solver_name,
):
    """
    Run zonal sensitivity by scaling all ATCs or one selected ATC link.

    Parameters
    ----------
    data_folder : str
        Path to the input data.
    hour : int
        Selected hour.
    zone_map : dict or None
        Mapping from node to zone.
    line_changes : dict or None
        Line-capacity changes applied before building ATCs.
    atc_factors : list[float]
        List of ATC scaling factors.
    target_link_id : str or None
        One selected link to scale. If ``None``, all ATCs are scaled.
    solver_name : str
        Solver name.

    Returns
    -------
    summary_df : pandas.DataFrame
        Summary table across zonal sensitivity cases.
    case_results : dict
        Detailed zonal results by factor.
    zone_map : dict
        Zone mapping used.
    """
    rows = []
    case_results = {}

    for factor in atc_factors:
        atc_changes = {target_link_id: factor} if target_link_id is not None else None
        atc_factor = 1.0 if target_link_id is not None else factor

        results, _, _, _, zone_map_used, _, _ = run_zonal_case(
            data_folder,
            hour,
            solver_name,
            zone_map,
            line_changes,
            atc_factor,
            atc_changes,
        )

        zones = results["zones"].copy()
        zonal_price_spread = zones["zonal_price"].max() - zones["zonal_price"].min()

        row = {
            "atc_factor": factor,
            "target_link_id": target_link_id if target_link_id is not None else "ALL",
            "total_welfare": results["summary"]["total_welfare"],
            "min_zonal_price": results["summary"]["min_zonal_price"],
            "max_zonal_price": results["summary"]["max_zonal_price"],
            "zonal_price_spread": zonal_price_spread,
            "number_of_congested_links": results["summary"]["number_of_congested_links"],
            "n_variables": results["summary"]["n_variables"],
            "n_constraints": results["summary"]["n_constraints"],
            "solve_time": results["summary"]["solve_time"],
        }

        for _, zrow in zones.iterrows():
            row[f"price_{zrow['zone']}"] = zrow["zonal_price"]

        rows.append(row)
        case_results[factor] = results
        zone_map = zone_map_used

    return pd.DataFrame(rows), case_results, zone_map


# ============================================================
# SUPPORT / COMPARISON
# ============================================================


def build_nodal_price_formation_table(nodal_results):
    """
    Build a support table for discussing nodal price formation.

    Parameters
    ----------
    nodal_results : dict
        Nodal result dictionary.

    Returns
    -------
    out : pandas.DataFrame
        Support table for nodal price interpretation.
    """
    generation = nodal_results["generation"].copy()
    buses = nodal_results["buses"][["node", "nodal_price"]].copy()

    local = (
        generation.loc[generation["dispatch"] > 1e-6]
        .groupby("node")
        .agg(
            local_min_dispatched_offer=("mcost", "min"),
            local_max_dispatched_offer=("mcost", "max"),
            local_generation=("dispatch", "sum"),
        )
        .reset_index()
    )

    out = buses.merge(local, on="node", how="left")
    out["price_minus_local_min_offer"] = (
        out["nodal_price"] - out["local_min_dispatched_offer"]
    )

    return out.sort_values("node")


def check_dispatch_feasibility_on_full_network(
    generation_dispatch,
    served_demand,
    generators,
    demands,
    buses,
    lines,
    solver_name,
):
    """
    Check whether a fixed dispatch is DC-feasible on the full nodal network.

    Parameters
    ----------
    generation_dispatch : pandas.Series
        Fixed generator dispatch by generator ID.
    served_demand : pandas.Series
        Fixed served demand by load ID.
    generators : pandas.DataFrame
        Generator table.
    demands : pandas.DataFrame
        Demand table.
    buses : pandas.DataFrame
        Bus table.
    lines : pandas.DataFrame
        Line table.
    solver_name : str
        Solver name.

    Returns
    -------
    results : dict
        Feasibility-check results including overload slack by line.
    """
    gen = generators.copy()
    gen["gen_id"] = gen["gen_id"].astype(str)
    dem = demands.copy()
    dem["load_id"] = dem["load_id"].astype(str)
    lines = lines.copy()
    lines["line_id"] = lines["line_id"].astype(str)

    gen_dispatch = generation_dispatch.copy()
    gen_dispatch.index = gen_dispatch.index.astype(str)
    demand_dispatch = served_demand.copy()
    demand_dispatch.index = demand_dispatch.index.astype(str)

    gen["fixed_dispatch"] = gen["gen_id"].map(gen_dispatch).fillna(0.0)
    dem["fixed_served_demand"] = dem["load_id"].map(demand_dispatch).fillna(0.0)

    injections = (
        gen.groupby("node")["fixed_dispatch"].sum()
        .sub(dem.groupby("node")["fixed_served_demand"].sum(), fill_value=0.0)
        .reindex(buses["node"], fill_value=0.0)
    )

    model = pyo.ConcreteModel()
    model.N = pyo.Set(initialize=buses["node"].tolist())
    model.L = pyo.Set(initialize=lines["line_id"].tolist())

    model.LineFrom = pyo.Param(
        model.L, initialize=lines.set_index("line_id")["from_node"].to_dict()
    )
    model.LineTo = pyo.Param(
        model.L, initialize=lines.set_index("line_id")["to_node"].to_dict()
    )
    model.LineB = pyo.Param(
        model.L, initialize=lines.set_index("line_id")["b"].to_dict()
    )
    model.LineFmax = pyo.Param(
        model.L, initialize=lines.set_index("line_id")["f_max"].to_dict()
    )
    model.Injection = pyo.Param(model.N, initialize=injections.to_dict())

    model.theta = pyo.Var(model.N, domain=pyo.Reals)
    model.f = pyo.Var(model.L, domain=pyo.Reals)
    model.overload_plus = pyo.Var(model.L, domain=pyo.NonNegativeReals)
    model.overload_minus = pyo.Var(model.L, domain=pyo.NonNegativeReals)

    # Minimize total overload slack.
    model.Objective = pyo.Objective(
        expr=sum(model.overload_plus[l] + model.overload_minus[l] for l in model.L),
        sense=pyo.minimize,
    )

    model.LineFlow = pyo.Constraint(
        model.L,
        rule=lambda m, l: m.f[l]
        == m.LineB[l] * (m.theta[m.LineFrom[l]] - m.theta[m.LineTo[l]]),
    )
    model.LineUpper = pyo.Constraint(
        model.L,
        rule=lambda m, l: m.f[l] <= m.LineFmax[l] + m.overload_plus[l],
    )
    model.LineLower = pyo.Constraint(
        model.L,
        rule=lambda m, l: m.f[l] >= -m.LineFmax[l] - m.overload_minus[l],
    )

    def nodal_balance_rule(m, n):
        inflow = sum(m.f[l] for l in m.L if m.LineTo[l] == n)
        outflow = sum(m.f[l] for l in m.L if m.LineFrom[l] == n)
        return m.Injection[n] + inflow == outflow

    model.NodalBalance = pyo.Constraint(model.N, rule=nodal_balance_rule)

    slack_bus = int(buses["node"].min())
    model.ReferenceAngle = pyo.Constraint(expr=model.theta[slack_bus] == 0.0)

    solve_pyomo_model(model, solver_name=solver_name)

    lines_out = lines.copy()
    lines_out["flow"] = lines_out["line_id"].map(lambda l: pyo.value(model.f[l]))
    lines_out["loading_abs"] = lines_out["flow"].abs()
    lines_out["loading_pct_vs_limit"] = (
        100 * lines_out["loading_abs"] / lines_out["f_max"]
    )
    lines_out["overload_plus"] = lines_out["line_id"].map(
        lambda l: pyo.value(model.overload_plus[l])
    )
    lines_out["overload_minus"] = lines_out["line_id"].map(
        lambda l: pyo.value(model.overload_minus[l])
    )
    lines_out["overload_mw"] = lines_out[["overload_plus", "overload_minus"]].sum(axis=1)
    lines_out["is_overloaded"] = lines_out["overload_mw"] > 1e-6

    total_overload = lines_out["overload_mw"].sum()
    max_overload = lines_out["overload_mw"].max()

    return {
        "lines": lines_out,
        "summary": {
            "total_line_overload_mw": total_overload,
            "max_line_overload_mw": max_overload,
            "number_of_overloaded_lines": int(lines_out["is_overloaded"].sum()),
            "is_feasible": total_overload <= 1e-6,
        },
    }


def compare_nodal_and_zonal(nodal_results, zonal_results):
    """
    Build a compact comparison table for nodal and zonal outcomes.

    Parameters
    ----------
    nodal_results : dict
        Nodal result dictionary.
    zonal_results : dict
        Zonal result dictionary.

    Returns
    -------
    comparison : pandas.DataFrame
        Compact comparison table.
    """
    comparison = pd.DataFrame(
        {
            "framework": ["nodal", "zonal"],
            "total_welfare": [
                nodal_results["summary"]["total_welfare"],
                zonal_results["summary"]["total_welfare"],
            ],
            "total_generation": [
                nodal_results["summary"]["total_generation"],
                zonal_results["summary"]["total_generation"],
            ],
            "total_served_demand": [
                nodal_results["summary"]["total_served_demand"],
                zonal_results["summary"]["total_served_demand"],
            ],
            "total_demand": [
                nodal_results["summary"]["total_demand"],
                zonal_results["summary"]["total_demand"],
            ],
            "min_price": [
                nodal_results["summary"]["min_nodal_price"],
                zonal_results["summary"]["min_zonal_price"],
            ],
            "max_price": [
                nodal_results["summary"]["max_nodal_price"],
                zonal_results["summary"]["max_zonal_price"],
            ],
            "total_generator_profit": [
                nodal_results["generation"]["profit"].sum(),
                zonal_results["generation"]["profit"].sum(),
            ],
        }
    )
    comparison["price_spread"] = comparison["max_price"] - comparison["min_price"]

    return comparison


def calculate_dispatch_adjustment_proxy(nodal_results, zonal_results):
    """
    Calculate a redispatch proxy based on absolute dispatch differences.

    Parameters
    ----------
    nodal_results : dict
        Nodal result dictionary.
    zonal_results : dict
        Zonal result dictionary.

    Returns
    -------
    diff : pandas.DataFrame
        Dispatch difference by generator.
    total_adjustment : float
        Sum of absolute dispatch differences.
    """
    nodal_dispatch = nodal_results["generation"].set_index("gen_id")["dispatch"]
    zonal_dispatch = zonal_results["generation"].set_index("gen_id")["dispatch"]

    diff = (
        nodal_dispatch.sub(zonal_dispatch, fill_value=0.0)
        .abs()
        .rename("dispatch_difference")
        .reset_index()
    )

    return diff, diff["dispatch_difference"].sum()


def build_profit_comparison_by_type(nodal_results, zonal_results):
    """
    Compare profits by generation technology.

    Parameters
    ----------
    nodal_results : dict
        Nodal result dictionary.
    zonal_results : dict
        Zonal result dictionary.

    Returns
    -------
    out : pandas.DataFrame
        Profit comparison by technology.
    """
    nodal_profit = (
        nodal_results["generation"].groupby("type")["profit"].sum().rename("nodal_profit")
    )
    zonal_profit = (
        zonal_results["generation"].groupby("type")["profit"].sum().rename("zonal_profit")
    )

    out = pd.concat([nodal_profit, zonal_profit], axis=1).fillna(0.0).reset_index()
    out["difference"] = out["zonal_profit"] - out["nodal_profit"]

    return out


# ============================================================
# PRINT HELPERS
# ============================================================


def print_nodal_summary(results, info, title):
    """
    Print a compact nodal case summary.

    Parameters
    ----------
    results : dict
        Nodal result dictionary.
    info : dict
        Input summary information.
    title : str
        Title printed above the summary.

    Returns
    -------
    None
    """
    conventional_profit = (
        results["generation"]
        .loc[results["generation"]["type"] == "Conventional", "profit"]
        .sum()
    )
    wind_profit = (
        results["generation"]
        .loc[results["generation"]["type"] == "Wind", "profit"]
        .sum()
    )
    price_spread = (
        results["summary"]["max_nodal_price"] - results["summary"]["min_nodal_price"]
    )

    print(f"\n{title}")
    print("-" * 60)
    print(f"Hour: {info['hour']}")
    print(f"Maximum demand: {results['demand']['q_max'].sum():.2f} MW")
    print(f"Served demand: {results['demand']['served_demand'].sum():.2f} MW")
    print(f"Curtailed demand: {results['demand']['curtailed_demand'].sum():.2f} MW")
    print(f"Average wind CF: {info['average_wind_cf']:.3f}")
    print(f"Total welfare: {results['summary']['total_welfare']:.2f}")
    print(
        f"Nodal price range: {results['summary']['min_nodal_price']:.2f} "
        f"to {results['summary']['max_nodal_price']:.2f} $/MWh"
    )
    print(f"Price spread: {price_spread:.2f} $/MWh")
    print(f"Total generator profit: {results['summary']['total_generator_profit']:.2f} $")
    print(f"Conventional profit: {conventional_profit:.2f} $")
    print(f"Wind profit: {wind_profit:.2f} $")
    print(f"Congested lines: {results['summary']['number_of_congested_lines']}")
    print(f"Number of variables: {results['summary']['n_variables']}")
    print(f"Number of constraints: {results['summary']['n_constraints']}")
    print(f"Solve time: {results['summary']['solve_time']:.6f} s")

    print("\nNodal prices")
    print(
        results["buses"][["node", "nodal_price"]]
        .sort_values("node")
        .to_string(index=False)
    )

    print("\nTop 10 most loaded lines")
    print(
        results["lines"][
            ["line_id", "from_node", "to_node", "flow", "f_max", "loading_pct", "is_congested"]
        ]
        .sort_values("loading_pct", ascending=False)
        .head(10)
        .to_string(index=False)
    )


def print_zonal_summary(results, info, title):
    """
    Print a compact zonal case summary.

    Parameters
    ----------
    results : dict
        Zonal result dictionary.
    info : dict
        Input summary information.
    title : str
        Title printed above the summary.

    Returns
    -------
    None
    """
    conventional_profit = (
        results["generation"]
        .loc[results["generation"]["type"] == "Conventional", "profit"]
        .sum()
    )
    wind_profit = (
        results["generation"]
        .loc[results["generation"]["type"] == "Wind", "profit"]
        .sum()
    )
    price_spread = (
        results["summary"]["max_zonal_price"] - results["summary"]["min_zonal_price"]
    )

    print(f"\n{title}")
    print("-" * 60)
    print(f"Hour: {info['hour']}")
    print(f"Total demand: {results['summary']['total_demand']:.2f} MW")
    print(f"Served demand: {results['summary']['total_served_demand']:.2f} MW")
    print(f"Curtailed demand: {results['summary']['total_curtailed_demand']:.2f} MW")
    print(f"Total welfare: {results['summary']['total_welfare']:.2f}")
    print(
        f"Zonal price range: {results['summary']['min_zonal_price']:.2f} "
        f"to {results['summary']['max_zonal_price']:.2f} $/MWh"
    )
    print(f"Price spread: {price_spread:.2f} $/MWh")
    print(f"Total generator profit: {results['summary']['total_generator_profit']:.2f} $")
    print(f"Conventional profit: {conventional_profit:.2f} $")
    print(f"Wind profit: {wind_profit:.2f} $")
    print(f"Congested links: {results['summary']['number_of_congested_links']}")
    print(f"Number of variables: {results['summary']['n_variables']}")
    print(f"Number of constraints: {results['summary']['n_constraints']}")
    print(f"Solve time: {results['summary']['solve_time']:.6f} s")

    print("\nZonal results")
    print(results["zones"].to_string(index=False))

    print("\nInterzonal flows")
    print(results["links"].to_string(index=False))


def print_framework_comparison(nodal_results, zonal_results, feasibility_results):
    """
    Print nodal-zonal comparison tables.

    Parameters
    ----------
    nodal_results : dict
        Nodal result dictionary.
    zonal_results : dict
        Zonal result dictionary.
    feasibility_results : dict
        Ex-post feasibility check results.

    Returns
    -------
    None
    """
    comparison = compare_nodal_and_zonal(nodal_results, zonal_results)
    profit_table = build_profit_comparison_by_type(nodal_results, zonal_results)
    dispatch_diff, dispatch_adjustment_proxy = calculate_dispatch_adjustment_proxy(
        nodal_results,
        zonal_results,
    )

    print("\nNODAL VS ZONAL COMPARISON")
    print("-" * 60)
    print(comparison.to_string(index=False))

    print(f"\nDispatch adjustment proxy [MW]: {dispatch_adjustment_proxy:.2f}")
    print(
        f"Ex-post feasibility on full network: {feasibility_results['summary']['is_feasible']} | "
        f"Overloaded lines: {feasibility_results['summary']['number_of_overloaded_lines']} | "
        f"Total overload slack [MW]: {feasibility_results['summary']['total_line_overload_mw']:.2f} | "
        f"Max line overload [MW]: {feasibility_results['summary']['max_line_overload_mw']:.2f}"
    )

    print("\nProfit comparison by technology")
    print(profit_table.to_string(index=False))

    print("\nTop 10 dispatch differences")
    print(
        dispatch_diff.sort_values("dispatch_difference", ascending=False)
        .head(10)
        .to_string(index=False)
    )

    if feasibility_results["summary"]["number_of_overloaded_lines"] > 0:
        print("\nTop overloaded lines in ex-post feasibility check")
        print(
            feasibility_results["lines"]
            .loc[
                lambda df: df["overload_mw"] > 1e-6,
                ["line_id", "from_node", "to_node", "f_max", "flow", "overload_mw"],
            ]
            .sort_values("overload_mw", ascending=False)
            .head(10)
            .to_string(index=False)
        )


def main():
    """
    Run the full Step 3 analysis.
    """
    data_folder = "data"
    solver_name = "gurobi"
    hour = 19

    congestion_cases = {
        "base": {},
        "reduced_corridors": {
            "L25": 400.0,
            "L23": 250.0,
            "L22": 250.0,
        },
    }

    zonal_atc_factors = [1.0, 0.5, 0.15, 0.01]
    selected_atc_factor = 0.15

    # 1) Base nodal case.
    base_nodal, info, generators, demands, buses, lines = run_nodal_case(
        data_folder,
        hour,
        solver_name,
        line_changes=None,
    )
    print_nodal_summary(base_nodal, info, title="BASE NODAL CASE")

    # 2) Nodal congestion sensitivity.
    sensitivity_df, nodal_cases = run_sensitivity_analysis(
        data_folder,
        hour,
        congestion_cases,
        solver_name,
    )
    print("\nNODAL CONGESTION CASES")
    print("-" * 60)
    print(sensitivity_df.to_string(index=False))

    stressed_nodal = nodal_cases["reduced_corridors"]
    print_nodal_summary(
        stressed_nodal,
        info,
        title="STRESSED NODAL CASE (L25=400 MW, L23=250 MW, L22=250 MW)",
    )

    print("\nNODAL PRICE FORMATION SUPPORT TABLE")
    print("-" * 60)
    print(build_nodal_price_formation_table(stressed_nodal).to_string(index=False))

    # 3) Base zonal case.
    base_zonal, _, _, _, zone_map, _, base_atc = run_zonal_case(
        data_folder,
        hour,
        solver_name,
        zone_map=None,
        line_changes=None,
        atc_factor=1.0,
        atc_changes=None,
    )
    print_zonal_summary(base_zonal, info, title="BASE ZONAL CASE")
    print("\nZONE MAPPING")
    print(zone_map)
    print("\nBASE INTERZONAL ATC TABLE")
    print(base_atc.to_string(index=False))

    # 4) Ex-post feasibility of the base zonal dispatch.
    base_zonal_gen_dispatch = base_zonal["generation"].set_index("gen_id")["dispatch"]
    base_zonal_served_demand = base_zonal["demand"].set_index("load_id")["served_demand"]

    base_feasibility_results = check_dispatch_feasibility_on_full_network(
        generation_dispatch=base_zonal_gen_dispatch,
        served_demand=base_zonal_served_demand,
        generators=generators,
        demands=demands,
        buses=buses,
        lines=lines,
        solver_name=solver_name,
    )

    print("\nBASE NODAL VS BASE ZONAL")
    print("=" * 60)
    print_framework_comparison(base_nodal, base_zonal, base_feasibility_results)

    # 5) Zonal ATC sensitivity.
    zonal_sensitivity_df, zonal_cases, _ = run_zonal_atc_sensitivity(
        data_folder,
        hour,
        zone_map=None,
        line_changes=None,
        atc_factors=zonal_atc_factors,
        target_link_id=None,
        solver_name=solver_name,
    )
    print("\nZONAL ATC SENSITIVITY RESULTS")
    print("-" * 60)
    print(zonal_sensitivity_df.to_string(index=False))

    selected_zonal = zonal_cases[selected_atc_factor]
    print_zonal_summary(
        selected_zonal,
        info,
        title=f"SELECTED ZONAL ATC CASE (all ATCs at {selected_atc_factor:.2f} of base)",
    )

    # 6) Ex-post feasibility of the selected zonal dispatch.
    selected_zonal_gen_dispatch = selected_zonal["generation"].set_index("gen_id")["dispatch"]
    selected_zonal_served_demand = selected_zonal["demand"].set_index("load_id")["served_demand"]

    selected_feasibility_results = check_dispatch_feasibility_on_full_network(
        generation_dispatch=selected_zonal_gen_dispatch,
        served_demand=selected_zonal_served_demand,
        generators=generators,
        demands=demands,
        buses=buses,
        lines=lines,
        solver_name=solver_name,
    )

    print("\nBASE NODAL VS SELECTED ZONAL ATC CASE")
    print("=" * 60)
    print_framework_comparison(base_nodal, selected_zonal, selected_feasibility_results)

    # 7) Plot
    plot_base_vs_stressed_nodal_prices(
         base_nodal,
         stressed_nodal,
         save_path="plots/base_vs_stressed_nodal_prices.pdf"
    )


if __name__ == "__main__":
    main()