"""Batch analysis jobs exposed as POST /api/jobs/{name}.

Each job takes the INP bytes and a params dict and returns a JSON-able dict.
Jobs run in the API process pool; raise BatchError for input problems
(reported as HTTP 400). Register new jobs in JOBS at the bottom.
"""
from __future__ import annotations

import numpy as np

import ctypes

from .batch import BatchError, Network, hhmm, parse_time
from .epanet import EN, EpanetException


def job_summary(inp: bytes, params: dict) -> dict:
    """Framework smoke job: min/max pressure per reporting period."""
    with Network(inp) as net:
        junctions = net.junction_indices()
        res = net.simulate(node_props=(EN.PRESSURE,), nodes=junctions)
        p = res.nodes[EN.PRESSURE]
        return {
            "times": res.times,
            "labels": [hhmm(t) for t in res.times],
            "minPressure": p.min(axis=1).round(4).tolist() if p.size else [],
            "maxPressure": p.max(axis=1).round(4).tolist() if p.size else [],
            "junctions": len(junctions),
            "warnings": [{"time": t, "code": c} for t, c in res.warnings],
        }


def _window_rows(times: list[int], start: int, end: int) -> np.ndarray:
    """Reporting periods inside [start, end); falls back to the one nearest start."""
    t = np.asarray(times)
    rows = np.where((t >= start) & (t < end))[0]
    if rows.size == 0:
        rows = np.array([int(np.argmin(np.abs(t - start)))])
    return rows


def job_fireflow(inp: bytes, params: dict) -> dict:
    """소화전 모의: extra demand at hydrant nodes for a time window.

    params: nodes [ids], flow (model flow units, per hydrant), start/end (h or h:mm),
    residual (model pressure units, optional: also find the available flow),
    threshold (optional: count junctions whose window-min pressure drops below it).
    """
    ids = [str(x) for x in params.get("nodes") or []]
    if not ids:
        raise BatchError("소화전 절점을 하나 이상 지정하세요.")
    flow = float(params.get("flow") or 0)
    if flow <= 0:
        raise BatchError("소화 유량은 0보다 커야 합니다.")
    with Network(inp) as net:
        start = parse_time(params.get("start"), 0)
        end = parse_time(params.get("end"), max(net.duration, net.report_step))
        if end <= start:
            end = start + net.report_step
        hyd = [net.node_index(i) for i in ids]
        junctions = net.junction_indices()
        base = net.simulate(nodes=junctions)
        rows = _window_rows(base.times, start, end)
        base_min = base.nodes[EN.PRESSURE][rows].min(axis=0)

        net.window_pattern("SPIPE_FIRE", start, end)
        cats = [net.add_demand(h, flow, "SPIPE_FIRE", "fire") for h in hyd]
        fire = net.simulate(nodes=junctions)
        fire_min = fire.nodes[EN.PRESSURE][rows].min(axis=0)
        for h, c in zip(hyd, cats):
            net.delete_demand(h, c)

        col = {j: k for k, j in enumerate(junctions)}
        # Compare the two runs period by period: tank/pump controls can shift
        # *when* the daily minimum happens, so "min before vs min after" may
        # even show a rise. Report each node at its largest drop instead.
        bw = base.nodes[EN.PRESSURE][rows]
        fw = fire.nodes[EN.PRESSURE][rows]
        drop = fw - bw
        worst = np.argmin(drop, axis=0)
        cols = np.arange(len(junctions))
        diff = drop[worst, cols]
        order = np.argsort(diff)
        threshold = params.get("threshold")
        out_nodes = [
            {
                "id": net.node_ids[junctions[k] - 1],
                "basePressure": round(float(bw[worst[k], k]), 3),
                "firePressure": round(float(fw[worst[k], k]), 3),
                "diff": round(float(diff[k]), 3),
                "time": hhmm(base.times[rows[worst[k]]]),
            }
            for k in order
            if abs(diff[k]) > 1e-4
        ]

        residual = params.get("residual")
        hydrants = []
        for nid, h in zip(ids, hyd):
            entry = {"id": nid}
            if h in col:
                k = col[h]
                # Hydrant: window minimum before/after (what the residual criterion is about).
                entry.update(basePressure=round(float(base_min[k]), 3), firePressure=round(float(fire_min[k]), 3))
            if residual not in (None, "") and h in col:
                entry["available"] = _available_fire_flow(net, h, col[h], junctions, rows, start, end, float(residual), flow)
            hydrants.append(entry)

        return {
            "window": [start, end],
            "windowLabel": f"{hhmm(start)}~{hhmm(end)}",
            "flow": flow,
            "hydrants": hydrants,
            "nodes": out_nodes,
            "belowThreshold": int(np.sum(fire_min < float(threshold))) if threshold not in (None, "") else None,
            "belowThresholdBefore": int(np.sum(base_min < float(threshold))) if threshold not in (None, "") else None,
            "warnings": [{"time": t, "code": c} for t, c in fire.warnings],
        }


def _available_fire_flow(net, h, col, junctions, rows, start, end, residual, hint) -> float:
    """Largest extra flow at node h keeping its window-min pressure >= residual (bisection)."""
    def min_p(q: float) -> float:
        net.window_pattern("SPIPE_AFF", start, end)
        c = net.add_demand(h, q, "SPIPE_AFF", "aff")
        try:
            r = net.simulate(nodes=[h])
        finally:
            net.delete_demand(h, c)
        return float(r.nodes[EN.PRESSURE][rows, 0].min())

    if min_p(0.0) < residual:
        return 0.0
    lo, hi = 0.0, max(hint, 1e-3)
    for _ in range(30):  # grow until the residual is violated
        if min_p(hi) < residual:
            break
        lo, hi = hi, hi * 2
    else:
        return round(hi, 3)
    for _ in range(25):
        mid = (lo + hi) / 2
        if min_p(mid) >= residual:
            lo = mid
        else:
            hi = mid
        if hi - lo <= max(1e-3, 1e-3 * hi):
            break
    return round(lo, 3)


def _window(net: Network, params: dict) -> tuple[int, int]:
    start = parse_time(params.get("start"), 0)
    end = parse_time(params.get("end"), max(net.duration, net.report_step))
    return start, (end if end > start else start + net.report_step)


MAX_CRITICALITY_LINKS = 4000


def job_criticality(inp: bytes, params: dict) -> dict:
    """관로 중요도: close each pipe in turn and measure who loses pressure.

    params: links [ids] (default all pipes), threshold (model pressure units),
    start/end (window, h). A junction is "affected" when its window-min pressure
    falls below the threshold while it was at/above it in the base run.
    """
    threshold = float(params.get("threshold") or 0)
    with Network(inp) as net:
        start, end = _window(net, params)
        ids = [str(x) for x in params.get("links") or []]
        if ids:
            links = [net.link_index(i) for i in ids]
        else:
            links = [k + 1 for k, t in enumerate(net.link_types) if t in (EN.CVPIPE, EN.PIPE)]
        if not links:
            raise BatchError("분석할 관로가 없습니다.")
        if len(links) > MAX_CRITICALITY_LINKS:
            raise BatchError(f"관로가 {len(links)}개입니다. {MAX_CRITICALITY_LINKS}개 이하로 선택해 나눠 실행하세요.")
        junctions = net.junction_indices()
        base = net.simulate(node_props=(EN.PRESSURE, EN.DEMAND), nodes=junctions)
        rows = _window_rows(base.times, start, end)
        bmin = base.nodes[EN.PRESSURE][rows].min(axis=0)
        bdem = base.nodes[EN.DEMAND][rows].mean(axis=0)
        ok_before = bmin >= threshold
        out = []
        for li in links:
            orig = net.get_link(li, EN.INITSTATUS)
            net.set_link(li, EN.INITSTATUS, 0)
            entry = {"id": net.link_ids[li - 1]}
            try:
                r = net.simulate(nodes=junctions)
                pmin = r.nodes[EN.PRESSURE][_window_rows(r.times, start, end)].min(axis=0)
                affected = np.where((pmin < threshold) & ok_before)[0]
                drop = bmin - pmin
                order = affected[np.argsort(-drop[affected])]
                entry.update(
                    affectedNodes=int(affected.size),
                    affectedDemand=round(float(np.sum(np.maximum(bdem[affected], 0))), 4),
                    minPressure=round(float(pmin.min()), 3),
                    maxDrop=round(float(drop.max()), 3),
                    nodes=[net.node_ids[junctions[k] - 1] for k in order[:200]],
                    warning=any(c in (3, 6) for _, c in r.warnings),
                    disconnected=bool(pmin.min() < -1000),
                )
            except (EpanetException, BatchError) as e:
                entry.update(error=str(e), affectedNodes=0, affectedDemand=0.0, minPressure=None, maxDrop=None, nodes=[])
            finally:
                net.set_link(li, EN.INITSTATUS, orig)
            out.append(entry)
        out.sort(key=lambda e: (-e["affectedDemand"], -e["affectedNodes"], -(e["maxDrop"] or 0)))
        return {
            "window": [start, end],
            "windowLabel": f"{hhmm(start)}~{hhmm(end)}",
            "threshold": threshold,
            "belowBefore": int(np.sum(~ok_before)),
            "totalDemand": round(float(np.sum(np.maximum(bdem, 0))), 4),
            "links": out,
        }


def job_servicearea(inp: bytes, params: dict) -> dict:
    """급수 구역: source tracing (EPANET TRACE) from each reservoir (tanks optional).

    Returns each node's average share (%) of water from each source over the
    analysis window (default: the last 24 h of the run; a run shorter than 24 h
    is extended to 72 h so traced water can reach the far ends).
    """
    with Network(inp) as net:
        notes = []
        if net.duration < 72 * 3600:
            # Traced water starts at 0 everywhere; give it time to reach the far ends.
            net.set_time_param(EN.DURATION, 72 * 3600)
            notes.append(f"해석 기간({net.duration // 3600}시간)을 72시간으로 늘려 추적하고 마지막 24시간을 평균했습니다.")
            net.duration = 72 * 3600
        # Default: reservoirs only. A traced tank counts water that merely passed
        # through it, so tank shares overlap the reservoir shares (sum > 100%).
        kinds = (EN.RESERVOIR, EN.TANK) if params.get("includeTanks") else (EN.RESERVOIR,)
        sources = [str(x) for x in params.get("sources") or []] or [
            nid for nid, t in zip(net.node_ids, net.node_types) if t in kinds
        ] or [nid for nid, t in zip(net.node_ids, net.node_types) if t == EN.TANK]
        if not sources:
            raise BatchError("배수지·탱크가 없습니다.")
        if len(sources) > 30:
            raise BatchError("수원이 30개를 넘습니다. 추적할 수원을 골라 주세요.")
        for sid in sources:
            net.node_index(sid)
        start = parse_time(params.get("start"), max(0, net.duration - 24 * 3600))
        end = parse_time(params.get("end"), net.duration + 1)
        net.lib.EN_solveH(net.ph)
        shares: dict[str, np.ndarray] = {}
        for sid in sources:
            net.lib.EN_setqualtype(net.ph, EN.QUAL_TRACE, b"", b"", sid.encode("utf-8"))
            times, q = net.run_quality()
            rows = _window_rows(times, start, end)
            shares[sid] = np.clip(q[rows].mean(axis=0), 0, 100)
        mat = np.vstack([shares[s] for s in sources])  # [source, node]
        dom = np.argmax(mat, axis=0)
        nodes = {}
        for k, nid in enumerate(net.node_ids):
            col = mat[:, k]
            if col.max() < 0.5:
                nodes[nid] = {"source": None, "share": 0.0, "shares": {}}
                continue
            nodes[nid] = {
                "source": sources[dom[k]],
                "share": round(float(col[dom[k]]), 1),
                "shares": {s: round(float(v), 1) for s, v in zip(sources, col) if v >= 0.5},
            }
        summary = []
        for sid in sources:
            mine = [nid for nid in net.node_ids if nodes[nid]["source"] == sid]
            summary.append({"source": sid, "nodes": len(mine), "mixed": sum(1 for nid in mine if nodes[nid]["share"] < 80)})
        return {
            "sources": sources,
            "window": [start, end],
            "windowLabel": f"{hhmm(start)}~{hhmm(end)}",
            "nodes": nodes,
            "summary": summary,
            "unreached": sum(1 for v in nodes.values() if v["source"] is None),
            "notes": notes,
        }


def job_vsp(inp: bytes, params: dict) -> dict:
    """VSP: pump speed giving a target pressure at a control node (bisection).

    Each requested hour is solved as a snapshot (DURATION 0, patterns shifted to
    that hour, tanks at their initial levels). params: pump, node, target
    (model pressure units), hours [list] (default 0..23), minSpeed, maxSpeed.
    """
    pump_id, node_id = str(params.get("pump") or ""), str(params.get("node") or "")
    target = float(params.get("target") or 0)
    lo_s, hi_s = float(params.get("minSpeed") or 0.5), float(params.get("maxSpeed") or 1.2)
    if not (0 < lo_s < hi_s):
        raise BatchError("속도 범위가 올바르지 않습니다 (0 < 최소 < 최대).")
    hours = [float(h) for h in (params.get("hours") or list(range(24)))][:168]
    with Network(inp) as net:
        pi = net.link_index(pump_id)
        if net.link_types[pi - 1] != EN.PUMP:
            raise BatchError(f"{pump_id}는 펌프가 아닙니다.")
        ni = net.node_index(node_id)
        ps0 = net.time_param(EN.PATTERNSTART)
        net.set_time_param(EN.DURATION, 0)

        def solve(speed: float) -> tuple[float, float, float]:
            net.set_link(pi, EN.INITSTATUS, 1)
            net.set_link(pi, EN.INITSETTING, speed)
            net.lib.EN_solveH(net.ph)
            return net.get_node(ni, EN.PRESSURE), net.get_link(pi, EN.FLOW), net.get_link(pi, EN.ENERGY)

        out = []
        for h in hours:
            net.set_time_param(EN.PATTERNSTART, ps0 + int(h * 3600))
            try:
                p_hi = solve(hi_s)
                p_lo = solve(lo_s)
                if p_hi[0] < target:
                    speed, res, status = hi_s, p_hi, "max"
                elif p_lo[0] >= target:
                    speed, res, status = lo_s, p_lo, "min"
                else:
                    a, b = lo_s, hi_s
                    for _ in range(30):
                        m = (a + b) / 2
                        if solve(m)[0] >= target:
                            b = m
                        else:
                            a = m
                        if b - a < 1e-4:
                            break
                    speed, res, status = b, solve(b), "ok"
                out.append({"hour": h, "speed": round(speed, 4), "pressure": round(res[0], 3),
                            "flow": round(res[1], 4), "power": round(res[2], 3), "status": status})
            except EpanetException as e:
                out.append({"hour": h, "error": str(e)})
        return {"pump": pump_id, "node": node_id, "target": target, "rows": out,
                "note": "시간별 정상상태 해석입니다 (탱크는 초기 수위, 그 시각의 시간 제어만 반영)."}


# m3/s per model flow unit, by EN_getflowunits code
_M3S = {0: 0.0283168, 1: 6.30902e-5, 2: 0.0438126, 3: 0.0526168, 4: 0.0142764, 5: 1e-3,
        6: 1e-3 / 60, 7: 1e3 / 86400, 8: 1 / 3600, 9: 1 / 86400, 10: 1.0}


def job_energy(inp: bytes, params: dict) -> dict:
    """에너지 보고서: per-pump utilisation, efficiency, kWh, kWh/m3, peak kW and cost,
    integrated over every hydraulic step (prices from [ENERGY] incl. patterns)."""
    with Network(inp) as net:
        pumps = [k + 1 for k, t in enumerate(net.link_types) if t == EN.PUMP]
        if not pumps:
            raise BatchError("펌프가 없습니다.")
        fu = ctypes.c_int()
        net.lib.EN_getflowunits(net.ph, ctypes.byref(fu))
        to_m3s = _M3S.get(fu.value, 1.0)
        g_price = net.option(EN.GLOBALPRICE)
        g_pat = int(net.option(EN.GLOBALPATTERN))
        d_charge = net.option(EN.DEMANDCHARGE)
        g_eff = net.option(EN.GLOBALEFFIC)
        ps0 = net.time_param(EN.PATTERNSTART)
        price = {p: (net.get_link(p, EN.PUMP_ECOST) or g_price) for p in pumps}
        pat = {p: (int(net.get_link(p, EN.PUMP_EPAT)) or g_pat) for p in pumps}
        acc = {p: {"on": 0.0, "kwh": 0.0, "m3": 0.0, "effw": 0.0, "peak": 0.0, "cost": 0.0} for p in pumps}
        peak_total = 0.0
        total_s = 0.0
        t, tstep = ctypes.c_long(), ctypes.c_long(1)
        steady = net.duration == 0
        net.lib.EN_openH(net.ph)
        try:
            net.lib.EN_initH(net.ph, EN.NOSAVE)
            while True:
                net.lib.EN_runH(net.ph, ctypes.byref(t))
                now = t.value
                # Values at `now` hold until the next hydraulic time.
                vals = [(p, abs(net.get_link(p, EN.FLOW)), max(0.0, net.get_link(p, EN.ENERGY)), net.get_link(p, EN.PUMP_EFFIC)) for p in pumps]
                net.lib.EN_nextH(net.ph, ctypes.byref(tstep))
                dt = 3600.0 if steady else float(tstep.value)
                kw_total = 0.0
                for p, q, kw, eff in vals:
                    a = acc[p]
                    if q > 1e-9 and kw > 0:
                        a["on"] += dt
                        a["effw"] += eff * dt
                    a["kwh"] += kw * dt / 3600
                    a["m3"] += q * to_m3s * dt
                    a["peak"] = max(a["peak"], kw)
                    a["cost"] += kw * dt / 3600 * price[p] * net.pattern_value(pat[p], now, ps0)
                    kw_total += kw
                peak_total = max(peak_total, kw_total)
                total_s += dt
                if tstep.value <= 0:
                    break
        finally:
            net.lib.EN_closeH(net.ph)
        total_s = max(total_s, 1.0)
        rows = []
        for p in pumps:
            a = acc[p]
            rows.append({
                "id": net.link_ids[p - 1],
                "utilization": round(100 * a["on"] / total_s, 2),
                "efficiency": round(100 * a["effw"] / a["on"], 2) if a["on"] else None,  # DLL gives a fraction
                "kwh": round(a["kwh"], 3),
                "kwhPerM3": round(a["kwh"] / a["m3"], 5) if a["m3"] > 0 else None,
                "volume": round(a["m3"], 3),
                "avgKw": round(a["kwh"] * 3600 / a["on"], 3) if a["on"] else 0.0,
                "peakKw": round(a["peak"], 3),
                "cost": round(a["cost"], 3),
                "costPerDay": round(a["cost"] * 86400 / total_s, 3),
            })
        return {
            "pumps": rows,
            "hours": round(total_s / 3600, 3),
            "globalPrice": g_price,
            "globalEfficiency": g_eff,
            "demandCharge": d_charge,
            "peakKw": round(peak_total, 3),
            "demandCost": round(peak_total * d_charge, 3),
            "totalCost": round(sum(r["cost"] for r in rows) + peak_total * d_charge, 3),
            "steady": steady,
        }


PARAM_PROPS = {"pressure": ("node", EN.PRESSURE), "head": ("node", EN.HEAD), "demand": ("node", EN.DEMAND),
               "quality": ("node", EN.QUALITY), "flow": ("link", EN.FLOW), "velocity": ("link", EN.VELOCITY)}


def job_calibrate(inp: bytes, params: dict) -> dict:
    """자동 보정: per-group pipe roughness multipliers ('roughness') or wall
    coefficients ('wall') fitted to observations by bounded least squares.

    params: mode, param, observations [{location, time (s), value (model units)}],
    groups {name: [pipe ids]}, bounds [lo, hi], maxEvaluations.
    """
    from scipy.optimize import least_squares

    mode = params.get("mode") or "roughness"
    param = params.get("param") or "pressure"
    if param not in PARAM_PROPS:
        raise BatchError(f"보정할 수 없는 항목입니다: {param}")
    kind, prop = PARAM_PROPS[param]
    if mode == "wall" and param != "quality":
        raise BatchError("관벽 반응계수 보정에는 수질(잔류염소) 실측값이 필요합니다.")
    if mode == "roughness" and param == "quality":
        raise BatchError("조도 보정에는 수압·수두·유량 실측값을 쓰세요.")
    obs = [o for o in params.get("observations") or [] if o.get("location") is not None and o.get("value") is not None]
    if len(obs) < 2:
        raise BatchError("실측값이 2개 이상 필요합니다.")
    groups = {str(k): [str(i) for i in v] for k, v in (params.get("groups") or {}).items() if v}
    if not groups:
        raise BatchError("보정할 관로 그룹이 없습니다.")
    if len(groups) > 30:
        raise BatchError("그룹이 30개를 넘습니다. 묶음을 줄이세요.")
    lo, hi = (float(b) for b in (params.get("bounds") or ((0.5, 1.5) if mode == "roughness" else (-5.0, 0.0))))
    if not lo < hi:
        raise BatchError("범위가 올바르지 않습니다.")
    max_eval = int(params.get("maxEvaluations") or 200)

    with Network(inp) as net:
        locs = sorted({str(o["location"]) for o in obs})
        index = {}
        for loc in locs:
            index[loc] = net.node_index(loc) if kind == "node" else net.link_index(loc)
        cols = {loc: k for k, loc in enumerate(locs)}
        o_col = np.array([cols[str(o["location"])] for o in obs])
        o_t = np.array([float(o["time"]) for o in obs])
        o_v = np.array([float(o["value"]) for o in obs])
        g_names = list(groups)
        g_links = [[net.link_index(i) for i in groups[g]] for g in g_names]
        if mode == "wall":
            qt = ctypes.c_int(); tr = ctypes.c_int()
            net.lib.EN_getqualinfo(net.ph, ctypes.byref(qt), ctypes.create_string_buffer(64), ctypes.create_string_buffer(64), ctypes.byref(tr))
            if qt.value != 1:
                raise BatchError("모델의 수질 해석이 화학물질(CHEMICAL)이 아닙니다. 모델 설정에서 수질 모드를 바꾸세요.")
            base = [np.array([net.get_link(li, EN.KWALL) for li in links]) for links in g_links]
            x0 = np.array([float(np.clip(np.mean(b) if len(b) else (lo + hi) / 2, lo, hi)) for b in base])
            net.lib.EN_solveH(net.ph)  # wall reactions do not change hydraulics
        else:
            base = [np.array([net.get_link(li, EN.ROUGHNESS) for li in links]) for links in g_links]
            x0 = np.clip(np.ones(len(g_names)), lo, hi)

        def apply(x):
            for links, b, v in zip(g_links, base, x):
                for li, bv in zip(links, b):
                    net.set_link(li, EN.KWALL if mode == "wall" else EN.ROUGHNESS, v if mode == "wall" else bv * v)

        def simulate_values() -> np.ndarray:
            if mode == "wall":
                times, q = net.run_quality(nodes=[index[l] for l in locs])
                times = np.asarray(times, dtype=float)
                mat = q
            else:
                r = net.simulate(node_props=(prop,) if kind == "node" else (), link_props=(prop,) if kind == "link" else (),
                                 nodes=[index[l] for l in locs] if kind == "node" else [], links=[index[l] for l in locs] if kind == "link" else [])
                times = np.asarray(r.times, dtype=float)
                mat = (r.nodes if kind == "node" else r.links)[prop]
            if param == "flow":
                mat = np.abs(mat)
            out = np.empty(len(obs))
            for c in range(len(locs)):
                sel = o_col == c
                out[sel] = np.interp(o_t[sel], times, mat[:, c]) if len(times) > 1 else mat[0, c]
            return out

        evals = {"n": 0}

        def residuals(x):
            evals["n"] += 1
            apply(x)
            try:
                return simulate_values() - o_v
            except (EpanetException, BatchError):
                return np.full(len(obs), 1e3)

        r0 = residuals(x0)
        fit = least_squares(residuals, x0, bounds=(np.full(len(x0), lo), np.full(len(x0), hi)),
                            diff_step=0.02 if mode == "roughness" else 0.05, max_nfev=max_eval, xtol=1e-4, ftol=1e-6)
        r1 = residuals(fit.x)
        apply(x0)  # leave the network as it was (not strictly needed)
        rmse = lambda r: float(np.sqrt(np.mean(r ** 2)))
        out_groups = []
        for name, links, b, v in zip(g_names, g_links, base, fit.x):
            out_groups.append({
                "name": name, "pipes": len(links),
                "value": round(float(v), 4),
                "meanBefore": round(float(np.mean(b)), 4) if len(b) else None,
                "meanAfter": round(float(np.mean(b * v)) if mode == "roughness" else float(v), 4) if len(b) else None,
                "atBound": bool(abs(v - lo) < 1e-6 or abs(v - hi) < 1e-6),
            })
        per_loc = []
        for c, loc in enumerate(locs):
            sel = o_col == c
            per_loc.append({"location": loc, "n": int(sel.sum()), "before": round(rmse(r0[sel]), 4), "after": round(rmse(r1[sel]), 4)})
        return {
            "mode": mode, "param": param, "groups": out_groups, "locations": per_loc,
            "rmseBefore": round(rmse(r0), 4), "rmseAfter": round(rmse(r1), 4),
            "evaluations": evals["n"], "converged": bool(fit.success), "message": str(fit.message),
            "observations": len(obs),
        }


def _pump_controls(net: Network, pump_idx: set) -> list[int]:
    """1-based indices of simple controls acting on the given pumps."""
    n = ctypes.c_int()
    net.lib.EN_getcount(net.ph, 5, ctypes.byref(n))
    out = []
    t, li, ni = ctypes.c_int(), ctypes.c_int(), ctypes.c_int()
    st, lv = ctypes.c_double(), ctypes.c_double()
    for k in range(1, n.value + 1):
        net.lib.EN_getcontrol(net.ph, k, ctypes.byref(t), ctypes.byref(li), ctypes.byref(st), ctypes.byref(ni), ctypes.byref(lv))
        if li.value in pump_idx:
            out.append(k)
    return out


def job_pumpschedule(inp: bytes, params: dict) -> dict:
    """펌프 운전 최적화: hourly (pattern-step) on/off schedules that cut energy cost
    while keeping junction pressure >= minPressure and tanks back to their start
    level at the end. Existing simple controls on those pumps are ignored.

    params: pumps [ids] (default all), minPressure (model units), maxEvaluations,
    tankTolerance (fraction of tank range allowed below the start level at the end).
    """
    min_p = float(params.get("minPressure") or 0)
    budget = int(params.get("maxEvaluations") or 400)
    tank_tol = float(params.get("tankTolerance") if params.get("tankTolerance") is not None else 0.02)
    with Network(inp) as net:
        all_pumps = [k + 1 for k, t in enumerate(net.link_types) if t == EN.PUMP]
        pumps = [net.link_index(str(p)) for p in params.get("pumps") or []] or all_pumps
        if not pumps:
            raise BatchError("펌프가 없습니다.")
        if net.duration < net.pattern_step:
            raise BatchError("연장 기간 해석(EPS) 모델이어야 합니다.")
        rules = ctypes.c_int()
        net.lib.EN_getcount(net.ph, 6, ctypes.byref(rules))
        controls = _pump_controls(net, set(pumps))
        slot = net.pattern_step
        n_slots = int(np.ceil(net.duration / slot))
        if n_slots > 168:
            raise BatchError("시간 구간이 168개를 넘습니다 (패턴 간격을 늘리세요).")
        junctions = net.junction_indices()
        tanks = [k + 1 for k, t in enumerate(net.node_types) if t == EN.TANK]
        t_min = np.array([net.get_node(t, 20) for t in tanks])
        t_max = np.array([net.get_node(t, 21) for t in tanks])
        t_init = np.array([net.get_node(t, EN.TANKLEVEL) for t in tanks])  # initial level
        t_elev = np.array([net.get_node(t, EN.ELEVATION) for t in tanks])
        ps0 = net.time_param(EN.PATTERNSTART)
        g_price = net.option(EN.GLOBALPRICE)
        g_pat = int(net.option(EN.GLOBALPATTERN))
        price = {p: (net.get_link(p, EN.PUMP_ECOST) or g_price or 1.0) for p in pumps}
        ppat = {p: (int(net.get_link(p, EN.PUMP_EPAT)) or g_pat) for p in pumps}
        priced = bool(g_price or any(net.get_link(p, EN.PUMP_ECOST) for p in pumps))

        def evaluate(sched: dict | None) -> dict:
            t, tstep = ctypes.c_long(), ctypes.c_long(1)
            cost = kwh = 0.0
            pmin = np.inf
            level_end = t_init.copy()
            net.lib.EN_openH(net.ph)
            try:
                net.lib.EN_initH(net.ph, EN.NOSAVE)
                while True:
                    net.lib.EN_runH(net.ph, ctypes.byref(t))
                    now = t.value
                    kw = [max(0.0, net.get_link(p, EN.ENERGY)) for p in pumps]
                    if net.is_report_time(now) and junctions:
                        pmin = min(pmin, min(net.get_node(j, EN.PRESSURE) for j in junctions))
                    if tanks:
                        # EN_TANKLEVEL reads back the *initial* level; the current one is head - bottom.
                        level_end = np.array([net.get_node(tk, EN.HEAD) for tk in tanks]) - t_elev
                    net.lib.EN_nextH(net.ph, ctypes.byref(tstep))
                    dt = float(tstep.value)
                    for p, k in zip(pumps, kw):
                        kwh += k * dt / 3600
                        cost += k * dt / 3600 * price[p] * net.pattern_value(ppat[p], now, ps0)
                    if tstep.value <= 0:
                        break
            finally:
                net.lib.EN_closeH(net.ph)
            short = np.maximum(0.0, (t_init - tank_tol * (t_max - t_min)) - level_end) if tanks else np.zeros(0)
            violation = max(0.0, min_p - pmin) + float(np.sum(short))
            return {"cost": cost, "kwh": kwh, "minPressure": float(pmin), "violation": violation, "tankEnd": level_end.tolist()}

        base = evaluate(None)
        # Baseline schedule = what the existing controls did (pump on in each slot?).
        status = {p: np.zeros(n_slots) for p in pumps}
        r = net.simulate(node_props=(), link_props=(EN.FLOW,), nodes=[], links=pumps)
        for c, p in enumerate(pumps):
            for k, tt in enumerate(r.times):
                s = min(n_slots - 1, int(tt // slot))
                if abs(r.links[EN.FLOW][k, c]) > 1e-9:
                    status[p][s] = 1
        # Take control of the pumps: drop their simple controls, drive them by a pattern.
        for k in sorted(controls, reverse=True):
            net.delete_control(k)
        for p in pumps:
            pid = f"SPSCH{p}"
            net.add_pattern(pid, [1.0] * n_slots)
            net.set_link(p, 15, float(_pattern_index(net, pid)))
            net.set_link(p, EN.INITSTATUS, 1)
        evals = {"n": 0}

        def run(sched):
            evals["n"] += 1
            for p in pumps:
                net.add_pattern(f"SPSCH{p}", [float(v) for v in sched[p]])
            return evaluate(sched)

        cur = {p: status[p].copy() for p in pumps}
        cur_res = run(cur)
        if cur_res["violation"] > 1e-6:
            # Pattern replay of the control schedule is not feasible (e.g. sub-slot switching): start from all on.
            cur = {p: np.ones(n_slots) for p in pumps}
            cur_res = run(cur)
        start_res = cur_res
        slot_price = {p: np.array([price[p] * net.pattern_value(ppat[p], k * slot, ps0) for k in range(n_slots)]) for p in pumps}

        improved = True
        while improved and evals["n"] < budget:
            improved = False
            # 1) switch off the most expensive running slots
            cand = sorted(((slot_price[p][k], p, k) for p in pumps for k in range(n_slots) if cur[p][k] > 0), reverse=True)
            for _, p, k in cand:
                if evals["n"] >= budget:
                    break
                trial = {q: cur[q].copy() for q in pumps}
                trial[p][k] = 0
                res = run(trial)
                if res["violation"] <= 1e-6 and res["cost"] < cur_res["cost"] - 1e-9:
                    cur, cur_res, improved = trial, res, True
            # 2) move running time from expensive to cheap slots
            for p in pumps:
                on = [k for k in range(n_slots) if cur[p][k] > 0]
                off = [k for k in range(n_slots) if cur[p][k] == 0]
                for k_on in sorted(on, key=lambda k: -slot_price[p][k]):
                    for k_off in sorted(off, key=lambda k: slot_price[p][k]):
                        if evals["n"] >= budget or slot_price[p][k_off] >= slot_price[p][k_on]:
                            break
                        trial = {q: cur[q].copy() for q in pumps}
                        trial[p][k_on], trial[p][k_off] = 0, 1
                        res = run(trial)
                        if res["violation"] <= 1e-6 and res["cost"] < cur_res["cost"] - 1e-9:
                            cur, cur_res, improved = trial, res, True
                            off.remove(k_off)
                            off.append(k_on)
                            break
        return {
            "pumps": [net.link_ids[p - 1] for p in pumps],
            "slotSeconds": slot,
            "schedules": {net.link_ids[p - 1]: [int(v) for v in cur[p]] for p in pumps},
            "before": {k: round(v, 3) if isinstance(v, float) else v for k, v in base.items()},
            "start": {k: round(v, 3) if isinstance(v, float) else v for k, v in start_res.items()},
            "after": {k: round(v, 3) if isinstance(v, float) else v for k, v in cur_res.items()},
            "priced": priced,
            "evaluations": evals["n"],
            "removedControls": len(controls),
            "rules": rules.value,
            "feasible": cur_res["violation"] <= 1e-6,
        }


def _pattern_index(net: Network, pid: str) -> int:
    idx = ctypes.c_int()
    net.lib.EN_getpatternindex(net.ph, pid.encode(), ctypes.byref(idx))
    return idx.value


JOBS = {
    "summary": job_summary,
    "fireflow": job_fireflow,
    "criticality": job_criticality,
    "servicearea": job_servicearea,
    "vsp": job_vsp,
    "energy": job_energy,
    "calibrate": job_calibrate,
    "pumpschedule": job_pumpschedule,
}


def get_job(name: str):
    fn = JOBS.get(name)
    if fn is None:
        raise BatchError(f"알 수 없는 분석입니다: {name}")
    return fn