import os
from typing import Any

import pandas as pd

from app.algorithms.leakage_identifier import LeakageIdentifier
from app.services.tjnetwork import (
    PARTITION_TYPE_KWAY,
    calculate_district_metering_area_for_nodes,
    dump_inp,
    get_all_scada_info,
    get_network_node_coords,
)


def run_leakage_identification(
    network: str,
    observed_pressure_data: str | pd.DataFrame | dict[str, list[Any]] | list[dict[str, Any]],
    start_time: float = 0,
    duration: float = 24,
    timestep: float = 5,
    q_sum: float = 0.2,
    q_sum_unit: str = "m3/s",
    output_dir: str = "Results",
    pop_size: int = 50,
    max_gen: int = 100,
    output_flow_unit: str = "m3/s",
) -> dict[str, Any]:
    os.makedirs(output_dir, exist_ok=True)
    inp_path = os.path.join(output_dir, f"{network}.leakage.inp")
    dump_inp(network, inp_path, "2")
    sensor_nodes = _get_pressure_sensor_nodes(network)
    area_map = _build_area_map_by_spectral_partition(network, sensor_nodes)

    q_sum_m3s = LeakageIdentifier._flow_to_m3s(q_sum, q_sum_unit)
    identifier = LeakageIdentifier(
        inp_path=inp_path,
        sensor_nodes=sensor_nodes,
        area_map=area_map,
        start_time=start_time,
        duration=duration,
        timestep=timestep,
        q_sum=q_sum_m3s,
    )
    result_df = identifier.run_identification(
        observed_pressure_data=observed_pressure_data,
        output_dir=output_dir,
        pop_size=pop_size,
        max_gen=max_gen,
        output_flow_unit=output_flow_unit,
        save_result=False,
    )
    return {
        "result_path": result_df.attrs.get("result_path"),
        "sensor_nodes": sensor_nodes,
        "area_count": len(set(area_map.values())),
        "rows": result_df.to_dict(orient="records"),
    }


def _get_pressure_sensor_nodes(network: str) -> list[str]:
    scada_info = get_all_scada_info(network)
    sensor_nodes: list[str] = []
    for item in scada_info:
        if item.get("type") != "pressure":
            continue
        node_id = item.get("associated_element_id")
        if isinstance(node_id, str) and node_id:
            sensor_nodes.append(node_id)
    sensor_nodes = list(dict.fromkeys(sensor_nodes))
    if not sensor_nodes:
        raise ValueError("未找到压力传感器对应节点（scada_info.type=pressure）。")
    return sensor_nodes


def _build_area_map_by_spectral_partition(
    network: str, sensor_nodes: list[str]
) -> dict[str, str]:
    node_coords = get_network_node_coords(network)
    all_nodes = list(node_coords.keys())
    if not all_nodes:
        raise ValueError("管网中未获取到可分区节点。")

    part_count = min(len(sensor_nodes), len(all_nodes))
    if part_count <= 0:
        raise ValueError("无可用压力传感器，无法生成虚拟分区。")

    groups = calculate_district_metering_area_for_nodes(
        network,
        all_nodes,
        part_count=part_count,
        part_type=PARTITION_TYPE_KWAY,
    )
    if not groups:
        raise ValueError("虚拟分区计算失败，未返回分区结果。")

    area_map: dict[str, str] = {}
    for idx, group_nodes in enumerate(groups, start=1):
        area_id = str(idx)
        for node_id in group_nodes:
            area_map[node_id] = area_id
    if not area_map:
        raise ValueError("虚拟分区结果为空，无法生成节点区域映射。")
    return area_map