design_optimizer.py

"""Deterministic, COC-constrained design optimizer.

Given a connected model's base INP, this module screens a bounded, fully
enumerated set of candidate designs -- variations on conduit diameter,
conduit roughness, and/or storage depth -- against a FIXED set of City of
Calgary design-standard constraints, and ranks the constraint-satisfying
candidates by a single VARIABLE, user-selected design objective.

Design intent
-------------
SWMM re-simulation is not free, and a black-box gradient/heuristic optimizer
over a discontinuous, constraint-heavy hydraulic model is hard to audit and
easy to trust more than it deserves. This module instead runs a **design of
experiments**: every candidate that gets evaluated is one you can see in the
returned ``evaluations`` list, with its own constraint pass/fail detail and
objective value. There is no hidden search path. Two candidate-generation
modes are supported:

  * ``variables``  -- a dict of {object_id: [candidate values]} per override
    category; the full Cartesian product is evaluated (grid search), capped
    by ``max_evaluations``.
  * ``candidates`` -- fully-formed override dicts supplied directly by the
    caller, for a hand-picked shortlist rather than a grid.

Fixed constraints -- confirm before treating as authoritative
---------------------------------------------------------------
``COC_DEFAULT_CONSTRAINTS`` below are commonly-used stormwater design
thresholds (SWMM-modelling continuity guidance, typical municipal
self-cleansing/non-erosive velocity bands, no-surcharge and no-surface-
flooding criteria). They are NOT transcribed from a specific, current
edition of the City of Calgary Stormwater Management & Design Manual or its
Industry Bulletins -- that document was not available to read while writing
this module. Every default is overridable via the ``constraints`` argument,
and the engineer of record must confirm each value against the current
SWMDM/Industry Bulletins and project-specific approval conditions before
relying on a "PASS" result for a real submission.
"""
from __future__ import annotations

from dataclasses import dataclass
from itertools import product
from pathlib import Path
from typing import Any, Iterable

import scenario_manager as sm

# ---------------------------------------------------------------------------
# Fixed constraint defaults -- see module docstring caveat above.
# ---------------------------------------------------------------------------

COC_DEFAULT_CONSTRAINTS: dict[str, Any] = {
    # No node may surcharge above the pipe crown at the design storm.
    # Expressed as the model's own "depth / diameter (or full_depth)" ratio.
    "max_depth_ratio": 1.0,
    # Self-cleansing minimum velocity, commonly-cited municipal default.
    "min_velocity_mps": 0.75,
    # Non-erosive maximum velocity for a typical rigid pipe material.
    "max_velocity_mps": 6.0,
    # No surface flooding at any node during the design storm.
    "max_node_flooding_cms": 0.0,
    # SWMM continuity error acceptance band (applies to both the runoff and
    # flow/routing continuity errors reported by the engine).
    "max_continuity_error_pct": 5.0,
    # Project-specific allowable peak release rate. Left as None (not
    # checked) unless the caller supplies a value and an outfall_link_id --
    # there is no generic default for this, it is genuinely project-specific.
    "max_allowable_outfall_flow_cms": None,
}

_OVERRIDE_CATEGORIES = (
    "conduit_diameter_overrides",
    "conduit_roughness_overrides",
    "storage_depth_overrides",
)

_DIRECT_OBJECTIVE_METRICS = {
    "maximum_storage_volume": "Maximum Storage Volume",
    "maximum_storage_depth": "Maximum Storage Depth",
    "maximum_link_velocity": "Maximum Link Velocity",
    "peak_link_flow": "Peak Link Flow",
    "maximum_modelled_depth_ratio": "Maximum Modelled Depth Ratio",
    "maximum_node_flooding": "Maximum Node Flooding",
    "peak_subcatchment_runoff": "Peak Subcatchment Runoff",
}

SUPPORTED_OBJECTIVE_METRICS = sorted(
    set(_DIRECT_OBJECTIVE_METRICS) | {"total_conduit_volume_m3"}
)


@dataclass
class OptimizationCandidate:
    candidate_id: str
    overrides: dict[str, dict[str, float]]
    summary: dict[str, Any]
    objective_value: float | None
    feasible: bool
    constraint_results: list[dict[str, Any]]
    scenario_id: str


def _read_conduit_geometry(inp_path: str | Path) -> dict[str, dict[str, float]]:
    """Read every conduit's length (CONDUITS) and diameter (XSECTIONS) from
    the base INP, for computing total_conduit_volume_m3. Only circular
    (Geom1 = diameter) links are included; other shapes are skipped rather
    than guessed at."""
    text = Path(inp_path).read_text(encoding="utf-8", errors="ignore")
    _, sections = sm._split_sections(text)  # noqa: SLF001 -- reuse the same INP splitter as the rest of this codebase
    geometry: dict[str, dict[str, float]] = {}
    for line in sections.get("CONDUITS", []):
        parsed = sm._data_tokens(line)  # noqa: SLF001
        if not parsed:
            continue
        tokens, _ = parsed
        if len(tokens) > 3:
            try:
                geometry[tokens[0]] = {"length": float(tokens[3]), "diameter": 0.0}
            except ValueError:
                continue
    for line in sections.get("XSECTIONS", []):
        parsed = sm._data_tokens(line)  # noqa: SLF001
        if not parsed:
            continue
        tokens, _ = parsed
        if len(tokens) > 2 and tokens[0] in geometry:
            try:
                geometry[tokens[0]]["diameter"] = float(tokens[2])
            except ValueError:
                continue
    return geometry


def _total_conduit_volume_m3(
    base_geometry: dict[str, dict[str, float]],
    diameter_overrides: dict[str, float],
) -> float:
    import math
    total = 0.0
    for conduit_id, geom in base_geometry.items():
        diameter = diameter_overrides.get(conduit_id, geom.get("diameter", 0.0))
        length = geom.get("length", 0.0)
        total += math.pi / 4.0 * diameter * diameter * length
    return total


def _build_candidates_from_variables(
    variables: dict[str, dict[str, list[float]]],
    max_evaluations: int,
) -> list[dict[str, dict[str, float]]]:
    """Cartesian-product a {category: {object_id: [values]}} spec into a
    bounded list of {category: {object_id: value}} override dicts."""
    axes: list[tuple[str, str, list[float]]] = []
    for category, per_object in (variables or {}).items():
        if category not in _OVERRIDE_CATEGORIES:
            raise ValueError(
                f"Unsupported variable category '{category}'. Supported: "
                f"{', '.join(_OVERRIDE_CATEGORIES)}"
            )
        for object_id, values in per_object.items():
            if not isinstance(values, list) or not values:
                raise ValueError(
                    f"variables['{category}']['{object_id}'] must be a non-empty list of candidate values"
                )
            axes.append((category, object_id, [float(v) for v in values]))

    if not axes:
        return [{}]

    combo_count = 1
    for _, _, values in axes:
        combo_count *= len(values)
    if combo_count > max_evaluations:
        raise ValueError(
            f"The variable grid has {combo_count} combinations, which exceeds "
            f"max_evaluations={max_evaluations}. Narrow the candidate lists, "
            f"raise max_evaluations, or switch to the 'candidates' shortlist mode."
        )

    combos: list[dict[str, dict[str, float]]] = []
    value_lists = [values for _, _, values in axes]
    for combo in product(*value_lists):
        overrides: dict[str, dict[str, float]] = {}
        for (category, object_id, _), value in zip(axes, combo):
            overrides.setdefault(category, {})[object_id] = value
        combos.append(overrides)
    return combos


def _evaluate_constraints(
    summary: dict[str, Any],
    constraints: dict[str, Any],
    outfall_flow_cms: float | None,
) -> list[dict[str, Any]]:
    results: list[dict[str, Any]] = []

    def check(name: str, actual: Any, passed: bool, limit: Any, note: str = "") -> None:
        results.append({
            "constraint": name, "actual": actual, "limit": limit,
            "passed": bool(passed), "note": note,
        })

    max_depth_ratio = constraints.get("max_depth_ratio")
    if max_depth_ratio is not None:
        actual = summary.get("Maximum Modelled Depth Ratio", 0.0)
        check("max_depth_ratio", actual, actual <= max_depth_ratio, max_depth_ratio,
              f"Controlling link: {summary.get('Depth-Ratio Link', '')}")

    min_v = constraints.get("min_velocity_mps")
    max_v = constraints.get("max_velocity_mps")
    if min_v is not None or max_v is not None:
        actual = summary.get("Maximum Link Velocity", 0.0)
        if max_v is not None:
            check("max_velocity_mps", actual, actual <= max_v, max_v,
                  f"Controlling link: {summary.get('Velocity Link', '')}")
        if min_v is not None:
            check("min_velocity_mps", actual, actual >= min_v, min_v,
                  "Self-cleansing screen uses the model's maximum observed "
                  "velocity per link as a proxy; confirm against your "
                  "project's actual low-flow/self-cleansing check.")

    max_flood = constraints.get("max_node_flooding_cms")
    if max_flood is not None:
        actual = summary.get("Maximum Node Flooding", 0.0)
        check("max_node_flooding_cms", actual, actual <= max_flood, max_flood)

    max_cont = constraints.get("max_continuity_error_pct")
    if max_cont is not None:
        runoff_err = abs(summary.get("Runoff Error (%)", 0.0))
        flow_err = abs(summary.get("Flow Error (%)", 0.0))
        check("max_continuity_error_pct (runoff)", runoff_err, runoff_err <= max_cont, max_cont)
        check("max_continuity_error_pct (flow/routing)", flow_err, flow_err <= max_cont, max_cont)

    max_outfall = constraints.get("max_allowable_outfall_flow_cms")
    if max_outfall is not None:
        if outfall_flow_cms is None:
            check("max_allowable_outfall_flow_cms", None, False, max_outfall,
                  "NOT EVALUATED: max_allowable_outfall_flow_cms was set but "
                  "outfall_link_id was not supplied, so no outfall-specific "
                  "flow could be measured. This candidate is scored infeasible "
                  "so an unchecked constraint can never silently pass.")
        else:
            check("max_allowable_outfall_flow_cms", outfall_flow_cms,
                  outfall_flow_cms <= max_outfall, max_outfall)

    return results


def _objective_value(
    objective_metric: str,
    summary: dict[str, Any],
    base_geometry: dict[str, dict[str, float]],
    overrides: dict[str, dict[str, float]],
) -> float | None:
    if objective_metric == "total_conduit_volume_m3":
        return _total_conduit_volume_m3(
            base_geometry, overrides.get("conduit_diameter_overrides", {})
        )
    field = _DIRECT_OBJECTIVE_METRICS.get(objective_metric)
    if field is None:
        raise ValueError(
            f"Unsupported objective metric '{objective_metric}'. Supported: "
            f"{', '.join(SUPPORTED_OBJECTIVE_METRICS)}"
        )
    value = summary.get(field)
    return float(value) if value is not None else None


def optimize(
    base_inp_path: str | Path,
    work_dir: str | Path,
    *,
    variables: dict[str, dict[str, list[float]]] | None,
    candidates: list[dict[str, dict[str, float]]] | None,
    objective_metric: str,
    objective_direction: str,
    constraints: dict[str, Any] | None,
    outfall_link_id: str,
    max_evaluations: int,
    scenario_prefix: str = "opt",
) -> dict[str, Any]:
    if objective_direction not in ("minimize", "maximize"):
        raise ValueError("objective_direction must be 'minimize' or 'maximize'")
    if bool(variables) == bool(candidates):
        if not variables and not candidates:
            raise ValueError("Supply either 'variables' (grid) or 'candidates' (shortlist).")
        # both supplied is allowed (candidates wins for explicitness) -- fallthrough

    resolved_constraints = dict(COC_DEFAULT_CONSTRAINTS)
    resolved_constraints.update(constraints or {})

    if candidates:
        if len(candidates) > max_evaluations:
            raise ValueError(
                f"{len(candidates)} candidates supplied, which exceeds "
                f"max_evaluations={max_evaluations}."
            )
        combos = candidates
    else:
        combos = _build_candidates_from_variables(variables or {}, max_evaluations)

    base_geometry = _read_conduit_geometry(base_inp_path)

    evaluations: list[OptimizationCandidate] = []
    for index, overrides in enumerate(combos, start=1):
        candidate_id = f"{scenario_prefix}_{index}"
        definition = sm.ScenarioDefinition(
            scenario_id=candidate_id,
            scenario_name=f"Optimization candidate {index}",
            conduit_diameter_overrides=dict(overrides.get("conduit_diameter_overrides", {})),
            conduit_roughness_overrides=dict(overrides.get("conduit_roughness_overrides", {})),
            storage_depth_overrides=dict(overrides.get("storage_depth_overrides", {})),
            review_status="Preliminary optimization candidate",
        )
        record = sm.run_scenario(base_inp_path, definition, work_dir=work_dir)
        summary = record["summary"]

        outfall_flow_cms = None
        if outfall_link_id:
            link_values = (record["results"].get("link_ts", {}) or {}).get(outfall_link_id)
            if link_values:
                flows = link_values.get("flow", []) or [0.0]
                outfall_flow_cms = max(abs(float(x)) for x in flows)

        constraint_results = _evaluate_constraints(summary, resolved_constraints, outfall_flow_cms)
        feasible = all(c["passed"] for c in constraint_results)
        objective_value = _objective_value(objective_metric, summary, base_geometry, overrides)

        evaluations.append(OptimizationCandidate(
            candidate_id=candidate_id,
            overrides=overrides,
            summary=summary,
            objective_value=objective_value,
            feasible=feasible,
            constraint_results=constraint_results,
            scenario_id=record["definition"]["scenario_id"],
        ))

    def sort_key(c: OptimizationCandidate):
        infeasible_penalty = 0 if c.feasible else 1
        value = c.objective_value
        if value is None:
            value = float("inf") if objective_direction == "minimize" else float("-inf")
        direction_value = value if objective_direction == "minimize" else -value
        return (infeasible_penalty, direction_value)

    ranked = sorted(evaluations, key=sort_key)
    best_feasible = next((c for c in ranked if c.feasible), None)

    return {
        "objective": {"metric": objective_metric, "direction": objective_direction},
        "constraints_used": resolved_constraints,
        "constraints_source_note": (
            "Fixed constraint defaults are commonly-used stormwater design "
            "thresholds, not a transcription of a specific SWMDM/Industry "
            "Bulletin edition. Confirm every value against current City of "
            "Calgary criteria and project approval conditions."
        ),
        "candidates_evaluated": len(evaluations),
        "feasible_count": sum(1 for c in evaluations if c.feasible),
        "best_feasible_candidate_id": best_feasible.candidate_id if best_feasible else None,
        "ranked_candidates": [
            {
                "rank": i + 1,
                "candidate_id": c.candidate_id,
                "feasible": c.feasible,
                "objective_value": c.objective_value,
                "overrides": c.overrides,
                "constraint_results": c.constraint_results,
                "summary": c.summary,
            }
            for i, c in enumerate(ranked)
        ],
    }