"""Automatic calibration (roughness / wall coefficient) and pump scheduling jobs, on Net1 with synthetic truth."""
import re
from pathlib import Path

import pytest

from backend.batch import BatchError, Network
from backend.epanet import EN
from backend.jobs import job_calibrate, job_pumpschedule

NET1 = (Path(__file__).resolve().parents[2] / "public" / "Net1.inp").read_text()
PIPES = [l.split()[0] for l in re.search(r"\[PIPES\](.*?)\[", NET1, re.S).group(1).splitlines() if l.strip() and not l.strip().startswith(";")]


def synthetic(factor=None, wall=None, prop=EN.PRESSURE, locs=("22", "31", "12", "23"), groups=None):
    with Network(NET1.encode()) as n:
        idx = [n.node_index(l) for l in locs]
        for pid in PIPES:
            li = n.link_index(pid)
            f = factor(pid) if callable(factor) else factor
            if f:
                n.set_link(li, EN.ROUGHNESS, n.get_link(li, EN.ROUGHNESS) * f)
            if wall is not None:
                n.set_link(li, EN.KWALL, wall)
        if prop == EN.QUALITY:
            n.lib.EN_solveH(n.ph)
            times, mat = n.run_quality(nodes=idx)
        else:
            r = n.simulate(node_props=(prop,), nodes=idx)
            times, mat = r.times, r.nodes[prop]
    return [{"location": l, "time": t, "value": float(mat[k, c])} for k, t in enumerate(times) if k % 2 == 0 for c, l in enumerate(locs)]


def test_roughness_recovers_two_group_truth():
    big = {"10", "11", "12", "111", "110"}
    obs = synthetic(factor=lambda p: 0.75 if p in big else 1.1, locs=("10", "11", "12", "13", "21", "22", "23", "31", "32"))
    groups = {"main": sorted(big), "rest": [p for p in PIPES if p not in big]}
    r = job_calibrate(NET1.encode(), {"mode": "roughness", "param": "pressure", "observations": obs, "groups": groups, "bounds": [0.5, 1.5]})
    got = {g["name"]: g["value"] for g in r["groups"]}
    assert got["main"] == pytest.approx(0.75, abs=0.02) and got["rest"] == pytest.approx(1.1, abs=0.05)
    assert r["rmseAfter"] < 0.05 * r["rmseBefore"]


def test_wall_coefficient_from_chlorine():
    obs = synthetic(wall=-0.4, prop=EN.QUALITY)
    r = job_calibrate(NET1.encode(), {"mode": "wall", "param": "quality", "observations": obs, "groups": {"all": PIPES}, "bounds": [-3, 0]})
    assert r["groups"][0]["value"] == pytest.approx(-0.4, abs=0.01) and r["groups"][0]["meanBefore"] == -1.0


def test_calibrate_input_errors():
    with pytest.raises(BatchError):
        job_calibrate(NET1.encode(), {"mode": "wall", "param": "pressure", "observations": synthetic(), "groups": {"a": PIPES}})
    with pytest.raises(BatchError):
        job_calibrate(NET1.encode(), {"observations": [], "groups": {"a": PIPES}})
    with pytest.raises(BatchError):
        job_calibrate(NET1.encode(), {"observations": synthetic(), "groups": {}})


def test_pump_schedule_keeps_constraints_and_reports_controls():
    priced = NET1.replace("[PATTERNS]", "[PATTERNS]\n PRICE 0.5 0.5 0.5 1.5 1.5 1.5 1.5 1.5 1.5 1.5 1.5 0.5")
    priced = re.sub(r"Global Price\s+0\.0", "Global Price 0.1\n Global Pattern PRICE", priced)
    r = job_pumpschedule(priced.encode(), {"minPressure": 20})
    assert r["pumps"] == ["9"] and r["removedControls"] == 2 and r["slotSeconds"] == 7200
    assert len(r["schedules"]["9"]) == 12 and r["feasible"]
    assert r["after"]["violation"] == 0 and r["after"]["minPressure"] >= 20
    assert r["after"]["tankEnd"][0] >= 120 - 0.02 * 50 - 1e-6
    # never worse than the feasible starting schedule
    assert r["after"]["cost"] <= r["start"]["cost"] + 1e-9
    # Net1's own controls leave the tank ~4.6 ft lower after a day: reported, not hidden.
    assert r["before"]["violation"] > 0