"""
Data validation functions for power flow data integrity and consistency checks.

This module contains all validation functions extracted from test_data_validation.py
to provide comprehensive validation of generated power flow data.
"""

import pandas as pd
import numpy as np
from typing import Dict, Iterable
from gridfm_datakit.utils.column_names import (
    BUS_COLUMNS,
    DC_BUS_COLUMNS,
    BRANCH_COLUMNS,
    DC_BRANCH_COLUMNS,
    GEN_COLUMNS,
    DC_GEN_COLUMNS,
    YBUS_COLUMNS,
    RUNTIME_COLUMNS,
    DC_RUNTIME_COLUMNS,
)
from gridfm_datakit.utils.power_balance import (
    compute_branch_powers_vectorized,
    compute_bus_balance,
)


def validate_generated_data(
    file_paths: Dict[str, str],
    mode: str,
    sn_mva: float,
    n_scenarios: int = 0,
) -> bool:
    """Run all validation tests on the generated data.

    Args:
        file_paths: Dictionary containing paths to data files (bus_data, branch_data, gen_data, y_bus_data).
        mode: Operating mode ("opf" or "pf").
        n_scenarios: Number of scenarios to sample for validation (0 for all scenarios).
        sn_mva: Base MVA used to scale power quantities

    Returns:
        True if all validations pass.

    Raises:
        AssertionError: If any validation fails.
    """
    # Load the generated data
    bus_data = pd.read_parquet(file_paths["bus_data"], engine="fastparquet")
    branch_data = pd.read_parquet(file_paths["branch_data"], engine="fastparquet")
    gen_data = pd.read_parquet(file_paths["gen_data"], engine="fastparquet")
    y_bus_data = pd.read_parquet(file_paths["y_bus_data"], engine="fastparquet")
    runtime_data = pd.read_parquet(file_paths["runtime_data"], engine="fastparquet")

    generated_data = {
        "bus_data": bus_data,
        "branch_data": branch_data,
        "gen_data": gen_data,
        "y_bus_data": y_bus_data,
        "runtime_data": runtime_data,
        "mode": mode,
        "file_paths": file_paths,
    }

    try:
        validate_scenario_indexing_consistency(generated_data)
    except Exception as e:
        raise AssertionError(f"Scenario indexing consistency validation failed: {e}")

    # Run Data Integrity Tests
    try:
        validate_bus_indexing_consistency(generated_data)
    except Exception as e:
        raise AssertionError(f"Bus indexing consistency validation failed: {e}")

    try:
        validate_data_completeness(generated_data)
    except Exception as e:
        raise AssertionError(f"Data completeness validation failed: {e}")

        # DC columns consistency: if any DC column has NaN, all DC columns must be NaN
    try:
        validate_dc_columns_consistency(generated_data)
    except Exception as e:
        raise AssertionError(f"DC columns consistency validation failed: {e}")

    # Sample scenarios if n_scenarios is provided to avoid too long validation times
    if n_scenarios > 0:
        max_scenarios = len(generated_data["bus_data"]["scenario"].unique())
        sampled_scenarios = np.random.choice(
            generated_data["bus_data"]["scenario"].unique(),
            size=min(n_scenarios, max_scenarios),
            replace=False,
        )
        generated_data["bus_data"] = generated_data["bus_data"][
            generated_data["bus_data"]["scenario"].isin(sampled_scenarios)
        ]
        generated_data["branch_data"] = generated_data["branch_data"][
            generated_data["branch_data"]["scenario"].isin(sampled_scenarios)
        ]
        generated_data["gen_data"] = generated_data["gen_data"][
            generated_data["gen_data"]["scenario"].isin(sampled_scenarios)
        ]
        generated_data["y_bus_data"] = generated_data["y_bus_data"][
            generated_data["y_bus_data"]["scenario"].isin(sampled_scenarios)
        ]
        generated_data["runtime_data"] = generated_data["runtime_data"][
            generated_data["runtime_data"]["scenario"].isin(sampled_scenarios)
        ]
        print(
            f"Sampled {len(sampled_scenarios)} scenarios for validation out of {max_scenarios}",
        )

    # Run Y-Bus Consistency Tests
    try:
        validate_ybus_diagonal_consistency(generated_data)
    except Exception as e:
        raise AssertionError(f"Y-bus diagonal consistency validation failed: {e}")

    # Run Branch Constraint Tests
    try:
        validate_deactivated_lines_zero_admittance(generated_data)
    except Exception as e:
        raise AssertionError(
            f"Deactivated lines zero admittance validation failed: {e}",
        )

    try:
        validate_computed_vs_stored_power_flows(generated_data, sn_mva)
    except Exception as e:
        raise AssertionError(f"Computed vs stored power flows validation failed: {e}")

    try:
        validate_tap_not_zero(generated_data)
    except Exception as e:
        raise AssertionError(f"Tap not zero validation failed: {e}")

    # Run branch loading validation for both OPF and PF modes
    # In OPF mode: asserts loading <= 1.01
    # In PF mode: computes statistics without asserting
    try:
        validate_branch_loading_opf_mode(generated_data)
    except Exception as e:
        if mode == "opf":
            raise AssertionError(f"Branch loading OPF mode validation failed: {e}")
        else:
            print(f"    Branch loading computation encountered errors: {e}")

    # Run Generator Constraint Tests
    try:
        validate_deactivated_generators_zero_output(generated_data)
    except Exception as e:
        raise AssertionError(
            f"Deactivated generators zero output validation failed: {e}",
        )

    try:
        validate_generator_limits(generated_data)
    except Exception as e:
        raise AssertionError(f"Generator limits validation failed: {e}")

    # Run OPF mode Constraints
    if mode == "opf":
        try:
            validate_voltage_magnitude_limits_opf_mode(generated_data)
        except Exception as e:
            raise AssertionError(
                f"Voltage magnitude limits OPF mode validation failed: {e}",
            )
        try:
            validate_branch_angle_difference_opf_mode(generated_data)
        except Exception as e:
            raise AssertionError(
                f"Branch angle difference limits OPF mode validation failed: {e}",
            )

    # Run Power Balance Tests
    try:
        validate_bus_generation_consistency(generated_data)
    except Exception as e:
        raise AssertionError(f"Bus generation consistency validation failed: {e}")

    # DC bus generation consistency (if DC fields present)
    try:
        validate_bus_generation_consistency_dc(generated_data)
    except Exception as e:
        raise AssertionError(f"Bus generation DC consistency validation failed: {e}")

    # Check Pg and Pg_dc match at slack nodes in PF mode
    if mode == "pf":
        try:
            validate_non_slack_pg_consistency(generated_data)
        except Exception as e:
            raise AssertionError(f"Slack Pg consistency validation failed: {e}")

    try:
        validate_power_balance_equations(generated_data, sn_mva)
    except Exception as e:
        raise AssertionError(f"Power balance equations validation failed: {e}")

    return True


def validate_ybus_diagonal_consistency(generated_data: Dict[str, pd.DataFrame]) -> None:
    """Test Y-bus diagonal consistency with bus and branch data (vectorized)."""
    bus_data = generated_data["bus_data"]
    branch_data = generated_data["branch_data"]
    y_bus_data = generated_data["y_bus_data"]

    scenarios = bus_data["scenario"].unique()
    total_buses = len(bus_data)
    print(
        f"    Y-bus diagonal consistency: validating {total_buses} bus entries across {len(scenarios)} scenarios",
    )

    # Aggregate Yff contributions by (scenario, from_bus)
    yff_sum = (
        branch_data.groupby(["scenario", "from_bus"], as_index=False)
        .agg({"Yff_r": "sum", "Yff_i": "sum"})
        .rename(columns={"from_bus": "bus", "Yff_r": "yff_sum_g", "Yff_i": "yff_sum_b"})
    )

    # Aggregate Ytt contributions by (scenario, to_bus)
    ytt_sum = (
        branch_data.groupby(["scenario", "to_bus"], as_index=False)
        .agg({"Ytt_r": "sum", "Ytt_i": "sum"})
        .rename(columns={"to_bus": "bus", "Ytt_r": "ytt_sum_g", "Ytt_i": "ytt_sum_b"})
    )

    # Prepare bus data with (scenario, bus) as key
    bus_keyed = bus_data[["scenario", "bus", "GS", "BS"]].copy()
    bus_keyed["scenario"] = bus_keyed["scenario"].astype(int)
    bus_keyed["bus"] = bus_keyed["bus"].astype(int)

    # Merge all contributions
    expected = (
        bus_keyed.merge(yff_sum, on=["scenario", "bus"], how="left")
        .merge(ytt_sum, on=["scenario", "bus"], how="left")
        .fillna(
            {"yff_sum_g": 0.0, "yff_sum_b": 0.0, "ytt_sum_g": 0.0, "ytt_sum_b": 0.0},
        )
    )

    # Compute expected G and B
    expected["expected_g"] = (
        expected["GS"] + expected["yff_sum_g"] + expected["ytt_sum_g"]
    )
    expected["expected_b"] = (
        expected["BS"] + expected["yff_sum_b"] + expected["ytt_sum_b"]
    )

    # Get actual G and B from y_bus_data (diagonal entries only)
    ybus_diagonal = y_bus_data[(y_bus_data["index1"] == y_bus_data["index2"])][
        ["scenario", "index1", "G", "B"]
    ].rename(columns={"index1": "bus"})
    ybus_diagonal["scenario"] = ybus_diagonal["scenario"].astype(int)
    ybus_diagonal["bus"] = ybus_diagonal["bus"].astype(int)

    # Merge expected with actual
    comparison = expected.merge(
        ybus_diagonal,
        on=["scenario", "bus"],
        how="left",
        suffixes=("", "_actual"),
    )

    # Vectorized comparison
    g_diff = np.abs(comparison["expected_g"] - comparison["G"])
    b_diff = np.abs(comparison["expected_b"] - comparison["B"])

    tolerance = 1e-6
    g_mismatches = comparison[g_diff >= tolerance]
    b_mismatches = comparison[b_diff >= tolerance]

    if len(g_mismatches) > 0:
        raise AssertionError(f"G mismatches: {g_mismatches}")
    if len(b_mismatches) > 0:
        raise AssertionError(f"B mismatches: {b_mismatches}")


def validate_deactivated_lines_zero_admittance(
    generated_data: Dict[str, pd.DataFrame],
) -> None:
    """Test that deactivated lines have zero power flows and admittances."""
    branch_data = generated_data["branch_data"]
    deactivated_branches = branch_data[branch_data["br_status"] == 0]

    print(
        f"    Deactivated lines zero admittance: validating {len(deactivated_branches)} deactivated branches",
    )
    if not deactivated_branches.empty:
        assert (deactivated_branches["pf"] == 0).all(), (
            "Deactivated branches should have zero pf"
        )
        assert (deactivated_branches["qf"] == 0).all(), (
            "Deactivated branches should have zero qf"
        )
        assert (deactivated_branches["pt"] == 0).all(), (
            "Deactivated branches should have zero pt"
        )
        assert (deactivated_branches["qt"] == 0).all(), (
            "Deactivated branches should have zero qt"
        )
        assert (deactivated_branches["Yff_r"] == 0).all(), (
            "Deactivated branches should have zero Yff_r"
        )
        assert (deactivated_branches["Yff_i"] == 0).all(), (
            "Deactivated branches should have zero Yff_i"
        )
        assert (deactivated_branches["Yft_r"] == 0).all(), (
            "Deactivated branches should have zero Yft_r"
        )
        assert (deactivated_branches["Yft_i"] == 0).all(), (
            "Deactivated branches should have zero Yft_i"
        )
        assert (deactivated_branches["Ytf_r"] == 0).all(), (
            "Deactivated branches should have zero Ytf_r"
        )
        assert (deactivated_branches["Ytf_i"] == 0).all(), (
            "Deactivated branches should have zero Ytf_i"
        )
        assert (deactivated_branches["Ytt_r"] == 0).all(), (
            "Deactivated branches should have zero Ytt_r"
        )
        assert (deactivated_branches["Ytt_i"] == 0).all(), (
            "Deactivated branches should have zero Ytt_i"
        )

    print("    Deactivated lines zero admittance: OK")


def validate_computed_vs_stored_power_flows(
    generated_data: Dict[str, pd.DataFrame],
    sn_mva: float,
) -> None:
    """Test that computed power flows match stored power flows."""

    print(
        f"    Validate computed vs stored power flows: validating {len(generated_data['branch_data'])} branches across {len(generated_data['branch_data']['scenario'].unique())} scenarios",
    )

    pf, qf, pt, qt = compute_branch_powers_vectorized(
        generated_data["branch_data"],
        generated_data["bus_data"],
        dc=False,
        sn_mva=sn_mva,
    )
    computed_flows = pd.DataFrame(
        {
            "pf": pf,
            "qf": qf,
            "pt": pt,
            "qt": qt,
            "scenario": generated_data["branch_data"]["scenario"],
            "from_bus": generated_data["branch_data"]["from_bus"],
            "to_bus": generated_data["branch_data"]["to_bus"],
        },
        index=generated_data["branch_data"].index,
    )

    flows_data = generated_data["branch_data"][
        ["pf", "qf", "pt", "qt", "scenario", "from_bus", "to_bus"]
    ]
    mismatch = ~np.isclose(computed_flows, flows_data, atol=1e-4, rtol=1e-4)
    # TODO investigate why atol has to be so large
    if mismatch.any():
        raise AssertionError(
            f"Computed power flows do not match stored power flows, stored: \n{flows_data[mismatch]}, computed: \n{computed_flows[mismatch]}",
        )

    print("    Computed vs stored power flows: OK")


def validate_tap_not_zero(generated_data: Dict[str, pd.DataFrame]) -> None:
    """Test that transformer tap ratio is not zero for active branches."""
    branch_data = generated_data["branch_data"]
    active_branches = branch_data[branch_data["br_status"] == 1]

    print(
        f"    Tap not zero: validating {len(active_branches)} active branches",
    )
    if not active_branches.empty:
        zero_tap_branches = active_branches[active_branches["tap"] == 0]
        assert len(zero_tap_branches) == 0, (
            f"Active branches should not have zero tap ratio. "
            f"Found {len(zero_tap_branches)} branches with tap=0"
        )

    print("    Tap not zero: OK")


def validate_branch_loading_opf_mode(generated_data: Dict[str, pd.DataFrame]) -> None:
    """Test branch loading limits in OPF mode, compute loading statistics in PF mode."""
    bus_data = generated_data["bus_data"]
    branch_data = generated_data["branch_data"]

    scenarios = bus_data["scenario"].unique()
    rated_branches = branch_data[
        (branch_data["br_status"] == 1) & (branch_data["rate_a"] > 0)
    ]

    mode_label = "opf" if generated_data["mode"] == "opf" else "pf"
    print(
        f"    Branch loading limits ({mode_label} mode): validating {len(rated_branches)} rated branches across {len(scenarios)} scenarios",
    )

    # Track binding constraints and overloads
    binding_loadings = []
    overloaded_branches = []

    for scenario in scenarios:
        branch_scenario = branch_data[branch_data["scenario"] == scenario]
        active_branches = branch_scenario[
            (branch_scenario["br_status"] == 1) & (branch_scenario["rate_a"] > 0)
        ]

        if active_branches.empty:
            continue

        for _, branch in active_branches.iterrows():
            from_bus = int(branch["from_bus"])
            to_bus = int(branch["to_bus"])
            rate_a = branch["rate_a"]

            # Get power flow values from branch data
            pf = branch["pf"]  # Real power from
            qf = branch["qf"]  # Reactive power from
            pt = branch["pt"]  # Real power to
            qt = branch["qt"]  # Reactive power to

            # Compute apparent power squared: S^2 = P^2 + Q^2
            s_from_sq = pf**2 + qf**2
            s_to_sq = pt**2 + qt**2

            # Loading is the ratio of apparent power to rate_a
            # PowerModels constraint: p[f_idx]^2 + q[f_idx]^2 <= rate_a^2
            loading_f = np.sqrt(s_from_sq) / rate_a if rate_a > 0 else 0
            loading_t = np.sqrt(s_to_sq) / rate_a if rate_a > 0 else 0
            loading = max(loading_f, loading_t)

            # Track if loading is binding (within 1% of limit)
            if loading >= 0.99:
                binding_loadings.append(loading)

            # Track overloads (loading > 1.0)
            if loading > 1.01:
                overloaded_branches.append((scenario, from_bus, to_bus, loading))

            # Only assert loading limits in OPF mode
            if generated_data["mode"] == "opf":
                assert loading <= 1.01, (
                    f"Scenario {scenario}, Branch {from_bus}->{to_bus}: "
                    f"Loading {loading:.3f} exceeds 1.01 in OPF mode"
                )

    print(
        f"    Binding loading constraints (>= 0.99): {len(binding_loadings)} branches",
    )
    if generated_data["mode"] == "pf":
        print(f"    Overloaded branches (> 1.0): {len(overloaded_branches)} branches")
        print("    Branch loading limits (PF mode): statistics computed")
    else:
        print("    Branch loading limits (OPF mode): OK")


def validate_deactivated_generators_zero_output(
    generated_data: Dict[str, pd.DataFrame],
) -> None:
    """Test that deactivated generators have zero output."""
    gen_data = generated_data["gen_data"]
    deactivated_gens = gen_data[gen_data["in_service"] == 0]

    print(
        f"    Deactivated generators zero output: validating {len(deactivated_gens)} deactivated generators",
    )
    if not deactivated_gens.empty:
        assert (deactivated_gens["p_mw"] == 0).all(), (
            "Deactivated generators should have zero p_mw"
        )
        assert (deactivated_gens["q_mvar"] == 0).all(), (
            "Deactivated generators should have zero q_mvar"
        )

        if "p_mw_dc" in deactivated_gens.columns:
            # zero or nan (if no solution was found)
            assert (
                (deactivated_gens["p_mw_dc"] == 0)
                | (deactivated_gens["p_mw_dc"].isna())
            ).all(), "Deactivated generators should have zero p_mw_dc or be NaN"

    print("    Deactivated generators zero output: OK")


def validate_generator_limits(generated_data: Dict[str, pd.DataFrame]) -> None:
    """Test that generator outputs respect their limits."""
    gen_data = generated_data["gen_data"]
    gen_data = gen_data[gen_data["in_service"] == 1]
    # keep only the ones with limits for p_mw
    filtered_gens = gen_data[
        gen_data["max_p_mw"].notna() & gen_data["min_p_mw"].notna()
    ]

    if generated_data["mode"] == "pf":
        filtered_gens = filtered_gens[filtered_gens["is_slack_gen"] == 0]

    print(
        f"    Generator limits: validating {len(filtered_gens)} active generators (mode: {generated_data['mode']})",
    )

    # Count binding P limits
    binding_p_min = 0
    binding_p_max = 0
    if not filtered_gens.empty:
        p_within_limits = (
            filtered_gens["p_mw"] >= filtered_gens["min_p_mw"] - 1e-2
        ) & (filtered_gens["p_mw"] <= filtered_gens["max_p_mw"] + 1e-2)

        # Check for binding minimum limits
        p_at_min = (filtered_gens["p_mw"] <= filtered_gens["min_p_mw"] + 1e-2) & (
            filtered_gens["p_mw"] >= filtered_gens["min_p_mw"] - 1e-2
        )
        binding_p_min = p_at_min.sum()

        # Check for binding maximum limits
        p_at_max = (filtered_gens["p_mw"] <= filtered_gens["max_p_mw"] + 1e-2) & (
            filtered_gens["p_mw"] >= filtered_gens["max_p_mw"] - 1e-2
        )
        binding_p_max = p_at_max.sum()

        assert p_within_limits.all(), (
            f"Generator active power should be within limits, current: \n{filtered_gens.loc[~p_within_limits, ['bus', 'p_mw']]}, \nmax: \n{filtered_gens.loc[~p_within_limits, ['bus', 'max_p_mw']]}"
        )

    # Count binding Q limits (only in OPF mode)
    binding_q_min = 0
    binding_q_max = 0
    if generated_data["mode"] == "opf":
        filtered_gens_q = filtered_gens[
            filtered_gens["max_q_mvar"].notna() & filtered_gens["min_q_mvar"].notna()
        ]
        q_within_limits = (
            filtered_gens_q["q_mvar"] >= filtered_gens_q["min_q_mvar"] - 1e-2
        ) & (filtered_gens_q["q_mvar"] <= filtered_gens_q["max_q_mvar"] + 1e-2)

        # Check for binding minimum limits
        q_at_min = (
            filtered_gens_q["q_mvar"] <= filtered_gens_q["min_q_mvar"] + 1e-2
        ) & (filtered_gens_q["q_mvar"] >= filtered_gens_q["min_q_mvar"] - 1e-2)
        binding_q_min = q_at_min.sum()

        # Check for binding maximum limits
        q_at_max = (
            filtered_gens_q["q_mvar"] <= filtered_gens_q["max_q_mvar"] + 1e-2
        ) & (filtered_gens_q["q_mvar"] >= filtered_gens_q["max_q_mvar"] - 1e-2)
        binding_q_max = q_at_max.sum()

        assert q_within_limits.all(), (
            f"Generator reactive power should be within limits, expected: {filtered_gens_q.loc[~q_within_limits, ['bus', 'q_mvar']]}, actual: {filtered_gens_q.loc[~q_within_limits, ['bus', 'q_mvar']]}, max: {filtered_gens_q.loc[~q_within_limits, ['bus', 'max_q_mvar']]}"
        )

    print(
        f"    Binding P limits: {binding_p_min} at minimum, {binding_p_max} at maximum",
    )
    if generated_data["mode"] == "opf":
        print(
            f"    Binding Q limits: {binding_q_min} at minimum, {binding_q_max} at maximum",
        )
    print("    Generator limits: OK")


def validate_voltage_magnitude_limits_opf_mode(
    generated_data: Dict[str, pd.DataFrame],
) -> None:
    """Test voltage magnitude limits in OPF mode."""
    if generated_data["mode"] != "opf":
        print("    Voltage magnitude limits: skipped (not in OPF mode)")
        return

    bus_data = generated_data["bus_data"]
    print(
        f"    Voltage magnitude limits (OPF mode): validating {len(bus_data)} bus voltage entries",
    )
    vm_within_limits = (bus_data["Vm"] >= bus_data["min_vm_pu"] - 1e-6) & (
        bus_data["Vm"] <= bus_data["max_vm_pu"] + 1e-6
    )
    assert vm_within_limits.all(), "Voltage magnitudes should be within limits"
    print("    Voltage magnitude limits (OPF mode): OK")


def validate_branch_angle_difference_opf_mode(
    generated_data: Dict[str, pd.DataFrame],
) -> None:
    """Validate branch angle difference limits in OPF mode.

    For each active branch, the difference in bus voltage angles must respect
    the branch angle limits [angmin, angmax].
    """
    if generated_data["mode"] != "opf":
        print("    Branch angle difference limits: skipped (not in OPF mode)")
        return

    bus_data = generated_data["bus_data"]
    branch_data = generated_data["branch_data"]
    scenarios = bus_data["scenario"].unique()
    print(
        f"    Branch angle difference limits (OPF mode): validating across {len(scenarios)} scenarios",
    )

    for scenario in scenarios:
        bus_scenario = bus_data[bus_data["scenario"] == scenario]
        branch_scenario = branch_data[branch_data["scenario"] == scenario]
        active_branches = branch_scenario[branch_scenario["br_status"] == 1]

        if active_branches.empty:
            continue

        va_deg = bus_scenario.set_index("bus")["Va"].to_dict()

        for _, br in active_branches.iterrows():
            fb = int(br["from_bus"])
            tb = int(br["to_bus"])
            angmin = br["ang_min"]
            angmax = br["ang_max"]

            if fb not in va_deg or tb not in va_deg:
                continue

            delta = va_deg[fb] - va_deg[tb]

            # Normalize delta into [-180, 180] for robust comparison
            while delta > 180.0:
                delta -= 360.0
            while delta < -180.0:
                delta += 360.0

            assert delta >= angmin - 1e-6, (
                f"Scenario {scenario}, Branch {fb}->{tb}: angle diff {delta:.3f} < angmin {angmin}"
            )
            assert delta <= angmax + 1e-6, (
                f"Scenario {scenario}, Branch {fb}->{tb}: angle diff {delta:.3f} > angmax {angmax}"
            )

    print("    Branch angle difference limits (OPF mode): OK")


def validate_bus_generation_consistency(
    generated_data: Dict[str, pd.DataFrame],
) -> None:
    """Test that Pg in bus data equals sum of generators at each bus (vectorized)."""
    bus_data = generated_data["bus_data"]
    gen_data = generated_data["gen_data"]

    scenarios = bus_data["scenario"].unique()
    print(
        f"    Bus generation consistency: validating {len(bus_data)} bus entries across {len(scenarios)} scenarios",
    )

    # Aggregate generator outputs by (scenario, bus)
    gen_sum = (
        gen_data.groupby(["scenario", "bus"], as_index=False)
        .agg({"p_mw": "sum", "q_mvar": "sum"})
        .rename(columns={"p_mw": "pg_gen_sum", "q_mvar": "qg_gen_sum"})
    )

    # Prepare bus data with (scenario, bus) as key
    bus_keyed = bus_data[["scenario", "bus", "Pg", "Qg"]].copy()
    bus_keyed["scenario"] = bus_keyed["scenario"].astype(int)
    bus_keyed["bus"] = bus_keyed["bus"].astype(int)

    # Merge bus data with generator sums
    comparison = bus_keyed.merge(gen_sum, on=["scenario", "bus"], how="left").fillna(
        {"pg_gen_sum": 0.0, "qg_gen_sum": 0.0},
    )

    # Vectorized comparison
    tolerance = 1e-6
    pg_diff = np.abs(comparison["Pg"] - comparison["pg_gen_sum"])
    qg_diff = np.abs(comparison["Qg"] - comparison["qg_gen_sum"])

    pg_mismatches = comparison[pg_diff >= tolerance]
    qg_mismatches = comparison[qg_diff >= tolerance]

    if len(pg_mismatches) > 0:
        raise AssertionError(f"Pg mismatches: {pg_mismatches}")
    if len(qg_mismatches) > 0:
        raise AssertionError(f"Qg mismatches: {qg_mismatches}")

    print("    Bus generation consistency: OK")


def validate_bus_generation_consistency_dc(
    generated_data: Dict[str, pd.DataFrame],
) -> None:
    """Test that Pg_dc in bus data equals sum of DC generator outputs at each bus (vectorized).

    Skips if required DC columns are not present.
    """
    bus_data = generated_data["bus_data"]
    gen_data = generated_data["gen_data"]

    # Require presence of DC columns
    if ("Pg_dc" not in bus_data.columns) or ("p_mw_dc" not in gen_data.columns):
        print("    Bus generation DC consistency: skipped (no DC columns)")
        return

    scenarios = bus_data["scenario"].unique()
    print(
        f"    Bus generation DC consistency: validating {len(bus_data)} bus entries across {len(scenarios)} scenarios",
    )

    # Aggregate DC generator outputs by (scenario, bus)
    gen_sum_dc = (
        gen_data.groupby(["scenario", "bus"], as_index=False)
        .agg({"p_mw_dc": "sum"})
        .rename(columns={"p_mw_dc": "pg_dc_gen_sum"})
    )

    # Prepare bus data with (scenario, bus) as key
    bus_keyed = bus_data[["scenario", "bus", "Pg_dc"]].copy()
    bus_keyed["scenario"] = bus_keyed["scenario"].astype(int)
    bus_keyed["bus"] = bus_keyed["bus"].astype(int)

    # Merge bus data with generator sums
    comparison = bus_keyed.merge(gen_sum_dc, on=["scenario", "bus"], how="left").fillna(
        {"pg_dc_gen_sum": 0.0},
    )

    # Filter out NaN Pg_dc rows (e.g., OPF mode without DC)
    comparison = comparison[comparison["Pg_dc"].notna()]

    if len(comparison) == 0:
        print("    Bus generation DC consistency: OK")
        return

    # Vectorized comparison
    tolerance = 1e-6
    pg_dc_diff = np.abs(comparison["Pg_dc"] - comparison["pg_dc_gen_sum"])

    pg_dc_mismatches = comparison[pg_dc_diff >= tolerance]

    if len(pg_dc_mismatches) > 0:
        raise AssertionError(f"Pg_dc mismatches: {pg_dc_mismatches}")

    print("    Bus generation DC consistency: OK")


def validate_dc_columns_consistency(generated_data: Dict[str, pd.DataFrame]) -> None:
    """Test that if any DC column has NaN for a scenario, all DC columns are NaN for that scenario (vectorized).

    This ensures that when DC power flow doesn't converge, all DC data is consistently NaN,
    not just some columns.
    """
    bus_data = generated_data["bus_data"]
    branch_data = generated_data["branch_data"]
    gen_data = generated_data["gen_data"]
    runtime_data = generated_data["runtime_data"]

    # check if there is dc data:
    if "Va_dc" not in bus_data.columns:
        print("    DC columns consistency: skipped (no Va_dc column)")
        return

    # compute scenarios that have nan for each DC column
    scenarios_with_nan = {}
    for col in DC_BUS_COLUMNS:
        scenarios_with_nan[col] = set(
            bus_data[bus_data[col].isna()]["scenario"].unique(),
        )
    for col in DC_BRANCH_COLUMNS:
        scenarios_with_nan[col] = set(
            branch_data[branch_data[col].isna()]["scenario"].unique(),
        )
    for col in DC_GEN_COLUMNS:
        scenarios_with_nan[col] = set(
            gen_data[gen_data[col].isna()]["scenario"].unique(),
        )
    for col in DC_RUNTIME_COLUMNS:
        scenarios_with_nan[col] = set(
            runtime_data[runtime_data[col].isna()]["scenario"].unique(),
        )

    # check all values of scenarios_with_nan are the same
    key_0 = list(scenarios_with_nan.keys())[0]
    for col in scenarios_with_nan:
        if scenarios_with_nan[col] != scenarios_with_nan[key_0]:
            # find the difference
            difference = (
                scenarios_with_nan[col]
                .difference(scenarios_with_nan[key_0])
                .union(scenarios_with_nan[key_0].difference(scenarios_with_nan[col]))
            )
            raise AssertionError(
                f"DC columns consistency: scenarios with nan for {col} are not the same, difference: {sorted([int(i) for i in difference])}",
            )

    print(
        "number of scenarios with nan: ",
        len(
            scenarios_with_nan[key_0].union(
                scenarios_with_nan[key_0].difference(scenarios_with_nan[col]),
            ),
        ),
    )
    print("    DC columns consistency: OK")


def validate_non_slack_pg_consistency(generated_data: Dict[str, pd.DataFrame]) -> None:
    """Test that Pg and Pg_dc are the same for all buses except slack in PF mode.

    In power flow mode, the dispatch (generator outputs) should not change between
    AC and DC power flow solutions. The slack generator may differ because it adjusts
    to balance the system, but all other buses should have Pg == Pg_dc.

    Skips if required DC columns are not present or if not in PF mode.
    """
    if generated_data["mode"] != "pf":
        print("    Slack Pg consistency: skipped (not in PF mode)")
        return

    bus_data = generated_data["bus_data"]

    # Require presence of DC columns
    if "Pg_dc" not in bus_data.columns:
        print("    Slack Pg consistency: skipped (no Pg_dc column)")
        return

    scenarios = bus_data["scenario"].unique()
    non_slack_buses = bus_data[bus_data["REF"] == 0]
    print(
        f"    Non-slack Pg consistency: validating {len(non_slack_buses)} non-slack bus entries across {len(scenarios)} scenarios",
    )

    for scenario in scenarios:
        bus_scenario = bus_data[bus_data["scenario"] == scenario]
        non_slack_buses_scenario = bus_scenario[bus_scenario["REF"] == 0]

        for _, bus_row in non_slack_buses_scenario.iterrows():
            bus_idx = int(bus_row["bus"])
            pg = bus_row["Pg"]
            pg_dc = bus_row["Pg_dc"]

            # Allow NaN Pg_dc rows to pass silently (e.g., if DC solution not available)
            if pd.isna(pg_dc):
                continue

            tolerance = 1e-6
            assert abs(pg - pg_dc) < tolerance, (
                f"Scenario {scenario}, Bus {bus_idx}: Pg and Pg_dc mismatch (dispatch should not change in PF mode), "
                f"Pg: {pg}, Pg_dc: {pg_dc}"
            )

    print("    Non-slack Pg consistency: OK")


def validate_power_balance_equations(
    generated_data: Dict[str, pd.DataFrame],
    sn_mva: float,
) -> None:
    """Test power balance equations (Kirchhoff's Current Law)."""
    bus_data = generated_data["bus_data"]
    branch_data = generated_data["branch_data"]

    scenarios = bus_data["scenario"].unique()
    print(
        f"    Power balance equations (Kirchhoff's Law): validating {len(bus_data)} bus entries across {len(scenarios)} scenarios",
    )

    power_balance_ac = compute_bus_balance(
        bus_data,
        branch_data,
        branch_data[["pf", "qf", "pt", "qt"]],
        False,
        sn_mva=sn_mva,
    )
    not_close_zero = ~np.isclose(0.0, power_balance_ac["P_mis_ac"], atol=1e-4)
    # TODO investigate why atol has to be so large
    if not_close_zero.any():
        raise AssertionError(
            f"Power balance equations (Kirchhoff's Law) do not hold, mismatches: {power_balance_ac[not_close_zero]}",
        )

    print("    Power balance equations (Kirchhoff's Law): OK")


def validate_scenario_indexing_consistency(
    generated_data: Dict[str, pd.DataFrame],
) -> None:
    """Test that scenario indices are consistent across all data files."""
    bus_data = generated_data["bus_data"]
    branch_data = generated_data["branch_data"]
    gen_data = generated_data["gen_data"]
    y_bus_data = generated_data["y_bus_data"]

    bus_scenarios = set(bus_data["scenario"].unique())
    branch_scenarios = set(branch_data["scenario"].unique())
    gen_scenarios = set(gen_data["scenario"].unique())
    ybus_scenarios = set(y_bus_data["scenario"].unique())

    print(
        f"    Scenario indexing consistency: validating {len(bus_scenarios)} scenarios across 4 data files",
    )

    assert bus_scenarios == branch_scenarios == gen_scenarios == ybus_scenarios, (
        "All data files should contain the same set of scenario indices"
    )

    print("    Scenario indexing consistency: OK")


def validate_bus_indexing_consistency(generated_data: Dict[str, pd.DataFrame]) -> None:
    """Test that bus indices are consistent across data files."""
    bus_data = generated_data["bus_data"]
    branch_data = generated_data["branch_data"]
    gen_data = generated_data["gen_data"]

    bus_indices = set(bus_data["bus"].unique())
    branch_bus_indices = set(branch_data["from_bus"].unique()) | set(
        branch_data["to_bus"].unique(),
    )
    gen_bus_indices = set(gen_data["bus"].unique())

    print(
        f"    Bus indexing consistency: validating {len(bus_indices)} buses across 3 data files",
    )

    assert gen_bus_indices.issubset(bus_indices), (
        "All generator buses should exist in bus data"
    )
    assert branch_bus_indices.issubset(bus_indices), (
        "All branch endpoint buses should exist in bus data"
    )

    print("    Bus indexing consistency: OK")


def _require_columns(df: pd.DataFrame, name: str, required: Iterable[str]) -> None:
    missing = set(required) - set(df.columns)
    assert not missing, f"{name}: missing required columns {sorted(missing)}"


def _check_no_nan(df: pd.DataFrame, name: str, required: Iterable[str]) -> None:
    if df[required].isna().any().any():
        for col in required:
            assert not df[col].isna().any(), (
                f"{name}: column '{col}' contains NaN values"
            )


def validate_data_completeness(generated_data: Dict[str, pd.DataFrame]) -> None:
    """Test that all required columns are present and contain no NaN values."""
    bus_data = generated_data["bus_data"]
    branch_data = generated_data["branch_data"]
    gen_data = generated_data["gen_data"]
    y_bus_data = generated_data["y_bus_data"]
    runtime_data = generated_data["runtime_data"]

    total_entries = len(bus_data) + len(branch_data) + len(gen_data) + len(y_bus_data)
    print(
        f"    Data completeness: validating {total_entries} total entries across 4 data files",
    )

    # 1) Ensure 'scenario' column exists everywhere
    for name, df in [
        ("Bus data", bus_data),
        ("Branch data", branch_data),
        ("Generator data", gen_data),
        ("Y-bus data", y_bus_data),
        ("Runtime data", runtime_data),
    ]:
        assert "scenario" in df.columns, f"{name} should have scenario column"

    dc = True if "Va_dc" in bus_data.columns else False

    # 2) Check required columns exist and contain no NaN values
    _require_columns(
        bus_data,
        "Bus data",
        BUS_COLUMNS + DC_BUS_COLUMNS if dc else BUS_COLUMNS,
    )
    _require_columns(
        branch_data,
        "Branch data",
        BRANCH_COLUMNS + DC_BRANCH_COLUMNS if dc else BRANCH_COLUMNS,
    )
    _require_columns(
        gen_data,
        "Generator data",
        GEN_COLUMNS + DC_GEN_COLUMNS if dc else GEN_COLUMNS,
    )
    _require_columns(y_bus_data, "Y-bus data", YBUS_COLUMNS)

    _check_no_nan(bus_data, "Bus data", BUS_COLUMNS)
    _check_no_nan(branch_data, "Branch data", BRANCH_COLUMNS)
    _check_no_nan(gen_data, "Generator data", GEN_COLUMNS)
    _check_no_nan(y_bus_data, "Y-bus data", YBUS_COLUMNS)
    _check_no_nan(runtime_data, "Runtime data", RUNTIME_COLUMNS)

    # 3) Non-emptiness
    assert len(bus_data) > 0, "Bus data should not be empty"
    assert len(branch_data) > 0, "Branch data should not be empty"
    assert len(gen_data) > 0, "Generator data should not be empty"
    assert len(y_bus_data) > 0, "Y-bus data should not be empty"

    print("    Data completeness: OK (all required columns present and NaN-free)")


if __name__ == "__main__":
    # debug with data_out/case24_ieee_rts/raw
    file_paths = {
        "bus_data": "data_out/case24_ieee_rts/raw/bus_data.parquet",
        "branch_data": "data_out/case24_ieee_rts/raw/branch_data.parquet",
        "gen_data": "data_out/case24_ieee_rts/raw/gen_data.parquet",
        "y_bus_data": "data_out/case24_ieee_rts/raw/y_bus_data.parquet",
        "runtime_data": "data_out/case24_ieee_rts/raw/runtime_data.parquet",
    }
    validate_generated_data(file_paths, "pf", 100.0, n_scenarios=10)
    print("Validation completed successfully!")
