"""Step 3 — 가중 그래프 구성, Step 4 — 하드 제약(강한 선호)을 가중치에 반영."""

from __future__ import annotations

import networkx as nx

from .config import Config
from .validate import BaselineResult


def build_weighted_graph(G: nx.Graph, baseline: BaselineResult, cfg: Config) -> nx.Graph:
    """엣지 가중치 = '이 관로를 자르기 싫은 정도'. 클수록 절단 회피."""
    Gw = G.copy()
    alpha = cfg.partition.flow_weight_alpha

    flows: dict[tuple, float] = {}
    for u, v, data in Gw.edges(data=True):
        pipe_ids = data.get("pipe_ids", [])
        flows[(u, v)] = max((baseline.link_flow_mean_abs.get(pid, 0.0) for pid in pipe_ids), default=0.0)

    max_flow = max(flows.values(), default=0.0)
    max_diameter = max((data.get("diameter", 0.0) for _, _, data in Gw.edges(data=True)), default=0.0)

    for u, v, data in Gw.edges(data=True):
        q_norm = (flows[(u, v)] / max_flow) if max_flow > 0 else 0.0
        d_norm = (data.get("diameter", 0.0) / max_diameter) if max_diameter > 0 else 0.0
        w = alpha * q_norm + (1 - alpha) * d_norm
        data["w"] = float(min(max(w, 1e-6), 1.0))

    return Gw


def apply_hard_constraints(Gw: nx.Graph, cfg: Config) -> nx.Graph:
    c = cfg.constraints

    # 1. 표고대 분리
    for u, v, data in Gw.edges(data=True):
        eu = Gw.nodes[u].get("elevation", 0.0)
        ev = Gw.nodes[v].get("elevation", 0.0)
        if abs(eu - ev) > c.elevation_band_m:
            data["w"] *= 0.1

    # 2. 대수용가 (base_demand는 SI[m3/s], 기준값은 L/s로 주어짐)
    threshold_m3s = c.large_consumer_demand_lps / 1000.0
    large_consumers = {
        n for n, ndata in Gw.nodes(data=True) if ndata.get("base_demand", 0.0) >= threshold_m3s
    }
    for n in Gw.nodes:
        Gw.nodes[n]["is_large_consumer"] = n in large_consumers

    factor = 0.1 if c.large_consumer_policy == "isolate" else 5.0
    for n in large_consumers:
        for _, _, data in Gw.edges(n, data=True):
            data["w"] *= factor

    # 3. 수원 분리: Reservoir/Tank 노드쌍 간 최단경로(연장 기준) 상 엣지
    if c.sources_in_separate_blocks:
        sources = [n for n, ndata in Gw.nodes(data=True) if ndata.get("node_type") in ("Reservoir", "Tank")]
        for i, s1 in enumerate(sources):
            for s2 in sources[i + 1 :]:
                if not nx.has_path(Gw, s1, s2):
                    continue
                path = nx.shortest_path(Gw, s1, s2, weight="length")
                for a, b in zip(path[:-1], path[1:]):
                    Gw[a][b]["w"] *= 0.5

    for _, _, data in Gw.edges(data=True):
        data["w"] = float(min(max(data["w"], 1e-6), 1.0))

    return Gw