{
  "version": "0.1",
  "parameters": {
    "pct_imperv_s1": {
      "type": "triangular",
      "lower": 15.0,
      "upper": 40.0,
      "baseline": "from_model",
      "precision": 3,
      "source": "User-defined land-cover uncertainty around the current model value."
    },
    "n_imperv_s1": {
      "type": "trapezoidal",
      "lower": 0.01,
      "upper": 0.03,
      "core_width": 0.004,
      "baseline": "from_model",
      "precision": 4,
      "source": "Baseline-centered plausible roughness plateau."
    },
    "ksat_s1": {
      "type": "triangular",
      "lower": 4.0,
      "upper": 14.0,
      "baseline": "from_model",
      "precision": 3,
      "source": "Soil hydraulic conductivity epistemic uncertainty."
    }
  }
}

{
  "method": "idf",
  "idf": {
    "type": "chicago",
    "duration_minutes": 360,
    "return_period_years": 100,
    "interval_minutes": 5,
    "chicago_peak_position": 0.4,
    "params": {
      "a": {"value": 30.0, "ci": [27.0, 33.0]},
      "b": {"value": 0.5, "ci": [0.4, 0.6]},
      "c": {"value": 0.7, "ci": [0.65, 0.75]}
    }
  },
  "n_realisations": 200
}

{
  "method": "perturbation",
  "perturbation": {
    "model": "autocorrelated",
    "sigma": 0.15,
    "ar1_coefficient": 0.6,
    "preserve_total_volume": false
  },
  "n_realisations": 200,
  "input_rainfall_path": "00_inputs/rainfall_obs.csv"
}

{
  "alpha_levels": [0.0, 0.25, 0.5, 0.75, 1.0],
  "sampling": {
    "method": "lhs",
    "samples_per_alpha": 6,
    "seed": 42
  },
  "outputs": {
    "swmm_node": "O1",
    "metrics": [
      "peak_flow",
      "runoff_continuity_error_percent",
      "flow_routing_continuity_error_percent"
    ]
  }
}

#!/usr/bin/env python3
from __future__ import annotations

import argparse
import json
import math
from dataclasses import dataclass
from pathlib import Path
from typing import Any


def load_json(path: Path) -> Any:
    return json.loads(path.read_text(encoding="utf-8"))


def write_json(path: Path, obj: Any) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_text(json.dumps(obj, indent=2, sort_keys=True), encoding="utf-8")


def finite_float(value: Any, *, field: str) -> float:
    if isinstance(value, bool):
        raise ValueError(f"{field} must be a finite number")
    try:
        out = float(value)
    except (TypeError, ValueError) as exc:
        raise ValueError(f"{field} must be a finite number") from exc
    if not math.isfinite(out):
        raise ValueError(f"{field} must be a finite number")
    return out


def normalize_type(spec: dict[str, Any]) -> str:
    raw = str(spec.get("type", spec.get("membership", ""))).strip().lower()
    aliases = {
        "fixed": "crisp",
        "constant": "crisp",
        "tri": "triangular",
        "triangle": "triangular",
        "trap": "trapezoidal",
        "trapezoid": "trapezoidal",
    }
    return aliases.get(raw, raw)


def read_baseline_values(inp_path: Path, patch_map: dict[str, Any]) -> dict[str, float]:
    lines = inp_path.read_text(encoding="utf-8", errors="ignore").splitlines()
    baselines: dict[str, float] = {}
    current_section: str | None = None

    for raw in lines:
        stripped = raw.strip()
        if not stripped:
            continue
        if stripped.startswith("[") and stripped.endswith("]"):
            current_section = stripped.upper()
            continue
        if stripped.startswith(";"):
            continue

        code = raw.split(";", 1)[0]
        tokens = code.split()
        if not tokens:
            continue

        for name, spec in patch_map.items():
            if name in baselines:
                continue
            if current_section != str(spec["section"]).upper():
                continue
            if tokens[0] != str(spec["object"]):
                continue
            idx = int(spec["field_index"])
            if idx >= len(tokens):
                raise IndexError(f"Field index {idx} out of range for {name} on line: {raw}")
            baselines[name] = finite_float(tokens[idx], field=f"baseline for {name}")

    missing = sorted(set(patch_map) - set(baselines))
    if missing:
        raise KeyError(f"Could not resolve baseline value(s) from INP: {missing}")
    return baselines


@dataclass(frozen=True)
class ResolvedFuzzyParameter:
    name: str
    kind: str
    lower: float
    upper: float
    baseline: float | None = None
    core_lower: float | None = None
    core_upper: float | None = None
    precision: int | None = None
    value_type: str = "float"
    source: str | None = None

    def alpha_interval(self, alpha: float) -> tuple[float, float]:
        if alpha < 0.0 or alpha > 1.0:
            raise ValueError(f"alpha must be within [0, 1], got {alpha}")

        if self.kind == "crisp":
            value = self.baseline if self.baseline is not None else self.lower
            return self._format(value), self._format(value)
        if self.kind == "interval":
            return self._format(self.lower), self._format(self.upper)
        if self.kind == "triangular":
            if self.baseline is None:
                raise ValueError(f"Triangular parameter '{self.name}' is missing baseline")
            lo = self.lower + alpha * (self.baseline - self.lower)
            hi = self.upper - alpha * (self.upper - self.baseline)
            return self._format(lo), self._format(hi)
        if self.kind == "trapezoidal":
            if self.core_lower is None or self.core_upper is None:
                raise ValueError(f"Trapezoidal parameter '{self.name}' is missing core bounds")
            lo = self.lower + alpha * (self.core_lower - self.lower)
            hi = self.upper - alpha * (self.upper - self.core_upper)
            return self._format(lo), self._format(hi)

        raise ValueError(f"Unsupported fuzzy parameter kind: {self.kind}")

    def _format(self, value: float) -> float | int:
        if self.value_type == "int":
            return int(round(value))
        if self.precision is not None:
            return round(float(value), self.precision)
        return float(value)

    def to_dict(self) -> dict[str, Any]:
        out: dict[str, Any] = {
            "type": self.kind,
            "lower": self.lower,
            "upper": self.upper,
            "baseline": self.baseline,
            "core_lower": self.core_lower,
            "core_upper": self.core_upper,
            "value_type": self.value_type,
            "precision": self.precision,
        }
        if self.source:
            out["source"] = self.source
        return out


def resolve_baseline(name: str, spec: dict[str, Any], baseline_values: dict[str, float]) -> float | None:
    raw = spec.get("baseline", "from_model")
    if raw in {None, "none"}:
        return None
    if raw == "from_model":
        if name not in baseline_values:
            raise KeyError(f"Parameter '{name}' was not found in the patch map / base INP baseline values")
        return baseline_values[name]
    return finite_float(raw, field=f"{name}.baseline")


def parse_value_type(name: str, spec: dict[str, Any]) -> tuple[str, int | None]:
    raw_type = str(spec.get("value_type", spec.get("type_hint", "float"))).strip().lower()
    if raw_type in {"int", "integer"}:
        return "int", None
    if raw_type not in {"float", "number"}:
        raise ValueError(f"{name}.value_type must be 'float' or 'int'")
    precision = spec.get("precision")
    if precision is None:
        return "float", None
    precision_int = int(precision)
    if precision_int < 0:
        raise ValueError(f"{name}.precision must be >= 0")
    return "float", precision_int


def _require_order(name: str, values: list[tuple[str, float]]) -> None:
    for (left_name, left), (right_name, right) in zip(values, values[1:]):
        if left > right:
            raise ValueError(f"{name} requires {left_name} <= {right_name}; got {left} > {right}")


def resolve_parameter(name: str, spec: dict[str, Any], baseline_values: dict[str, float]) -> ResolvedFuzzyParameter:
    kind = normalize_type(spec)
    value_type, precision = parse_value_type(name, spec)
    source = spec.get("source")

    if kind == "crisp":
        baseline = resolve_baseline(name, spec, baseline_values)
        if baseline is None:
            baseline = finite_float(spec.get("value"), field=f"{name}.value")
        return ResolvedFuzzyParameter(
            name=name,
            kind=kind,
            lower=baseline,
            upper=baseline,
            baseline=baseline,
            precision=precision,
            value_type=value_type,
            source=source,
        )

    lower = finite_float(spec.get("lower", spec.get("a")), field=f"{name}.lower")
    upper = finite_float(spec.get("upper", spec.get("c" if kind == "triangular" else "d")), field=f"{name}.upper")
    if lower > upper:
        raise ValueError(f"{name}.lower must be <= upper")

    if kind == "interval":
        return ResolvedFuzzyParameter(
            name=name,
            kind=kind,
            lower=lower,
            upper=upper,
            precision=precision,
            value_type=value_type,
            source=source,
        )

    if kind == "triangular":
        baseline = resolve_baseline(name, spec, baseline_values)
        if baseline is None:
            baseline = finite_float(spec.get("mode", spec.get("b")), field=f"{name}.mode")
        _require_order(name, [("lower", lower), ("baseline", baseline), ("upper", upper)])
        return ResolvedFuzzyParameter(
            name=name,
            kind=kind,
            lower=lower,
            upper=upper,
            baseline=baseline,
            precision=precision,
            value_type=value_type,
            source=source,
        )

    if kind == "trapezoidal":
        baseline = resolve_baseline(name, spec, baseline_values)
        if "core_lower" in spec or "core_upper" in spec or "b" in spec:
            core_lower = finite_float(spec.get("core_lower", spec.get("b")), field=f"{name}.core_lower")
            core_upper = finite_float(spec.get("core_upper", spec.get("c")), field=f"{name}.core_upper")
        elif "core_width" in spec:
            if baseline is None:
                raise ValueError(f"{name}.core_width requires a baseline")
            width = finite_float(spec["core_width"], field=f"{name}.core_width")
            if width < 0:
                raise ValueError(f"{name}.core_width must be >= 0")
            core_lower = baseline - width / 2.0
            core_upper = baseline + width / 2.0
        else:
            if baseline is None:
                raise ValueError(f"{name} trapezoidal spec requires core bounds or core_width")
            core_lower = baseline
            core_upper = baseline

        _require_order(
            name,
            [("lower", lower), ("core_lower", core_lower), ("core_upper", core_upper), ("upper", upper)],
        )
        return ResolvedFuzzyParameter(
            name=name,
            kind=kind,
            lower=lower,
            upper=upper,
            baseline=baseline,
            core_lower=core_lower,
            core_upper=core_upper,
            precision=precision,
            value_type=value_type,
            source=source,
        )

    raise ValueError(f"Unsupported fuzzy membership type for '{name}': {kind}")


def resolve_fuzzy_space(fuzzy_space: dict[str, Any], baseline_values: dict[str, float]) -> dict[str, ResolvedFuzzyParameter]:
    raw_params = fuzzy_space.get("parameters", fuzzy_space)
    if not isinstance(raw_params, dict) or not raw_params:
        raise ValueError("Fuzzy space must contain a non-empty 'parameters' object")

    out: dict[str, ResolvedFuzzyParameter] = {}
    for name, raw_spec in raw_params.items():
        if not isinstance(raw_spec, dict):
            raise ValueError(f"Fuzzy parameter '{name}' must be an object")
        out[name] = resolve_parameter(name, raw_spec, baseline_values)
    return out


def build_alpha_intervals(
    parameters: dict[str, ResolvedFuzzyParameter],
    alpha_levels: list[float],
) -> dict[str, Any]:
    out: dict[str, Any] = {"alpha_levels": alpha_levels, "parameters": {}}
    for alpha in alpha_levels:
        if alpha < 0.0 or alpha > 1.0:
            raise ValueError(f"alpha level must be within [0, 1], got {alpha}")
    for name, param in parameters.items():
        out["parameters"][name] = {
            "resolved": param.to_dict(),
            "alpha_cuts": [
                {
                    "alpha": alpha,
                    "lower": param.alpha_interval(alpha)[0],
                    "upper": param.alpha_interval(alpha)[1],
                }
                for alpha in alpha_levels
            ],
        }
    return out


def parse_args() -> argparse.Namespace:
    ap = argparse.ArgumentParser(description="Resolve fuzzy SWMM parameter membership functions into alpha-cut intervals.")
    ap.add_argument("--base-inp", required=True, type=Path)
    ap.add_argument("--patch-map", required=True, type=Path)
    ap.add_argument("--fuzzy-space", required=True, type=Path)
    ap.add_argument("--alpha-levels", default="0,0.25,0.5,0.75,1")
    ap.add_argument("--out-resolved", required=True, type=Path)
    ap.add_argument("--out-alpha-intervals", required=True, type=Path)
    return ap.parse_args()


def main() -> None:
    args = parse_args()
    patch_map = load_json(args.patch_map)
    fuzzy_space = load_json(args.fuzzy_space)
    alpha_levels = [float(item.strip()) for item in args.alpha_levels.split(",") if item.strip()]
    baselines = read_baseline_values(args.base_inp, patch_map)
    resolved = resolve_fuzzy_space(fuzzy_space, baselines)
    write_json(args.out_resolved, {"parameters": {name: p.to_dict() for name, p in resolved.items()}})
    write_json(args.out_alpha_intervals, build_alpha_intervals(resolved, alpha_levels))


if __name__ == "__main__":
    main()

#!/usr/bin/env python3
"""INP-aware parameter recommender for Monte Carlo priors (issue #52).

The recommender answers: "given this INP, which parameters should I put
in my prior parameter_space.json?". The answer has three parts:

* ``core_required`` — a hardcoded 6-element list of the SWMM-sensitive
  parameters that should always be perturbed regardless of what the INP
  contains (N-Imperv, S-Imperv, Pct-Imperv, Width, MaxRate, MinRate).
  Per #52 this is the lock-in list; do not change it without revisiting
  the PRD.
* ``recommended`` — ``core_required`` plus extras detected from the INP
  (e.g. Decay if HORTON, Suction/K/IMD if GREEN_AMPT, Slope if there's
  any non-trivial spread). Always a superset of ``core_required``.
* ``rationale`` — a ``{param: prose}`` map. Every parameter in
  ``recommended`` that is **not** in ``core_required`` carries a
  non-empty rationale string so the modeller knows why each extra was
  added. Core parameters get a rationale too when there's a useful
  evidence boundary to call out (e.g. baseline range derived from
  default ±20%), but the contract only requires non-empty rationale
  for extras.

Detection is intentionally simple and authoritative: the
``[OPTIONS]`` block carries an ``INFILTRATION`` keyword which is the
SWMM5-defined source of truth. We honor that keyword verbatim and only
fall back to the column count of the ``[INFILTRATION]`` body when
``[OPTIONS]`` is missing or unparseable — that fallback is enough for
the unit tests but is not a substitute for a well-formed INP.

CLI::

    python parameter_recommender.py --inp <path>

emits the structured object on stdout as JSON. Stdout is the contract;
stderr is reserved for warnings.
"""

from __future__ import annotations

import argparse
import json
import sys
from pathlib import Path
from typing import Any


# The six-parameter SWMM-sensitive core. Per the issue spec this list
# is fixed regardless of the INP — even for Green-Ampt INPs, the
# Horton-tinted parameters MaxRate/MinRate stay in the core. The
# infiltration-method detection drives the *extras* in `recommended`,
# not the contents of `core_required`.
CORE_REQUIRED: tuple[str, ...] = (
    "N-Imperv",
    "S-Imperv",
    "Pct-Imperv",
    "Width",
    "MaxRate",
    "MinRate",
)


# Per-method extras. The intersection between these and CORE_REQUIRED
# (MaxRate / MinRate for Horton) is allowed and silently deduplicated
# when we compose `recommended`.
INFILTRATION_EXTRAS: dict[str, tuple[str, ...]] = {
    "horton": ("MaxRate", "MinRate", "Decay"),
    "modified_horton": ("MaxRate", "MinRate", "Decay"),
    "green_ampt": ("Suction", "K", "IMD"),
    "modified_green_ampt": ("Suction", "K", "IMD"),
    "curve_number": ("CurveNum", "Ksat", "DryTime"),
}


# Static rationale fragments. These describe *why* each extra is in the
# recommended set; the prose stays terse because the broader hydrology
# rationale lives in docs/hitl-thresholds.md and PRD-Z.
#
# PRD-GF-CORE: as of the gap-fill refactor, citations on a per-parameter
# basis are loaded from ``defaults_table.yaml`` (single source of truth
# across the runtime). The static templates below remain as fallbacks
# for parameters that do not have a defaults_table entry yet.
RATIONALE_TEMPLATES: dict[str, str] = {
    "Decay": "Horton infiltration detected in [INFILTRATION] section.",
    "Suction": "Green-Ampt infiltration detected in [INFILTRATION] section.",
    "K": "Green-Ampt infiltration detected in [INFILTRATION] section.",
    "IMD": "Green-Ampt infiltration detected in [INFILTRATION] section.",
    "CurveNum": "SCS Curve-Number infiltration detected in [INFILTRATION] section.",
    "Ksat": "Curve-Number infiltration detected in [INFILTRATION] section.",
    "DryTime": "Curve-Number infiltration detected in [INFILTRATION] section.",
    "Slope": "Default range [-20%, +20%] of measured value.",
    "MaxRate": "Horton infiltration detected in [INFILTRATION] section.",
    "MinRate": "Horton infiltration detected in [INFILTRATION] section.",
}


# PRD-GF-CORE: alias map from SWMM-canonical parameter names used in
# this recommender to the registry entry names in
# ``defaults_table.yaml``. Parallel to the alias map in
# ``agentic_swmm.gap_fill.proposer`` but local to this script — the
# recommender is shipped as a standalone CLI under skills/, so we
# cannot import the agentic_swmm package.
_DEFAULTS_TABLE_ALIASES: dict[str, str] = {
    "MaxRate": "horton_max_infiltration_rate",
    "MinRate": "horton_min_infiltration_rate",
    "Decay": "horton_decay_constant",
}


def _defaults_table_path() -> Path:
    """Resolve the project-root ``defaults_table.yaml``.

    The recommender script lives at
    ``skills/swmm-uncertainty/scripts/parameter_recommender.py``;
    the table is at the repo root three levels up. Tests can
    override via ``AISWMM_DEFAULTS_TABLE``.
    """
    import os

    override = os.environ.get("AISWMM_DEFAULTS_TABLE")
    if override:
        return Path(override)
    return Path(__file__).resolve().parents[3] / "defaults_table.yaml"


def _load_defaults_table() -> dict[str, dict[str, Any]]:
    """Return the ``entries`` map of the defaults table, or empty.

    Missing file / missing PyYAML / malformed YAML all map to an
    empty dict. The recommender then falls back to the static
    ``RATIONALE_TEMPLATES`` so behaviour is preserved when the table
    is unavailable.
    """
    path = _defaults_table_path()
    if not path.is_file():
        return {}
    try:
        import yaml
    except ImportError:  # pragma: no cover - defensive
        return {}
    try:
        payload = yaml.safe_load(path.read_text(encoding="utf-8"))
    except (OSError, yaml.YAMLError):
        return {}
    if not isinstance(payload, dict):
        return {}
    entries = payload.get("entries")
    if not isinstance(entries, dict):
        return {}
    return {str(k): dict(v) for k, v in entries.items() if isinstance(v, dict)}


def _rationale_for(param: str, method: str, table: dict[str, dict[str, Any]]) -> str:
    """Return the rationale string for ``param``.

    Lookup order:

    1. ``defaults_table.yaml`` entry pointed at by the alias map —
       returns ``"<reason from static template> Source: <citation>"``
       so the modeller sees both the trigger and the literature ref.
    2. Static ``RATIONALE_TEMPLATES`` entry — preserves the pre-PRD
       behaviour for parameters we do not have a defaults entry for.
    3. Generic ``f"Detected from [INFILTRATION] section ({method})"``
       fallback.
    """
    base = RATIONALE_TEMPLATES.get(param, "")
    entry_name = _DEFAULTS_TABLE_ALIASES.get(param)
    if entry_name and entry_name in table:
        citation = table[entry_name].get("source")
        if citation:
            if base:
                return f"{base} Source: {citation}"
            return f"Source: {citation}"
    if base:
        return base
    return f"Detected from [INFILTRATION] section ({method})."


def _normalise_method(token: str) -> str:
    """Map a SWMM ``INFILTRATION`` token to our snake_case identifier."""

    cleaned = token.strip().upper()
    table = {
        "HORTON": "horton",
        "MODIFIED_HORTON": "modified_horton",
        "GREEN_AMPT": "green_ampt",
        "MODIFIED_GREEN_AMPT": "modified_green_ampt",
        "CURVE_NUMBER": "curve_number",
    }
    return table.get(cleaned, cleaned.lower())


def _read_sections(inp_text: str) -> dict[str, list[str]]:
    """Slice an INP into a ``{SECTION_NAME: [lines]}`` map.

    Lines inside a section keep their original whitespace so downstream
    parsers can split by columns; comment-only lines (``;;``) are kept
    as-is because some [INFILTRATION] headers carry parameter names in
    a comment row, which the fallback detector consults.
    """

    sections: dict[str, list[str]] = {}
    current: str | None = None
    for raw in inp_text.splitlines():
        stripped = raw.strip()
        if stripped.startswith("[") and stripped.endswith("]"):
            current = stripped[1:-1].upper()
            sections.setdefault(current, [])
            continue
        if current is None:
            continue
        sections[current].append(raw)
    return sections


def _detect_method(sections: dict[str, list[str]]) -> str:
    """Return the snake_case infiltration method or ``"unknown"``.

    Primary source: ``[OPTIONS]`` ``INFILTRATION <token>`` line. This is
    the SWMM-canonical place to declare the method and the only one
    used by ``build_swmm_inp.py``. Fallback (rare): inspect the
    ``[INFILTRATION]`` comment header for parameter names.
    """

    options = sections.get("OPTIONS", [])
    for raw in options:
        stripped = raw.strip()
        if not stripped or stripped.startswith(";"):
            continue
        parts = stripped.split()
        if len(parts) >= 2 and parts[0].upper() == "INFILTRATION":
            return _normalise_method(parts[1])

    # Fallback: peek at the comment row of [INFILTRATION].
    body = sections.get("INFILTRATION", [])
    for raw in body:
        s = raw.strip()
        if not s.startswith(";"):
            continue
        upper = s.upper()
        if "SUCTION" in upper or "KSAT" in upper or "IMD" in upper:
            return "green_ampt"
        if "MAXRATE" in upper or "DECAY" in upper:
            return "horton"
        if "CURVENUM" in upper or "CURVE_NUM" in upper:
            return "curve_number"
    return "unknown"


def recommend(inp_path: Path) -> dict[str, Any]:
    """Return the structured recommender payload for ``inp_path``.

    The dict shape is::

        {
          "core_required":         list[str],
          "recommended":           list[str],
          "rationale":             dict[str, str],
          "infiltration_method":   str,
        }

    ``recommended`` preserves insertion order: ``core_required`` first,
    then the method-specific extras, then any always-on extras (Slope).
    """

    inp_text = Path(inp_path).read_text(encoding="utf-8", errors="ignore")
    sections = _read_sections(inp_text)
    method = _detect_method(sections)

    recommended: list[str] = list(CORE_REQUIRED)
    extras = list(INFILTRATION_EXTRAS.get(method, ()))
    # Always include Slope as a default-range extra. The static rationale
    # explains that the prior is heuristic (±20%) rather than measured.
    if "Slope" not in extras:
        extras.append("Slope")
    for extra in extras:
        if extra not in recommended:
            recommended.append(extra)

    rationale: dict[str, str] = {}
    # Every parameter in `recommended` that is not in `core_required`
    # must have a rationale (the test asserts this directly). We also
    # provide rationale entries for core parameters where there is a
    # useful evidence-boundary note (e.g. MaxRate/MinRate for Horton).
    #
    # PRD-GF-CORE: each rationale string is enriched with the citation
    # from defaults_table.yaml when available. The yaml lookup is
    # one-shot — we read once and pass the dict to `_rationale_for`.
    defaults = _load_defaults_table()
    core_set = set(CORE_REQUIRED)
    for param in recommended:
        if param in core_set:
            # Only fill rationale for core params that match the
            # detected method, to keep the output noise-free.
            if param in ("MaxRate", "MinRate") and method.endswith("horton"):
                rationale[param] = _rationale_for(param, method, defaults)
            continue
        rationale[param] = _rationale_for(param, method, defaults)

    return {
        "core_required": list(CORE_REQUIRED),
        "recommended": recommended,
        "rationale": rationale,
        "infiltration_method": method,
    }


def _build_argparser() -> argparse.ArgumentParser:
    ap = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    ap.add_argument(
        "--inp",
        required=True,
        type=Path,
        help="Path to the SWMM INP file to inspect.",
    )
    return ap


def main() -> None:
    args = _build_argparser().parse_args()
    payload = recommend(args.inp)
    json.dump(payload, sys.stdout, indent=2)
    sys.stdout.write("\n")


if __name__ == "__main__":
    main()

#!/usr/bin/env python3
"""Rainfall ensemble generator (issue #51, slice 5).

Two methods, one CLI:

  Method A — `perturbation`
      Take one observed rainfall timeseries and synthesise N noisy
      realisations from it. Four noise models are supported:
          * gaussian_iid          — additive zero-mean Gaussian
          * multiplicative        — log-normal-style scalar(s) * pattern
          * autocorrelated        — AR(1) noise with configurable phi
          * intensity_scaling     — sigma proportional to intensity (peaks
                                    vary more than troughs)
      Optional `preserve_total_volume=True` rescales each realisation so
      that the integrated rainfall matches the observed total.

  Method B — `idf`
      Sample IDF parameters `(a, b, c)` from their confidence intervals
      and synthesise a design hyetograph for each draw. Three storm types
      are supported:
          * chicago    — Keifer-Chu (1957)
          * huff       — 4 quartiles, Huff (1967), default 1st quartile
          * scs_type_ii — SCS 24-hr Type II canonical mass curve

Outputs:
  * `runs/<case>/09_audit/rainfall_realisations/realisation_<NNN>.csv`
  * `runs/<case>/09_audit/rainfall_ensemble_summary.json`

If a `--base-inp` + `--patch-rainfall-series` pair is supplied, each
realisation is patched into the [TIMESERIES] block of a copy of the base
INP and run through swmm5; peak flow / total volume are aggregated in the
summary.

The Python entry points are also exposed as importable functions so unit
tests can exercise them without invoking swmm5.
"""
from __future__ import annotations

import argparse
import csv
import json
import math
import re
import shutil
import subprocess
import sys
import time
from dataclasses import dataclass
from datetime import datetime, timedelta, timezone
from pathlib import Path
from typing import Any, Iterable, Sequence

import numpy as np


SCRIPT_DIR = Path(__file__).resolve().parent
REPO_ROOT = SCRIPT_DIR.parents[2]


# ---------------------------------------------------------------------------
# IO helpers
# ---------------------------------------------------------------------------


def load_json(path: Path) -> Any:
    return json.loads(Path(path).read_text(encoding="utf-8"))


def write_json(path: Path, obj: Any) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_text(json.dumps(obj, indent=2, sort_keys=True), encoding="utf-8")


def utc_now() -> str:
    return datetime.now(timezone.utc).isoformat(timespec="seconds")


@dataclass
class RainfallSeries:
    """A rainfall timeseries with timestamps and values (mm/hr)."""

    timestamps: list[datetime]
    values: np.ndarray

    @property
    def interval_minutes(self) -> int:
        if len(self.timestamps) < 2:
            return 0
        deltas = [
            int((self.timestamps[i + 1] - self.timestamps[i]).total_seconds() / 60.0)
            for i in range(len(self.timestamps) - 1)
        ]
        if not deltas:
            return 0
        # Use the modal delta
        return max(set(deltas), key=deltas.count)


# ---------------------------------------------------------------------------
# Rainfall series IO — CSV + SWMM .dat
# ---------------------------------------------------------------------------


_CSV_TIMESTAMP_FORMATS = (
    "%Y-%m-%d %H:%M:%S",
    "%Y-%m-%d %H:%M",
    "%Y-%m-%dT%H:%M:%S",
    "%Y-%m-%dT%H:%M",
    "%Y/%m/%d %H:%M:%S",
    "%Y/%m/%d %H:%M",
    "%m/%d/%Y %H:%M:%S",
    "%m/%d/%Y %H:%M",
)


def _parse_timestamp(value: str) -> datetime:
    s = value.strip()
    for fmt in _CSV_TIMESTAMP_FORMATS:
        try:
            return datetime.strptime(s, fmt)
        except ValueError:
            continue
    raise ValueError(f"unsupported timestamp format: '{value}'")


def read_rainfall_series(path: Path) -> RainfallSeries:
    """Read a rainfall timeseries from CSV or SWMM .dat.

    CSV: header row with `timestamp` (or `time`/`date`) and a rainfall
    column. Any column that is not the timestamp is taken as the value.

    SWMM .dat: SWMM external rainfall file with columns
        gauge year month day hour minute value
    """
    p = Path(path)
    suffix = p.suffix.lower()
    if suffix in {".csv", ".tsv"}:
        return _read_csv(p)
    if suffix in {".dat", ".txt"}:
        # try CSV first if it looks like a header
        sample = p.read_text(encoding="utf-8", errors="ignore").splitlines()[:3]
        if sample and ("," in sample[0] or "timestamp" in sample[0].lower()):
            return _read_csv(p)
        return _read_swmm_dat(p)
    # default: try CSV
    return _read_csv(p)


def _read_csv(path: Path) -> RainfallSeries:
    timestamps: list[datetime] = []
    values: list[float] = []
    with path.open("r", encoding="utf-8") as fh:
        reader = csv.reader(fh)
        rows = list(reader)
    if not rows:
        raise ValueError(f"empty rainfall CSV: {path}")
    header = [c.strip() for c in rows[0]]
    lower = [c.lower() for c in header]
    # Column inference walks CANDIDATE PRIORITY, not file order, and a
    # split Date+Time pair is combined into one timestamp (same defect
    # class as the calibration obs_reader, found 2026-08-08). The old
    # first-in-file scan also let the OTHER time column be chosen as
    # the value column for "Date,Time,Rainfall" headers.
    name_to_idx: dict[str, int] = {}
    for i, c in enumerate(lower):
        name_to_idx.setdefault(c, i)
    combined_idx = next(
        (name_to_idx[n] for n in ("timestamp", "datetime", "date_time") if n in name_to_idx),
        None,
    )
    date_idx = name_to_idx.get("date")
    time_idx = name_to_idx.get("time")
    ts_idx: int | None = None
    ts_pair: tuple[int, int] | None = None
    if combined_idx is not None:
        ts_idx = combined_idx
    elif date_idx is not None and time_idx is not None:
        ts_pair = (date_idx, time_idx)
    elif time_idx is not None:
        ts_idx = time_idx
    elif date_idx is not None:
        ts_idx = date_idx
    if ts_idx is None and ts_pair is None:
        # assume column 0 is timestamp, column 1 is value
        ts_idx, val_idx = 0, 1
    else:
        time_columns = {
            i for i in (combined_idx, date_idx, time_idx) if i is not None
        }
        val_idx = next(
            (i for i in range(len(header)) if i not in time_columns),
            1,
        )
    for row in rows[1:]:
        if not row or all(not c.strip() for c in row):
            continue
        try:
            if ts_pair is not None:
                d_idx, t_idx = ts_pair
                ts = _parse_timestamp(f"{row[d_idx].strip()} {row[t_idx].strip()}")
            else:
                ts = _parse_timestamp(row[ts_idx])
        except (IndexError, ValueError) as exc:
            raise ValueError(f"row {row}: cannot parse timestamp ({exc})") from exc
        try:
            val = float(row[val_idx])
        except (IndexError, ValueError) as exc:
            raise ValueError(f"row {row}: cannot parse value ({exc})") from exc
        timestamps.append(ts)
        values.append(val)
    return RainfallSeries(timestamps=timestamps, values=np.array(values, dtype=float))


def _read_swmm_dat(path: Path) -> RainfallSeries:
    timestamps: list[datetime] = []
    values: list[float] = []
    with path.open("r", encoding="utf-8") as fh:
        for raw in fh:
            line = raw.strip()
            if not line or line.startswith(";"):
                continue
            parts = line.split()
            if len(parts) < 7:
                continue
            try:
                year, month, day, hour, minute = (int(x) for x in parts[1:6])
                val = float(parts[6])
            except ValueError:
                continue
            timestamps.append(datetime(year, month, day, hour, minute))
            values.append(val)
    if not timestamps:
        raise ValueError(f"empty or malformed SWMM .dat: {path}")
    return RainfallSeries(timestamps=timestamps, values=np.array(values, dtype=float))


def write_rainfall_csv(path: Path, series: RainfallSeries) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    with path.open("w", encoding="utf-8", newline="") as fh:
        writer = csv.writer(fh)
        writer.writerow(["timestamp", "rainfall_mm_per_hr"])
        for ts, val in zip(series.timestamps, series.values):
            writer.writerow([ts.strftime("%Y-%m-%d %H:%M:%S"), float(val)])


# ---------------------------------------------------------------------------
# Method A — Perturbation
# ---------------------------------------------------------------------------


PERTURBATION_MODELS = ("gaussian_iid", "multiplicative", "autocorrelated", "intensity_scaling")


def _ensure_rng(rng: np.random.Generator | int | None) -> np.random.Generator:
    if rng is None:
        return np.random.default_rng()
    if isinstance(rng, np.random.Generator):
        return rng
    return np.random.default_rng(int(rng))


def perturb_series(
    *,
    observed: Sequence[float] | np.ndarray,
    config: dict[str, Any],
    n_realisations: int,
    rng: np.random.Generator | int | None = None,
) -> np.ndarray:
    """Return a `(n_realisations, len(observed))` matrix of perturbed series.

    `config` keys:
        model: one of PERTURBATION_MODELS
        sigma: noise scale
        ar1_coefficient: only for `autocorrelated`
        preserve_total_volume: bool (default False)

    All realisations are clipped at zero — rainfall is non-negative.
    If `preserve_total_volume=True`, each realisation is rescaled so its
    sum matches the observed total (when the observed total is positive).
    """
    rng = _ensure_rng(rng)
    if n_realisations <= 0:
        raise ValueError("n_realisations must be >= 1")
    observed = np.asarray(observed, dtype=float)
    if observed.ndim != 1:
        raise ValueError("observed must be a 1-D series")

    model = str(config.get("model", "")).lower()
    if model not in PERTURBATION_MODELS:
        raise ValueError(
            f"unknown perturbation model '{config.get('model')}'. "
            f"Supported: {', '.join(PERTURBATION_MODELS)}"
        )
    sigma = float(config.get("sigma", 0.1))
    preserve_total = bool(config.get("preserve_total_volume", False))

    n = len(observed)
    realisations = np.empty((n_realisations, n), dtype=float)

    if model == "gaussian_iid":
        noise = rng.normal(loc=0.0, scale=sigma, size=(n_realisations, n))
        realisations[:] = observed + noise

    elif model == "multiplicative":
        # log-normal multipliers per timestep, mean 1, std ~ sigma
        # use log-normal with mu = -sigma^2/2 so that E[exp(noise)] = 1
        log_mean = -0.5 * (sigma ** 2)
        log_noise = rng.normal(loc=log_mean, scale=sigma, size=(n_realisations, n))
        multipliers = np.exp(log_noise)
        realisations[:] = observed * multipliers

    elif model == "autocorrelated":
        phi = float(config.get("ar1_coefficient", 0.5))
        if not -1.0 < phi < 1.0:
            raise ValueError("ar1_coefficient must be in (-1, 1)")
        # innovation variance to keep stationary AR(1) variance = sigma^2
        innov_std = sigma * math.sqrt(1.0 - phi ** 2)
        innov = rng.normal(loc=0.0, scale=innov_std, size=(n_realisations, n))
        # initialise with stationary sigma
        noise = np.empty_like(innov)
        noise[:, 0] = rng.normal(loc=0.0, scale=sigma, size=n_realisations)
        for t in range(1, n):
            noise[:, t] = phi * noise[:, t - 1] + innov[:, t]
        realisations[:] = observed + noise

    elif model == "intensity_scaling":
        # variance proportional to intensity ^ 2 (relative noise)
        # noise_t ~ Normal(0, sigma * intensity_t)
        scales = sigma * observed
        # Avoid zero everywhere if observed has zeros — keep a small floor
        # so that low-intensity steps don't have literally zero spread but
        # the peaks still dominate.
        floor = sigma * max(float(observed.max()), 1e-9) * 0.01
        scales = np.maximum(scales, floor)
        noise = rng.normal(loc=0.0, scale=1.0, size=(n_realisations, n)) * scales
        realisations[:] = observed + noise

    # Non-negativity floor
    realisations = np.maximum(realisations, 0.0)

    if preserve_total:
        observed_total = float(observed.sum())
        if observed_total > 0.0:
            row_totals = realisations.sum(axis=1, keepdims=True)
            # rescale rows that have a non-zero sum
            safe = np.where(row_totals > 0.0, row_totals, 1.0)
            realisations = realisations * (observed_total / safe)

    return realisations


# ---------------------------------------------------------------------------
# Method B — IDF design storm hyetographs
# ---------------------------------------------------------------------------


DESIGN_STORM_TYPES = ("chicago", "huff", "scs_type_ii")


def idf_intensity_mm_per_hr(duration_minutes: float, a: float, b: float, c: float) -> float:
    """IDF curve i = a / (duration_hours + b)^c, returns mm/hr.

    Duration is in minutes -> we use hours for the canonical IDF form so that
    `a` has units of mm/hr.
    """
    d_hr = max(duration_minutes / 60.0, 1e-9)
    return float(a / ((d_hr + b) ** c))


def synthesise_design_hyetograph(
    *,
    storm_type: str,
    duration_minutes: int,
    interval_minutes: int,
    a: float,
    b: float,
    c: float,
    huff_quartile: int = 1,
    chicago_peak_position: float = 0.4,
    start_time: datetime | None = None,
) -> RainfallSeries:
    """Build a hyetograph for the given design storm type.

    All hyetographs are returned in mm/hr at `interval_minutes` resolution.
    """
    storm_type = storm_type.lower()
    if storm_type not in DESIGN_STORM_TYPES:
        raise ValueError(
            f"unsupported storm_type '{storm_type}'. Supported: {', '.join(DESIGN_STORM_TYPES)}"
        )
    if duration_minutes <= 0 or interval_minutes <= 0:
        raise ValueError("duration and interval must be positive")
    if duration_minutes % interval_minutes != 0:
        raise ValueError("duration_minutes must be a multiple of interval_minutes")

    n_steps = duration_minutes // interval_minutes
    # Reference total rainfall depth (mm) implied by the IDF curve for the
    # full design duration.
    total_intensity = idf_intensity_mm_per_hr(duration_minutes, a, b, c)
    total_depth_mm = total_intensity * (duration_minutes / 60.0)

    if storm_type == "chicago":
        intensities = _chicago_hyetograph(
            n_steps=n_steps,
            interval_minutes=interval_minutes,
            a=a,
            b=b,
            c=c,
            peak_position=chicago_peak_position,
        )
    elif storm_type == "huff":
        intensities = _huff_hyetograph(
            n_steps=n_steps,
            interval_minutes=interval_minutes,
            total_depth_mm=total_depth_mm,
            quartile=huff_quartile,
        )
    else:  # scs_type_ii
        intensities = _scs_type_ii_hyetograph(
            n_steps=n_steps,
            interval_minutes=interval_minutes,
            total_depth_mm=total_depth_mm,
        )

    # Build timestamps
    t0 = start_time or datetime(2024, 1, 1, 0, 0, 0)
    timestamps = [t0 + timedelta(minutes=interval_minutes * i) for i in range(n_steps)]
    return RainfallSeries(timestamps=timestamps, values=intensities)


def _chicago_hyetograph(
    *,
    n_steps: int,
    interval_minutes: int,
    a: float,
    b: float,
    c: float,
    peak_position: float,
) -> np.ndarray:
    """Keifer-Chu Chicago hyetograph (1957).

    Average intensity over a duration `d` from the IDF curve:
        i_avg(d) = a / (d_hr + b)^c
    Total depth for duration d:  P(d) = i_avg(d) * d_hr
    Average intensity left of peak (over a duration of t_b):
        i_left = a*( (1-c) * t_b/r + b ) / ( t_b/r + b )^(c+1)
    Similarly right of peak.

    Here we use the canonical instantaneous intensity formula (derivative
    of P with respect to duration):
        i(d) = a * ( (1-c)*d_hr + b ) / (d_hr + b)^(c+1)
    placed symmetrically around the peak at fraction `peak_position`.
    """
    if not 0.0 < peak_position < 1.0:
        raise ValueError("chicago peak_position must be in (0, 1)")
    duration_hr = (n_steps * interval_minutes) / 60.0
    t_peak_hr = peak_position * duration_hr
    intensities = np.zeros(n_steps, dtype=float)
    for i in range(n_steps):
        # midpoint of bin i in hours
        t_mid_hr = ((i + 0.5) * interval_minutes) / 60.0
        # offset from the peak in hours (positive)
        dt = abs(t_mid_hr - t_peak_hr)
        # Convert offset to duration relative to the appropriate side
        if t_mid_hr < t_peak_hr:
            d_eff = dt / peak_position
        else:
            d_eff = dt / (1.0 - peak_position)
        d_eff = max(d_eff, 1e-6)
        intensities[i] = a * ((1.0 - c) * d_eff + b) / ((d_eff + b) ** (c + 1.0))
    return intensities


# Huff cumulative distributions (fraction of total rainfall vs fraction of
# total time) for each quartile, from Huff (1967). 21 points (0..1 in 0.05
# steps) gives enough resolution for hyetograph interpolation. Each
# quartile peaks (steepest slope) within its named quartile of duration.
_HUFF_CUMULATIVE = {
    1: (
        0.000, 0.032, 0.110, 0.245, 0.420, 0.580, 0.700, 0.770, 0.815,
        0.840, 0.860, 0.876, 0.890, 0.903, 0.915, 0.926, 0.938, 0.952,
        0.968, 0.985, 1.000,
    ),
    2: (
        0.000, 0.020, 0.060, 0.110, 0.180, 0.270, 0.400, 0.570, 0.720,
        0.820, 0.880, 0.910, 0.930, 0.945, 0.958, 0.968, 0.978, 0.986,
        0.992, 0.997, 1.000,
    ),
    3: (
        0.000, 0.015, 0.040, 0.070, 0.105, 0.145, 0.195, 0.260, 0.345,
        0.460, 0.610, 0.760, 0.860, 0.910, 0.940, 0.960, 0.975, 0.985,
        0.992, 0.997, 1.000,
    ),
    4: (
        0.000, 0.010, 0.025, 0.045, 0.070, 0.100, 0.135, 0.175, 0.220,
        0.270, 0.325, 0.385, 0.450, 0.520, 0.595, 0.680, 0.775, 0.875,
        0.955, 0.990, 1.000,
    ),
}


def _huff_hyetograph(
    *,
    n_steps: int,
    interval_minutes: int,
    total_depth_mm: float,
    quartile: int,
) -> np.ndarray:
    if quartile not in _HUFF_CUMULATIVE:
        raise ValueError(f"huff_quartile must be one of {sorted(_HUFF_CUMULATIVE)}")
    cum = np.array(_HUFF_CUMULATIVE[quartile], dtype=float)
    # Sample cumulative fraction at each step's right edge
    fractions = np.linspace(0.0, 1.0, n_steps + 1)
    base = np.linspace(0.0, 1.0, len(cum))
    cum_depth = np.interp(fractions, base, cum) * total_depth_mm
    depths_per_step = np.diff(cum_depth)
    # convert depth_mm to mm/hr
    interval_hr = interval_minutes / 60.0
    intensities = depths_per_step / interval_hr
    return intensities


# SCS 24-hr Type II canonical cumulative-rainfall ratios (fraction of total
# rainfall at fraction of total duration). 25 points (every hour for 24 hrs).
_SCS_TYPE_II = (
    0.000, 0.011, 0.022, 0.034, 0.048, 0.063, 0.080, 0.098, 0.120, 0.147,
    0.181, 0.235, 0.663, 0.772, 0.820, 0.850, 0.880, 0.898, 0.916, 0.934,
    0.952, 0.964, 0.976, 0.988, 1.000,
)


def _scs_type_ii_hyetograph(
    *,
    n_steps: int,
    interval_minutes: int,
    total_depth_mm: float,
) -> np.ndarray:
    base_fractions = np.linspace(0.0, 1.0, len(_SCS_TYPE_II))
    cum_ratios = np.array(_SCS_TYPE_II, dtype=float)
    edges = np.linspace(0.0, 1.0, n_steps + 1)
    cum_depth = np.interp(edges, base_fractions, cum_ratios) * total_depth_mm
    depths_per_step = np.diff(cum_depth)
    interval_hr = interval_minutes / 60.0
    intensities = depths_per_step / interval_hr
    return intensities


def sample_idf_param(
    *, value: float, ci: tuple[float, float] | list[float], rng: np.random.Generator
) -> float:
    """Sample an IDF parameter from a normal distribution implied by its CI.

    The 95% CI -> sigma = (upper - lower) / (2 * 1.96).
    """
    lo, hi = float(ci[0]), float(ci[1])
    if hi < lo:
        lo, hi = hi, lo
    sigma = (hi - lo) / (2.0 * 1.959963984540054)
    return float(rng.normal(loc=value, scale=max(sigma, 1e-12)))


def build_idf_realisations(
    *,
    idf_config: dict[str, Any],
    n_realisations: int,
    rng: np.random.Generator | int | None = None,
) -> list[RainfallSeries]:
    """Build N hyetograph realisations by sampling IDF parameters.

    `idf_config` keys (matches issue spec):
        type:              chicago | huff | scs_type_ii
        duration_minutes:  int
        return_period_years: int (metadata)
        interval_minutes:  int (default 5)
        huff_quartile:     int (default 1, only for huff)
        chicago_peak_position: float (default 0.4)
        start_time:        ISO timestamp (optional)
        params:
            a: {value, ci: [lo, hi]}
            b: {value, ci: [lo, hi]}
            c: {value, ci: [lo, hi]}
    """
    rng = _ensure_rng(rng)
    storm_type = str(idf_config["type"]).lower()
    duration = int(idf_config["duration_minutes"])
    interval = int(idf_config.get("interval_minutes", 5))
    huff_quartile = int(idf_config.get("huff_quartile", 1))
    peak_position = float(idf_config.get("chicago_peak_position", 0.4))
    start_iso = idf_config.get("start_time")
    if start_iso:
        start_time = datetime.fromisoformat(str(start_iso))
    else:
        start_time = datetime(2024, 1, 1, 0, 0, 0)
    params = idf_config["params"]
    out: list[RainfallSeries] = []
    for _ in range(n_realisations):
        a = sample_idf_param(value=params["a"]["value"], ci=params["a"]["ci"], rng=rng)
        b = sample_idf_param(value=params["b"]["value"], ci=params["b"]["ci"], rng=rng)
        c = sample_idf_param(value=params["c"]["value"], ci=params["c"]["ci"], rng=rng)
        # Guardrails: IDF requires a > 0 and (d_hr + b) > 0. Clamp.
        a = max(a, 1e-6)
        b = max(b, 0.0)
        c = max(c, 0.01)
        series = synthesise_design_hyetograph(
            storm_type=storm_type,
            duration_minutes=duration,
            interval_minutes=interval,
            a=a,
            b=b,
            c=c,
            huff_quartile=huff_quartile,
            chicago_peak_position=peak_position,
            start_time=start_time,
        )
        out.append(series)
    return out


# ---------------------------------------------------------------------------
# SWMM INP patching for ensemble runs
# ---------------------------------------------------------------------------


_SECTION_RE = re.compile(r"^\s*\[(?P<name>[A-Z_]+)\]\s*$")


def _split_sections(inp_text: str) -> list[tuple[str | None, list[str]]]:
    """Return ordered (section_name | None, lines) groups (preserves blanks)."""
    sections: list[tuple[str | None, list[str]]] = []
    current_name: str | None = None
    current_lines: list[str] = []
    for line in inp_text.splitlines():
        m = _SECTION_RE.match(line)
        if m:
            sections.append((current_name, current_lines))
            current_name = m.group("name")
            current_lines = [line]
        else:
            current_lines.append(line)
    sections.append((current_name, current_lines))
    return sections


def patch_rainfall_timeseries(
    *,
    base_inp_text: str,
    series_name: str,
    series: RainfallSeries,
) -> str:
    """Replace `series_name` rows inside the [TIMESERIES] block."""
    sections = _split_sections(base_inp_text)
    out_lines: list[str] = []
    for name, lines in sections:
        if name != "TIMESERIES":
            out_lines.extend(lines)
            continue
        kept: list[str] = []
        for line in lines:
            stripped = line.strip()
            # Keep header / comment lines untouched; drop existing rows
            # for this series.
            if not stripped or stripped.startswith(";") or stripped.startswith("["):
                kept.append(line)
                continue
            parts = stripped.split(None, 1)
            if parts and parts[0] == series_name:
                continue
            kept.append(line)
        # Append the realisation rows
        for ts, val in zip(series.timestamps, series.values):
            kept.append(
                f"{series_name:<18} {ts.strftime('%m/%d/%Y')} {ts.strftime('%H:%M')} {float(val):.6f}"
            )
        out_lines.extend(kept)
    return "\n".join(out_lines) + ("\n" if not base_inp_text.endswith("\n") else "")


def _parse_peak_total_from_rpt(rpt_path: Path, node: str) -> tuple[float | None, float | None]:
    """Extract (peak_flow_cms, total_volume_m3) for `node` from a SWMM .rpt.

    Reads two complementary tables:

    1. Node Inflow Summary — rows like
           ``O1   OUTFALL   0.000   3.366   2  10:28   0   46.3   0.000``
       cols: name type maxLat maxTotal days hr:min latVol totalVol err%

    2. Outfall Loading Summary — rows like
           ``O1   52.35   0.205   3.366   46.402``
       cols: name flowFreq avgFlow maxFlow totalVolume(10^6 ltr)

    Preference order: (a) Outfall Loading Summary for outfalls (max flow,
    total volume), then (b) Node Inflow Summary as a fallback.
    """
    if not rpt_path.exists():
        return (None, None)
    text = rpt_path.read_text(encoding="utf-8", errors="ignore").splitlines()
    peak: float | None = None
    total: float | None = None

    _terminator_titles = (
        "Link Flow Summary",
        "Flow Classification Summary",
        "Conduit Surcharge Summary",
        "Pumping Summary",
        "Storage Volume Summary",
        "Subcatchment Runoff Summary",
        "Node Surcharge Summary",
        "Node Flooding Summary",
        "Continuity Error",
        "Analysis begun",
    )

    def _is_terminator(stripped: str, current_title: str) -> bool:
        for other in _terminator_titles:
            if other == current_title:
                continue
            if other in stripped:
                return True
        return False

    # Pass 1: Outfall Loading Summary -> data row contains name, freq, avg,
    # max, total. We start scanning AFTER the title's trailing banner of
    # asterisks, so the "***" terminator only fires when we hit the next
    # section's banner.
    in_outfall = False
    rows_seen = 0
    for line in text:
        if "Outfall Loading Summary" in line:
            in_outfall = True
            rows_seen = 0
            continue
        if in_outfall:
            stripped = line.strip()
            if _is_terminator(stripped, "Outfall Loading Summary"):
                break
            if not stripped:
                continue
            if set(stripped) <= {"-", "*"}:
                # Decorative separator: dashes or trailing title banner.
                continue
            if stripped.startswith("Outfall Node") or stripped.startswith("Flow") or stripped.startswith("Freq") or stripped.startswith("Pcnt"):
                continue
            parts = stripped.split()
            if not parts:
                continue
            if parts[0] == "System":
                break
            if parts[0] == node and len(parts) >= 5:
                try:
                    peak = float(parts[3])  # Max flow CMS
                    total = float(parts[4]) * 1000.0  # 10^6 L -> m^3
                    return (peak, total)
                except (IndexError, ValueError):
                    pass
            rows_seen += 1

    # Pass 2: Node Inflow Summary (catches non-outfall nodes)
    in_inflow = False
    for line in text:
        if "Node Inflow Summary" in line:
            in_inflow = True
            continue
        if in_inflow:
            stripped = line.strip()
            if _is_terminator(stripped, "Node Inflow Summary"):
                break
            if not stripped:
                continue
            if set(stripped) <= {"-", "*"}:
                continue
            if any(
                stripped.startswith(p)
                for p in ("Maximum", "Lateral", "Inflow", "Node ", "Time of")
            ):
                continue
            parts = stripped.split()
            if not parts:
                continue
            if parts[0] == node and len(parts) >= 8:
                try:
                    peak = float(parts[3])
                except (IndexError, ValueError):
                    pass
                try:
                    total = float(parts[7]) * 1000.0
                except (IndexError, ValueError):
                    pass
                break
    return (peak, total)


# ---------------------------------------------------------------------------
# Orchestration
# ---------------------------------------------------------------------------


def _series_to_records(series: RainfallSeries) -> list[dict[str, Any]]:
    return [
        {"timestamp": ts.strftime("%Y-%m-%d %H:%M:%S"), "rainfall_mm_per_hr": float(v)}
        for ts, v in zip(series.timestamps, series.values)
    ]


def _summarise_realisation(series: RainfallSeries) -> dict[str, float]:
    values = np.asarray(series.values, dtype=float)
    interval_hr = max(series.interval_minutes, 1) / 60.0
    return {
        "peak_intensity_mm_per_hr": float(values.max() if values.size else 0.0),
        "total_volume_mm": float(values.sum() * interval_hr),
        "interval_minutes": int(series.interval_minutes),
        "n_steps": int(values.size),
    }


def run_swmm_for_realisation(
    *,
    base_inp_text: str,
    series_name: str,
    series: RainfallSeries,
    realisation_dir: Path,
    swmm_node: str,
) -> dict[str, Any]:
    """Patch + run swmm5 for a single realisation. Returns metrics + status."""
    realisation_dir.mkdir(parents=True, exist_ok=True)
    inp_path = realisation_dir / "model.inp"
    rpt_path = realisation_dir / "model.rpt"
    out_path = realisation_dir / "model.out"
    patched = patch_rainfall_timeseries(
        base_inp_text=base_inp_text,
        series_name=series_name,
        series=series,
    )
    inp_path.write_text(patched, encoding="utf-8")

    if not shutil.which("swmm5"):
        return {
            "status": "skipped",
            "reason": "swmm5 binary not on PATH",
            "files": {"inp": str(inp_path)},
        }

    proc = subprocess.run(
        ["swmm5", str(inp_path), str(rpt_path), str(out_path)],
        capture_output=True,
        text=True,
    )
    (realisation_dir / "stdout.txt").write_text(proc.stdout, encoding="utf-8", errors="ignore")
    (realisation_dir / "stderr.txt").write_text(proc.stderr, encoding="utf-8", errors="ignore")
    if proc.returncode != 0:
        return {
            "status": "failed",
            "reason": f"swmm5 returned {proc.returncode}",
            "files": {
                "inp": str(inp_path),
                "rpt": str(rpt_path),
                "stdout": str(realisation_dir / "stdout.txt"),
                "stderr": str(realisation_dir / "stderr.txt"),
            },
        }
    peak, total_vol = _parse_peak_total_from_rpt(rpt_path, swmm_node)
    return {
        "status": "ok",
        "metrics": {
            "peak_flow": peak,
            "total_volume_m3": total_vol,
        },
        "files": {
            "inp": str(inp_path),
            "rpt": str(rpt_path),
            "out": str(out_path),
        },
    }


def aggregate_metrics(values: Iterable[Any]) -> dict[str, Any]:
    nums = [float(v) for v in values if isinstance(v, (int, float))]
    if not nums:
        return {"count": 0, "min": None, "max": None, "mean": None, "p05": None, "p50": None, "p95": None}
    arr = np.array(nums, dtype=float)
    return {
        "count": len(nums),
        "min": float(arr.min()),
        "max": float(arr.max()),
        "mean": float(arr.mean()),
        "p05": float(np.percentile(arr, 5)),
        "p50": float(np.percentile(arr, 50)),
        "p95": float(np.percentile(arr, 95)),
    }


def generate_realisations(
    *,
    method: str,
    config: dict[str, Any],
    rng: np.random.Generator | int | None = None,
) -> tuple[list[RainfallSeries], dict[str, Any]]:
    """Return (realisations, controls) for the requested method.

    For `perturbation`, the realisations inherit timestamps from the
    observed input. For `idf`, timestamps are synthesised from
    `start_time` + `interval_minutes`.
    """
    rng = _ensure_rng(rng)
    method = method.lower()
    n = int(config.get("n_realisations", 100))
    if method == "perturbation":
        pert_cfg = config["perturbation"]
        path = Path(config["input_rainfall_path"])
        observed = read_rainfall_series(path)
        if observed.values.size == 0:
            raise ValueError(f"observed series is empty: {path}")
        matrix = perturb_series(
            observed=observed.values,
            config=pert_cfg,
            n_realisations=n,
            rng=rng,
        )
        realisations = [
            RainfallSeries(timestamps=list(observed.timestamps), values=row)
            for row in matrix
        ]
        controls = {
            "input_rainfall_path": str(path),
            "interval_minutes": observed.interval_minutes,
            "observed_n_steps": int(observed.values.size),
            "perturbation": pert_cfg,
        }
        return realisations, controls

    if method == "idf":
        idf_cfg = config["idf"]
        # honour `n_realisations` at the top level as the spec says, but
        # also accept it inside `idf` (the config example in the issue
        # places it under `idf`).
        n_real = int(config.get("n_realisations", idf_cfg.get("n_realisations", n)))
        realisations = build_idf_realisations(
            idf_config=idf_cfg,
            n_realisations=n_real,
            rng=rng,
        )
        controls = {
            "idf": idf_cfg,
        }
        return realisations, controls

    raise ValueError(f"unknown method '{method}' (expected: perturbation | idf)")


def run_ensemble(
    *,
    method: str,
    config: dict[str, Any],
    run_root: Path,
    base_inp: Path | None,
    series_name: str,
    swmm_node: str,
    seed: int,
    dry_run: bool,
) -> dict[str, Any]:
    """Top-level orchestration. Writes realisation CSVs + summary JSON."""
    started_at = utc_now()
    rng = np.random.default_rng(int(seed))
    audit_dir = run_root / "09_audit"
    realisations_dir = audit_dir / "rainfall_realisations"
    audit_dir.mkdir(parents=True, exist_ok=True)
    realisations_dir.mkdir(parents=True, exist_ok=True)

    realisations, controls = generate_realisations(method=method, config=config, rng=rng)

    # Persist each realisation as CSV + summarise
    per_realisation: list[dict[str, Any]] = []
    base_inp_text: str | None = None
    if base_inp is not None and not dry_run:
        base_inp_text = Path(base_inp).read_text(encoding="utf-8", errors="ignore")

    n_digits = max(len(str(len(realisations) - 1)), 3)
    for idx, series in enumerate(realisations):
        rel_name = f"realisation_{idx:0{n_digits}d}"
        csv_path = realisations_dir / f"{rel_name}.csv"
        write_rainfall_csv(csv_path, series)
        rec: dict[str, Any] = {
            "index": idx,
            "name": rel_name,
            "csv": str(csv_path),
            "summary": _summarise_realisation(series),
            "status": "csv_written",
        }
        if base_inp_text is not None:
            run_dir = audit_dir / "swmm_realisations" / rel_name
            swmm_result = run_swmm_for_realisation(
                base_inp_text=base_inp_text,
                series_name=series_name,
                series=series,
                realisation_dir=run_dir,
                swmm_node=swmm_node,
            )
            rec["swmm"] = swmm_result
            rec["status"] = swmm_result.get("status", "unknown")
        per_realisation.append(rec)

    swmm_metrics_present = [
        (r.get("swmm") or {}).get("metrics", {}) for r in per_realisation if r.get("swmm")
    ]
    payload: dict[str, Any] = {
        "schema": "swmm-uncertainty/rainfall-ensemble/v1",
        "method": method,
        "created_at_utc": started_at,
        "finished_at_utc": utc_now(),
        "seed": int(seed),
        "n_realisations": len(realisations),
        "controls": controls,
        "outputs": {
            "realisations_dir": str(realisations_dir),
            "summary_json": str(audit_dir / "rainfall_ensemble_summary.json"),
        },
        "rainfall_ensemble_stats": {
            "peak_intensity_mm_per_hr": aggregate_metrics(
                r["summary"]["peak_intensity_mm_per_hr"] for r in per_realisation
            ),
            "total_volume_mm": aggregate_metrics(
                r["summary"]["total_volume_mm"] for r in per_realisation
            ),
        },
        "swmm_ensemble_stats": {
            "peak_flow": aggregate_metrics(m.get("peak_flow") for m in swmm_metrics_present),
            "total_volume_m3": aggregate_metrics(m.get("total_volume_m3") for m in swmm_metrics_present),
        },
        "realisations": per_realisation,
    }
    write_json(audit_dir / "rainfall_ensemble_summary.json", payload)
    return payload


# ---------------------------------------------------------------------------
# CLI
# ---------------------------------------------------------------------------


def parse_args(argv: list[str] | None = None) -> argparse.Namespace:
    p = argparse.ArgumentParser(
        description="Generate a rainfall ensemble via perturbation of an observed series or IDF-based design storms.",
    )
    p.add_argument("--method", required=True, choices=("perturbation", "idf"))
    p.add_argument("--config", required=True, type=Path, help="JSON config (see SKILL.md).")
    p.add_argument("--run-root", required=True, type=Path)
    p.add_argument("--base-inp", type=Path, default=None, help="If given, each realisation is patched + run through swmm5.")
    p.add_argument("--series-name", default="TS_RAIN", help="Name of the SWMM [TIMESERIES] block to replace.")
    p.add_argument("--swmm-node", default="O1", help="Node to extract peak flow / total volume from.")
    p.add_argument("--seed", type=int, default=42)
    p.add_argument("--dry-run", action="store_true", help="Generate realisations + CSVs but skip swmm5.")
    return p.parse_args(argv)


def main(argv: list[str] | None = None) -> int:
    args = parse_args(argv)
    config = load_json(args.config)
    if not isinstance(config, dict):
        raise SystemExit("config JSON must be an object")

    # Allow --method to override config["method"] (and vice versa)
    method = str(args.method or config.get("method", "")).lower()

    payload = run_ensemble(
        method=method,
        config=config,
        run_root=Path(args.run_root),
        base_inp=args.base_inp,
        series_name=str(args.series_name),
        swmm_node=str(args.swmm_node),
        seed=int(args.seed),
        dry_run=bool(args.dry_run),
    )
    print(json.dumps({
        "method": payload["method"],
        "n_realisations": payload["n_realisations"],
        "summary_json": payload["outputs"]["summary_json"],
        "rainfall_ensemble_stats": payload["rainfall_ensemble_stats"],
        "swmm_ensemble_stats": payload["swmm_ensemble_stats"],
    }, indent=2, sort_keys=True))
    return 0


if __name__ == "__main__":
    raise SystemExit(main())

#!/usr/bin/env python3
from __future__ import annotations

import argparse
import itertools
import json
import random
from pathlib import Path
from typing import Any


def load_json(path: Path) -> Any:
    return json.loads(path.read_text(encoding="utf-8"))


def write_json(path: Path, obj: Any) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_text(json.dumps(obj, indent=2, sort_keys=True), encoding="utf-8")


def _format_value(value: float, resolved: dict[str, Any] | None, template: Any) -> float | int:
    resolved = resolved or {}
    value_type = str(resolved.get("value_type", "")).lower()
    if value_type in {"int", "integer"} or (not value_type and isinstance(template, int) and not isinstance(template, bool)):
        return int(round(value))
    precision = resolved.get("precision")
    if precision is not None:
        return round(float(value), int(precision))
    return float(value)


def sample_lhs(
    intervals: dict[str, tuple[Any, Any]],
    resolved: dict[str, dict[str, Any]],
    count: int,
    rng: random.Random,
) -> list[dict[str, float | int]]:
    if count <= 0:
        raise ValueError("samples_per_alpha must be >= 1")

    unit_vectors: dict[str, list[float]] = {}
    for name in intervals:
        vals = [(i + rng.random()) / count for i in range(count)]
        rng.shuffle(vals)
        unit_vectors[name] = vals

    samples: list[dict[str, float | int]] = []
    for sample_idx in range(count):
        params: dict[str, float | int] = {}
        for name, (lo, hi) in intervals.items():
            value = float(lo) + (float(hi) - float(lo)) * unit_vectors[name][sample_idx]
            params[name] = _format_value(value, resolved.get(name), lo)
        samples.append(params)
    return samples


def sample_random(
    intervals: dict[str, tuple[Any, Any]],
    resolved: dict[str, dict[str, Any]],
    count: int,
    rng: random.Random,
) -> list[dict[str, float | int]]:
    if count <= 0:
        raise ValueError("samples_per_alpha must be >= 1")
    samples: list[dict[str, float | int]] = []
    for _ in range(count):
        params: dict[str, float | int] = {}
        for name, (lo, hi) in intervals.items():
            value = rng.uniform(float(lo), float(hi))
            params[name] = _format_value(value, resolved.get(name), lo)
        samples.append(params)
    return samples


def sample_boundary(intervals: dict[str, tuple[Any, Any]]) -> list[dict[str, float | int]]:
    names = list(intervals)
    choices = [[intervals[name][0], intervals[name][1]] for name in names]
    samples: list[dict[str, float | int]] = []
    seen: set[tuple[tuple[str, Any], ...]] = set()
    for values in itertools.product(*choices):
        sample = {name: value for name, value in zip(names, values)}
        key = tuple(sorted(sample.items()))
        if key in seen:
            continue
        seen.add(key)
        samples.append(sample)
    return samples


def intervals_for_alpha(alpha_intervals: dict[str, Any], alpha: float) -> dict[str, tuple[Any, Any]]:
    intervals: dict[str, tuple[Any, Any]] = {}
    for name, record in alpha_intervals["parameters"].items():
        cuts = record.get("alpha_cuts") or []
        matching = [cut for cut in cuts if abs(float(cut["alpha"]) - alpha) < 1e-12]
        if not matching:
            raise KeyError(f"No alpha-cut for parameter '{name}' at alpha={alpha}")
        cut = matching[0]
        intervals[name] = (cut["lower"], cut["upper"])
    return intervals


def resolved_specs(alpha_intervals: dict[str, Any]) -> dict[str, dict[str, Any]]:
    return {
        name: record.get("resolved") or {}
        for name, record in alpha_intervals["parameters"].items()
    }


def generate_parameter_sets(
    alpha_intervals: dict[str, Any],
    *,
    method: str,
    samples_per_alpha: int,
    seed: int,
) -> list[dict[str, Any]]:
    rng = random.Random(seed)
    method = method.lower().strip()
    alpha_levels = [float(alpha) for alpha in alpha_intervals["alpha_levels"]]
    resolved = resolved_specs(alpha_intervals)
    trials: list[dict[str, Any]] = []
    trial_idx = 1

    for alpha in alpha_levels:
        intervals = intervals_for_alpha(alpha_intervals, alpha)
        if all(float(lo) == float(hi) for lo, hi in intervals.values()):
            samples = [{name: lo for name, (lo, _hi) in intervals.items()}]
        elif method == "lhs":
            samples = sample_lhs(intervals, resolved, samples_per_alpha, rng)
        elif method == "random":
            samples = sample_random(intervals, resolved, samples_per_alpha, rng)
        elif method == "boundary":
            samples = sample_boundary(intervals)
        else:
            raise ValueError(f"Unsupported sampling method: {method}")

        for local_idx, params in enumerate(samples, start=1):
            trials.append(
                {
                    "name": f"alpha_{alpha:.2f}_trial_{local_idx:03d}",
                    "params": params,
                    "metadata": {
                        "alpha": alpha,
                        "sample_index": local_idx,
                        "global_sample_index": trial_idx,
                        "sampling_method": method,
                    },
                }
            )
            trial_idx += 1

    return trials


def parse_args() -> argparse.Namespace:
    ap = argparse.ArgumentParser(description="Generate parameter sets from fuzzy alpha-cut intervals.")
    ap.add_argument("--alpha-intervals", required=True, type=Path)
    ap.add_argument("--method", default="lhs", choices=["lhs", "random", "boundary"])
    ap.add_argument("--samples-per-alpha", default=20, type=int)
    ap.add_argument("--seed", default=42, type=int)
    ap.add_argument("--out", required=True, type=Path)
    return ap.parse_args()


def main() -> None:
    args = parse_args()
    alpha_intervals = load_json(args.alpha_intervals)
    trials = generate_parameter_sets(
        alpha_intervals,
        method=args.method,
        samples_per_alpha=args.samples_per_alpha,
        seed=args.seed,
    )
    write_json(args.out, {"parameter_sets": trials})


if __name__ == "__main__":
    main()

#!/usr/bin/env python3
"""Unified sensitivity-analysis entry point for swmm-uncertainty.

Slice 4 (#49) ports `swmm-calibration/scripts/parameter_scout.py` to
`swmm-uncertainty/scripts/sensitivity.py` and extends it with two
variance-based methods backed by SALib:

    --method oat      One-at-a-time perturbation around a baseline; ranks
                      parameters by an RMSE+peak-error importance score.
                      Output mirrors the legacy parameter_scout summary so
                      downstream consumers (calibration scaffold, audit
                      run folder) keep working.
    --method morris   Morris elementary-effects (mu_star, sigma) at sample
                      budget r * (k + 1).
    --method sobol    Sobol' indices (first-order S_i, total-effect S_T_i)
                      at sample budget N * (2k + 2).

All three modes share an `inp_patch` workflow: a base INP, a patch-map
JSON describing where each named parameter sits in the INP, and an
observed series used to score each trial. The OAT branch consumes a
`scan_spec.json` (the old parameter_scout shape); the Morris and Sobol'
branches consume a `parameter_space.json` keyed by name with `min`/`max`.

Outputs are written to `--summary-json` (the issue spec puts this at
`runs/<case>/09_audit/sensitivity_indices.json`).

The scoring objective for Morris/Sobol' is RMSE between simulated and
observed flow at the target node. RMSE is monotone in deviation, so
larger output sensitivity directly maps to larger index values without
needing a sign convention.
"""

from __future__ import annotations

import argparse
import json
import math
import subprocess
import sys
from pathlib import Path
from typing import Any

import pandas as pd

# Reuse calibration helpers (inp_patch, obs_reader) so the move does not
# duplicate code. CONTEXT.md treats those as cross-skill primitives.
SCRIPT_DIR = Path(__file__).resolve().parent
REPO_ROOT = SCRIPT_DIR.parents[2]
CALIBRATION_SCRIPTS = REPO_ROOT / "skills" / "swmm-calibration" / "scripts"
if str(CALIBRATION_SCRIPTS) not in sys.path:
    sys.path.insert(0, str(CALIBRATION_SCRIPTS))

from inp_patch import patch_inp_text  # noqa: E402
from obs_reader import read_series  # noqa: E402

from swmmtoolbox import swmmtoolbox  # noqa: E402


# ---------------------------------------------------------------------------
# Shared utilities
# ---------------------------------------------------------------------------


def load_json(path: str | Path) -> Any:
    return json.loads(Path(path).read_text())


def filter_series_window(df: pd.DataFrame, start: str | None, end: str | None) -> pd.DataFrame:
    out = df.copy()
    out["timestamp"] = pd.to_datetime(out["timestamp"])
    if start:
        out = out[out["timestamp"] >= pd.Timestamp(start)]
    if end:
        out = out[out["timestamp"] <= pd.Timestamp(end)]
    return out.reset_index(drop=True)


def run_swmm(inp: Path, run_dir: Path) -> tuple[int, Path, Path]:
    run_dir.mkdir(parents=True, exist_ok=True)
    rpt = run_dir / "model.rpt"
    out = run_dir / "model.out"
    proc = subprocess.run(
        ["swmm5", str(inp), str(rpt), str(out)],
        capture_output=True,
        text=True,
    )
    (run_dir / "stdout.txt").write_text(proc.stdout, encoding="utf-8", errors="ignore")
    (run_dir / "stderr.txt").write_text(proc.stderr, encoding="utf-8", errors="ignore")
    return proc.returncode, rpt, out


def extract_simulated_series(
    out_path: Path,
    swmm_node: str,
    swmm_attr: str,
    aggregate: str,
) -> pd.DataFrame:
    label = f"node,{swmm_node},{swmm_attr}"
    series = swmmtoolbox.extract(str(out_path), label)
    df = series.reset_index()
    df.columns = ["timestamp", "flow"]
    df["timestamp"] = pd.to_datetime(df["timestamp"])
    if aggregate == "daily_mean":
        df = df.set_index("timestamp").resample("D").mean(numeric_only=True).reset_index()
    return df


def evaluate_trial(
    base_inp: Path,
    patch_map: dict,
    params: dict,
    observed: pd.DataFrame,
    run_dir: Path,
    swmm_node: str,
    swmm_attr: str,
    aggregate: str,
) -> dict:
    """Patch the INP, run swmm5, compute RMSE/peak/mean errors vs observed.

    Same scoring shape parameter_scout used. Returning a partial record on
    swmm5 failure (return_code != 0) is intentional: the caller can mark
    the trial invalid without aborting the whole sensitivity scan.
    """

    run_dir.mkdir(parents=True, exist_ok=True)
    patched = patch_inp_text(base_inp.read_text(errors="ignore"), patch_map, params)
    inp = run_dir / "model.inp"
    inp.write_text(patched, encoding="utf-8")
    rc, _, out_path = run_swmm(inp, run_dir)
    rec: dict[str, Any] = {"params": params, "run_dir": str(run_dir), "return_code": rc}
    if rc != 0:
        return rec
    sim = extract_simulated_series(out_path, swmm_node, swmm_attr, aggregate)
    merged = pd.merge(observed, sim, on="timestamp", how="inner", suffixes=("_obs", "_sim"))
    if merged.empty:
        return rec
    obs = merged["flow_obs"].astype(float)
    simv = merged["flow_sim"].astype(float)
    diff = simv - obs
    den = float(((obs - obs.mean()) ** 2).sum())
    nse = None if den == 0 else float(1 - ((diff.pow(2).sum()) / den))
    rmse = float(math.sqrt((diff.pow(2).mean())))
    peak_err_abs = float(abs(simv.max() - obs.max()))
    mean_err_abs = float(abs(simv.mean() - obs.mean()))
    rec.update({
        "nse": nse,
        "rmse": rmse,
        "peak_err_abs": peak_err_abs,
        "mean_err_abs": mean_err_abs,
        "max_sim": float(simv.max()),
        "max_obs": float(obs.max()),
        "mean_sim": float(simv.mean()),
        "mean_obs": float(obs.mean()),
    })
    return rec


def write_summary(path: Path, payload: dict) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_text(json.dumps(payload, indent=2), encoding="utf-8")


# ---------------------------------------------------------------------------
# OAT (port of parameter_scout)
# ---------------------------------------------------------------------------


def run_oat(args: argparse.Namespace, observed: pd.DataFrame) -> dict:
    patch_map = load_json(args.patch_map)
    base_params = load_json(args.base_params)
    scan_spec = load_json(args.scan_spec)
    summary: dict[str, Any] = {"method": "oat", "parameters": []}
    for pname, values in scan_spec.items():
        trials: list[dict] = []
        for idx, val in enumerate(values):
            params = dict(base_params)
            params[pname] = val
            run_dir = args.run_root / f"{pname}_{idx}"
            trials.append(
                evaluate_trial(
                    args.base_inp,
                    patch_map,
                    params,
                    observed,
                    run_dir,
                    args.swmm_node,
                    args.swmm_attr,
                    args.aggregate,
                )
            )
        valid = [t for t in trials if t.get("return_code") == 0 and t.get("rmse") is not None]
        if valid:
            rmses = [t["rmse"] for t in valid]
            peaks = [t["peak_err_abs"] for t in valid]
            importance = (max(rmses) - min(rmses)) + (max(peaks) - min(peaks))
            best = min(valid, key=lambda x: x["rmse"] + x["peak_err_abs"] + x["mean_err_abs"])
            base_val = base_params[pname]
            if best["params"][pname] < base_val:
                direction = "down"
                next_range = [min(values), base_val]
            elif best["params"][pname] > base_val:
                direction = "up"
                next_range = [base_val, max(values)]
            else:
                direction = "stay"
                next_range = [min(values), max(values)]
        else:
            importance = None
            best = None
            direction = "unclear"
            next_range = [min(values), max(values)]
        summary["parameters"].append({
            "parameter": pname,
            "tested_values": values,
            "importance": importance,
            "recommended_direction": direction,
            "suggested_next_range": next_range,
            "best_trial": best,
            "trials": valid,
        })

    summary["parameters"].sort(
        key=lambda x: x["importance"] if x["importance"] is not None else -1,
        reverse=True,
    )
    return summary


# ---------------------------------------------------------------------------
# Variance-based helpers (Morris, Sobol')
# ---------------------------------------------------------------------------


def _build_salib_problem(parameter_space: dict) -> tuple[dict, list[str]]:
    """Turn `parameter_space.json` into a SALib `problem` dict.

    Names are read in insertion order so the sample matrix columns line up
    with the parameter names we report back.
    """

    if not parameter_space:
        raise ValueError("parameter_space.json must contain at least one parameter")
    names: list[str] = []
    bounds: list[list[float]] = []
    for name, spec in parameter_space.items():
        if "min" not in spec or "max" not in spec:
            raise ValueError(f"Parameter {name!r} must define both 'min' and 'max'")
        names.append(name)
        bounds.append([float(spec["min"]), float(spec["max"])])
    return {"num_vars": len(names), "names": names, "bounds": bounds}, names


def _score_rmse(rec: dict) -> float:
    """Convert an evaluate_trial record to an RMSE for SA scoring.

    Failed runs become +inf so SALib has a numeric to work with; this is
    consistent with treating a crashed swmm5 as "infinite deviation".
    """

    if rec.get("return_code") != 0 or rec.get("rmse") is None:
        return float("inf")
    return float(rec["rmse"])


def _propagate_samples(
    args: argparse.Namespace,
    samples,
    names: list[str],
    observed: pd.DataFrame,
) -> tuple[list[dict], list[float]]:
    """Run SWMM for each sample row; return per-row records and RMSE vector.

    `samples` is a 2-D NumPy array of shape (N_trials, k); row `i`
    becomes `dict(zip(names, samples[i]))` after patching.
    """

    patch_map = load_json(args.patch_map)
    records: list[dict] = []
    scores: list[float] = []
    for i, row in enumerate(samples):
        params = {name: float(value) for name, value in zip(names, row)}
        run_dir = args.run_root / f"trial_{i:04d}"
        rec = evaluate_trial(
            args.base_inp,
            patch_map,
            params,
            observed,
            run_dir,
            args.swmm_node,
            args.swmm_attr,
            args.aggregate,
        )
        records.append(rec)
        scores.append(_score_rmse(rec))
    return records, scores


def run_morris(args: argparse.Namespace, observed: pd.DataFrame) -> dict:
    """Morris elementary-effects via SALib.

    Sample budget = r * (k + 1), with `r = args.morris_r` trajectories and
    `k = len(parameter_space)` parameters. Outputs `mu_star` + `sigma`
    per parameter.
    """

    import numpy as np
    from SALib.sample import morris as morris_sampler
    from SALib.analyze import morris as morris_analyzer

    parameter_space = load_json(args.parameter_space)
    problem, names = _build_salib_problem(parameter_space)
    k = problem["num_vars"]
    r = int(args.morris_r)
    num_levels = int(args.morris_levels)

    samples = morris_sampler.sample(
        problem,
        N=r,
        num_levels=num_levels,
        seed=args.seed,
    )
    expected_budget = r * (k + 1)
    actual_budget = int(samples.shape[0])
    # SALib aligns the trajectory budget to r*(k+1) exactly; assert it so
    # the acceptance criterion never silently drifts.
    if actual_budget != expected_budget:
        raise RuntimeError(
            f"Morris sample budget mismatch: r*(k+1)={expected_budget} but SALib returned {actual_budget}"
        )

    records, scores = _propagate_samples(args, samples, names, observed)
    y = np.asarray(scores, dtype=float)

    # SALib's analyze chokes on non-finite scores. Replace inf with the
    # finite max so the index calculation can still complete; the trial
    # records preserve the original failure signal.
    finite_mask = np.isfinite(y)
    if not finite_mask.any():
        raise RuntimeError("All Morris trials failed; cannot compute indices")
    if not finite_mask.all():
        finite_max = float(y[finite_mask].max())
        y = np.where(finite_mask, y, finite_max)

    si = morris_analyzer.analyze(problem, samples, y, num_levels=num_levels, print_to_console=False)
    indices: dict[str, dict[str, float]] = {}
    for i, name in enumerate(names):
        indices[name] = {
            "mu": float(si["mu"][i]),
            "mu_star": float(si["mu_star"][i]),
            "sigma": float(si["sigma"][i]),
            "mu_star_conf": float(si["mu_star_conf"][i]),
        }

    return {
        "method": "morris",
        "objective": "rmse",
        "parameters": names,
        "sample_budget": expected_budget,
        "morris": {"r": r, "num_levels": num_levels},
        "indices": indices,
        "trials": records,
    }


def run_sobol(args: argparse.Namespace, observed: pd.DataFrame) -> dict:
    """Sobol' indices via SALib (saltelli + sobol.analyze).

    Sample budget = N * (2k + 2), with `N = args.sobol_n`. Outputs first-
    order `S_i` + total-effect `S_T_i`. We use the new `SALib.sample.sobol`
    helper (Saltelli sampling), which lines up with the budget formula and
    avoids the deprecation warning attached to the legacy import path.
    """

    import numpy as np
    from SALib.sample import sobol as sobol_sampler
    from SALib.analyze import sobol as sobol_analyzer

    parameter_space = load_json(args.parameter_space)
    problem, names = _build_salib_problem(parameter_space)
    k = problem["num_vars"]
    N = int(args.sobol_n)

    # The Saltelli budget formula in the issue spec — N * (2k + 2) — is the
    # variant with second-order interactions enabled. We don't *expose* the
    # S_ij matrix in sensitivity_indices.json (only first-order S_i and
    # total-effect S_T_i), but using calc_second_order=True keeps the
    # budget mathematically aligned with the acceptance criterion.
    samples = sobol_sampler.sample(problem, N, calc_second_order=True, seed=args.seed)
    expected_budget = N * (2 * k + 2)
    actual_budget = int(samples.shape[0])
    if actual_budget != expected_budget:
        raise RuntimeError(
            f"Sobol sample budget mismatch: N*(2k+2)={expected_budget} but SALib returned {actual_budget}"
        )

    records, scores = _propagate_samples(args, samples, names, observed)
    y = np.asarray(scores, dtype=float)
    finite_mask = np.isfinite(y)
    if not finite_mask.any():
        raise RuntimeError("All Sobol' trials failed; cannot compute indices")
    if not finite_mask.all():
        finite_max = float(y[finite_mask].max())
        y = np.where(finite_mask, y, finite_max)

    si = sobol_analyzer.analyze(problem, y, calc_second_order=True, print_to_console=False)
    indices: dict[str, dict[str, float]] = {}
    for i, name in enumerate(names):
        indices[name] = {
            "S_i": float(si["S1"][i]),
            "S_i_conf": float(si["S1_conf"][i]),
            "S_T_i": float(si["ST"][i]),
            "S_T_i_conf": float(si["ST_conf"][i]),
        }

    return {
        "method": "sobol",
        "objective": "rmse",
        "parameters": names,
        "sample_budget": expected_budget,
        "sobol": {"N": N, "calc_second_order": True},
        "indices": indices,
        "trials": records,
    }


# ---------------------------------------------------------------------------
# CLI
# ---------------------------------------------------------------------------


def build_argparser() -> argparse.ArgumentParser:
    ap = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    ap.add_argument(
        "--method",
        required=True,
        choices=["oat", "morris", "sobol"],
        help="Which sensitivity-analysis sub-method to run.",
    )
    ap.add_argument("--base-inp", required=True, type=Path)
    ap.add_argument("--patch-map", required=True, type=Path)
    ap.add_argument("--observed", required=True, type=Path)
    ap.add_argument("--run-root", required=True, type=Path)
    ap.add_argument("--summary-json", required=True, type=Path)
    ap.add_argument("--swmm-node", default="O1")
    ap.add_argument("--swmm-attr", default="Total_inflow")
    ap.add_argument("--aggregate", choices=["none", "daily_mean"], default="none")
    ap.add_argument("--timestamp-col", default=None)
    ap.add_argument("--flow-col", default=None)
    ap.add_argument("--time-format", default=None)
    ap.add_argument("--obs-start", default=None)
    ap.add_argument("--obs-end", default=None)
    ap.add_argument("--seed", type=int, default=42)

    # OAT-specific.
    ap.add_argument(
        "--base-params",
        type=Path,
        help="JSON object with the baseline parameter values (OAT only).",
    )
    ap.add_argument(
        "--scan-spec",
        type=Path,
        help="JSON object: parameter -> list of trial values (OAT only).",
    )

    # Variance-based (Morris/Sobol').
    ap.add_argument(
        "--parameter-space",
        type=Path,
        help="JSON object: parameter -> {min, max, type?} (Morris and Sobol').",
    )
    ap.add_argument(
        "--morris-r",
        type=int,
        default=10,
        help="Number of Morris trajectories (sample budget = r*(k+1)).",
    )
    ap.add_argument(
        "--morris-levels",
        type=int,
        default=4,
        help="Number of grid levels for Morris (default 4, the SALib default).",
    )
    ap.add_argument(
        "--sobol-n",
        type=int,
        default=256,
        help="Saltelli base sample size (sample budget = N*(2k+2)).",
    )
    return ap


def main() -> None:
    args = build_argparser().parse_args()

    observed = read_series(
        args.observed,
        timestamp_col=args.timestamp_col,
        flow_col=args.flow_col,
        time_format=args.time_format,
    )
    observed = filter_series_window(observed, args.obs_start, args.obs_end)

    args.run_root.mkdir(parents=True, exist_ok=True)

    if args.method == "oat":
        if args.base_params is None or args.scan_spec is None:
            raise SystemExit("--method oat requires --base-params and --scan-spec")
        summary = run_oat(args, observed)
    elif args.method == "morris":
        if args.parameter_space is None:
            raise SystemExit("--method morris requires --parameter-space")
        summary = run_morris(args, observed)
    elif args.method == "sobol":
        if args.parameter_space is None:
            raise SystemExit("--method sobol requires --parameter-space")
        summary = run_sobol(args, observed)
    else:  # pragma: no cover - argparse choices guard
        raise SystemExit(f"Unsupported --method: {args.method}")

    write_summary(args.summary_json, summary)
    # Stay quiet on stdout for non-OAT methods so the per-trial swmm5
    # output never floods CI logs; the JSON file is the contract.
    print(json.dumps({"method": summary["method"], "summary_json": str(args.summary_json)}, indent=2))


if __name__ == "__main__":
    main()

#!/usr/bin/env python3
"""Integrated uncertainty source decomposition (issue #55).

This module is the FINAL impl deliverable for the PRD
"Uncertainty and Calibration Strengthening" track. It reads whichever
raw uncertainty outputs are present in ``<run_dir>/09_audit/`` and
emits two files alongside them:

* ``uncertainty_source_summary.md`` — paper-reviewer-facing report with
  five fixed sections + an Evidence Boundary header that lists every
  potential method with ✓/✗ so no method is silently absent.
* ``uncertainty_source_decomposition.json`` — machine-readable mirror
  of the markdown ( ``schema_version == "1.0"`` ).

The function is **pure over the filesystem state**: same inputs always
produce the same outputs, no SWMM execution, no network, no global
state mutation outside the audit dir. Re-invoking the function
overwrites the two output files in-place so the latest state of a run
is always represented.

Inputs consumed (each is optional, all detected by filename):

* ``sensitivity_indices.json`` — Sobol' or Morris (Slice 4 / #49).
* ``posterior_samples.csv`` + ``chain_convergence.json`` — DREAM-ZS
  (Slice 2 / #53).
* ``candidate_calibration.json`` — SCE-UA / DREAM-ZS candidate handover
  (Slice 6 / #54). ``strategy`` distinguishes the method that produced
  it.
* ``rainfall_ensemble_summary.json`` — Rainfall ensemble (Slice 5 /
  #51). ``method`` ∈ {``perturbation``, ``idf``} tells us which sub-
  method (A or B) actually ran.
* ``uncertainty_summary.json`` — MC propagation summary (legacy).

Schema (JSON file)::

    {
      "schema_version": "1.0",
      "generated_at_utc": "...",
      "run_id": "...",
      "evidence_boundary": {
        "sobol":             {"ran": bool, "source": "<rel path or null>"},
        "morris":            {"ran": bool, "source": "<rel path or null>"},
        "dream_zs":          {"ran": bool, "source": "<rel path or null>"},
        "sce_ua":            {"ran": bool, "source": "<rel path or null>"},
        "rainfall_ensemble": {
            "ran": bool,
            "method": "perturbation" | "idf" | null,
            "source": "<rel path or null>"
        },
        "mc_propagation":    {"ran": bool, "source": "<rel path or null>"}
      },
      "output_envelope":         {...} | null,
      "parameter_contribution":  {...} | null,
      "input_contribution":      {...} | null,
      "structural_assumptions":  [<str>, ...],
      "cross_references":        {"<artefact id>": "<rel path>", ...}
    }
"""

from __future__ import annotations

import argparse
import json
import sys
from dataclasses import dataclass
from datetime import datetime, timezone
from pathlib import Path
from typing import Any, Mapping


SCHEMA_VERSION = "1.0"


# ---------------------------------------------------------------------------
# Pure helpers
# ---------------------------------------------------------------------------


def _utc_now_iso() -> str:
    return datetime.now(timezone.utc).isoformat(timespec="seconds")


def _read_json(path: Path) -> dict[str, Any] | None:
    """Return JSON contents at ``path`` or ``None`` if missing/unreadable.

    We swallow malformed JSON because the caller's responsibility is to
    report "method present" or "method absent"; if a file exists but is
    corrupt, treating it as absent yields a more useful audit signal
    than crashing the whole report.
    """
    if not path.is_file():
        return None
    try:
        return json.loads(path.read_text(encoding="utf-8"))
    except (OSError, json.JSONDecodeError):
        return None


def _rel(path: Path, root: Path) -> str:
    """Return ``path`` relative to ``root`` (posix-style); fallback to str."""
    try:
        return path.resolve().relative_to(root.resolve()).as_posix()
    except ValueError:
        return str(path)


@dataclass(frozen=True)
class DecompositionResult:
    """Return value from :func:`decompose`.

    Keeping the two output paths in a small dataclass means callers
    (CLI, tests, audit-pipeline hook) can introspect what was written
    without re-deriving the names.
    """

    markdown_path: Path
    json_path: Path
    payload: dict[str, Any]
    methods_present: list[str]
    methods_absent: list[str]


# ---------------------------------------------------------------------------
# Evidence-boundary detection
# ---------------------------------------------------------------------------


def _detect_evidence(audit_dir: Path, run_dir: Path) -> dict[str, dict[str, Any]]:
    """Return the structured ``evidence_boundary`` dict.

    Each row reports ``ran`` (bool), and a relative-to-run-dir
    ``source`` path when present. The sensitivity-indices file is
    one-of {Sobol', Morris}; the Morris vs Sobol detection key is the
    ``method`` field inside it. Candidate calibration covers both
    SCE-UA and DREAM-ZS strategies; the ``strategy`` field discriminates.
    """
    out: dict[str, dict[str, Any]] = {
        "sobol": {"ran": False, "source": None},
        "morris": {"ran": False, "source": None},
        "dream_zs": {"ran": False, "source": None},
        "sce_ua": {"ran": False, "source": None},
        "rainfall_ensemble": {"ran": False, "method": None, "source": None},
        "mc_propagation": {"ran": False, "source": None},
    }

    sens = _read_json(audit_dir / "sensitivity_indices.json")
    if sens:
        method = str(sens.get("method", "")).lower()
        if method == "sobol":
            out["sobol"] = {
                "ran": True,
                "source": _rel(audit_dir / "sensitivity_indices.json", run_dir),
            }
        elif method == "morris":
            out["morris"] = {
                "ran": True,
                "source": _rel(audit_dir / "sensitivity_indices.json", run_dir),
            }
        # method=="oat" is screening, not a variance-decomposition
        # method; the integrated report does not surface it as a
        # parameter-contribution evidence slot, but we still flag the
        # file in cross-references.

    posterior_csv = audit_dir / "posterior_samples.csv"
    convergence = audit_dir / "chain_convergence.json"
    if posterior_csv.is_file():
        # DREAM-ZS writes both files together; require the CSV (the
        # primary evidence) and let the convergence JSON be optional.
        out["dream_zs"] = {
            "ran": True,
            "source": _rel(posterior_csv, run_dir),
            "convergence_source": _rel(convergence, run_dir) if convergence.is_file() else None,
        }

    candidate = _read_json(audit_dir / "candidate_calibration.json")
    if candidate:
        strategy = str(candidate.get("strategy", "")).lower()
        rel = _rel(audit_dir / "candidate_calibration.json", run_dir)
        if strategy == "sce-ua":
            out["sce_ua"] = {"ran": True, "source": rel}
        elif strategy == "dream-zs":
            # DREAM-ZS already detected via the CSV; record the candidate
            # ref for cross-references but don't overwrite the dream_zs
            # slot's "source".
            pass
        else:
            # Unknown strategy — still flag SCE-UA as the closest mode so
            # the agent doesn't silently lose the candidate evidence.
            out["sce_ua"] = {"ran": True, "source": rel, "strategy_label": strategy or "unknown"}

    rainfall = _read_json(audit_dir / "rainfall_ensemble_summary.json")
    if rainfall:
        rmethod = str(rainfall.get("method", "")).lower()
        out["rainfall_ensemble"] = {
            "ran": True,
            "method": rmethod or None,
            "source": _rel(audit_dir / "rainfall_ensemble_summary.json", run_dir),
        }

    mc = _read_json(audit_dir / "uncertainty_summary.json")
    if mc:
        out["mc_propagation"] = {
            "ran": True,
            "source": _rel(audit_dir / "uncertainty_summary.json", run_dir),
        }

    return out


# ---------------------------------------------------------------------------
# Section builders
# ---------------------------------------------------------------------------


def _build_output_envelope(mc: dict[str, Any] | None, rainfall: dict[str, Any] | None) -> dict[str, Any] | None:
    """Combine the MC envelope + rainfall ensemble peak flow envelope.

    Returns ``None`` when neither input is present. When both are
    present we keep them as two separate sub-blocks (they represent
    different perturbation sources, so collapsing them would mislead
    readers).
    """
    if mc is None and rainfall is None:
        return None
    block: dict[str, Any] = {}
    if mc is not None:
        block["mc_propagation"] = {
            "samples": mc.get("samples"),
            "node": mc.get("node"),
            "peak_cms_envelope": mc.get("peak_cms_envelope"),
            "peak_percent_change_envelope": mc.get("peak_percent_change_envelope"),
        }
    if rainfall is not None:
        swmm_stats = rainfall.get("swmm_ensemble_stats") or {}
        block["rainfall_ensemble"] = {
            "method": rainfall.get("method"),
            "n_realisations": rainfall.get("n_realisations"),
            "peak_flow": swmm_stats.get("peak_flow"),
            "total_volume_m3": swmm_stats.get("total_volume_m3"),
        }
    return block


def _sorted_sobol_indices(sens: dict[str, Any] | None) -> list[dict[str, Any]]:
    """Return ``indices`` rows sorted by ``S_T_i`` descending.

    Returns an empty list when ``sens`` is None or method != "sobol".
    """
    if not sens or str(sens.get("method", "")).lower() != "sobol":
        return []
    rows: list[dict[str, Any]] = []
    indices = sens.get("indices") or {}
    if not isinstance(indices, Mapping):
        return []
    for name, row in indices.items():
        if not isinstance(row, Mapping):
            continue
        rows.append(
            {
                "parameter": name,
                "S_i": row.get("S_i"),
                "S_i_conf": row.get("S_i_conf"),
                "S_T_i": row.get("S_T_i"),
                "S_T_i_conf": row.get("S_T_i_conf"),
            }
        )
    rows.sort(key=lambda r: float(r.get("S_T_i") or 0.0), reverse=True)
    return rows


def _sorted_morris_indices(sens: dict[str, Any] | None) -> list[dict[str, Any]]:
    """Return ``indices`` rows sorted by ``mu_star`` descending."""
    if not sens or str(sens.get("method", "")).lower() != "morris":
        return []
    rows: list[dict[str, Any]] = []
    indices = sens.get("indices") or {}
    if not isinstance(indices, Mapping):
        return []
    for name, row in indices.items():
        if not isinstance(row, Mapping):
            continue
        rows.append(
            {
                "parameter": name,
                "mu": row.get("mu"),
                "mu_star": row.get("mu_star"),
                "sigma": row.get("sigma"),
                "mu_star_conf": row.get("mu_star_conf"),
            }
        )
    rows.sort(key=lambda r: abs(float(r.get("mu_star") or 0.0)), reverse=True)
    return rows


def _build_parameter_contribution(sens: dict[str, Any] | None) -> dict[str, Any] | None:
    """Build the "Parameter contribution" block.

    When Sobol' indices are present we surface the sorted ``S_T_i``
    ranking. When only Morris is present we surface the ``mu_star``
    ranking and flag this fact so a paper-reviewer reading the report
    knows the values are screening-quality, not full variance
    decomposition.
    """
    if sens is None:
        return None
    method = str(sens.get("method", "")).lower()
    if method == "sobol":
        return {
            "method": "sobol",
            "sample_budget": sens.get("sample_budget"),
            "sobol_total_effect_sorted": _sorted_sobol_indices(sens),
        }
    if method == "morris":
        return {
            "method": "morris",
            "sample_budget": sens.get("sample_budget"),
            "morris_mu_star_sorted": _sorted_morris_indices(sens),
            "note": (
                "Morris elementary-effects is a screening method; "
                "rankings are reliable but the magnitudes are not "
                "variance-decomposition quality."
            ),
        }
    if method == "oat":
        return {
            "method": "oat",
            "note": "OAT is one-at-a-time screening; not a variance decomposition.",
        }
    return None


def _build_input_contribution(
    rainfall: dict[str, Any] | None,
    parameter_contribution: dict[str, Any] | None,
) -> dict[str, Any]:
    """Compare rainfall-induced and parameter-induced uncertainty.

    The integrated narrative is "how big is the rainfall envelope
    relative to the variance attributable to parameters?". We surface
    both sides so the reader can eyeball the split; we do **not**
    compute a single % attribution number because that requires an
    apples-to-apples reduction (variance of peak flow under rainfall
    perturbation vs Sobol' total-effect variance), and the underlying
    artefacts are not always emitted on the same target node.
    """
    block: dict[str, Any] = {
        "rainfall_ensemble": None,
        "parameter": None,
        "comparison_note": None,
    }
    if rainfall is not None:
        swmm_stats = rainfall.get("swmm_ensemble_stats") or {}
        rainfall_block: dict[str, Any] = {
            "method": rainfall.get("method"),
            "n_realisations": rainfall.get("n_realisations"),
            "rainfall_stats": rainfall.get("rainfall_ensemble_stats"),
        }
        peak_flow = swmm_stats.get("peak_flow") or {}
        if any(peak_flow.get(k) is not None for k in ("p05", "p50", "p95")):
            rainfall_block["peak_flow_envelope"] = peak_flow
        block["rainfall_ensemble"] = rainfall_block
    if parameter_contribution is not None:
        block["parameter"] = {
            "method": parameter_contribution.get("method"),
            "top": (
                parameter_contribution.get("sobol_total_effect_sorted")
                or parameter_contribution.get("morris_mu_star_sorted")
                or []
            )[:3],
        }
    if block["rainfall_ensemble"] and block["parameter"]:
        block["comparison_note"] = (
            "Rainfall-input uncertainty and parameter uncertainty are both "
            "quantified; compare the rainfall peak-flow envelope width "
            "against the top-ranked parameter's variance contribution to "
            "judge which source dominates for this case."
        )
    elif block["rainfall_ensemble"]:
        block["comparison_note"] = (
            "Only rainfall-input uncertainty was quantified; "
            "parameter contribution was not run."
        )
    elif block["parameter"]:
        block["comparison_note"] = (
            "Only parameter sensitivity was quantified; "
            "rainfall ensemble was not run."
        )
    else:
        block["comparison_note"] = (
            "Neither rainfall-input uncertainty nor a variance-based "
            "parameter sensitivity were run for this case."
        )
    return block


def _structural_assumptions(
    evidence: dict[str, dict[str, Any]],
    candidate: dict[str, Any] | None,
) -> list[str]:
    """Document the assumptions that are NOT quantified by this run.

    The reviewer-facing point is that uncertainty quantification is
    bounded: model-structural uncertainty (which conceptual model is
    correct), boundary-condition uncertainty (downstream BCs, time-
    series gaps), and observation-noise uncertainty are not propagated
    by any of the methods listed in the Evidence Boundary table.
    """
    items: list[str] = [
        "Model-structural uncertainty (choice of conceptual model, e.g. "
        "kinematic vs dynamic wave) is not quantified.",
        "Boundary-condition uncertainty (downstream stage, lateral "
        "inflows) is not quantified.",
        "Observation-noise uncertainty is not propagated through the "
        "likelihood — DREAM-ZS treats the observed series as exact.",
    ]
    if not evidence["rainfall_ensemble"]["ran"]:
        items.append(
            "Rainfall-input uncertainty was not quantified (no rainfall "
            "ensemble run); the report is parameter-uncertainty-only."
        )
    if not evidence["sobol"]["ran"] and not evidence["morris"]["ran"]:
        items.append(
            "No global sensitivity / variance decomposition was run; "
            "parameter contribution rankings are unavailable."
        )
    if not evidence["dream_zs"]["ran"] and not evidence["sce_ua"]["ran"]:
        items.append(
            "No calibration (DREAM-ZS or SCE-UA) was run; posterior or "
            "best-fit parameter sets are unavailable."
        )
    if candidate and candidate.get("evidence_boundary") == "candidate_not_accepted_yet":
        items.append(
            "The calibration candidate is recorded but the canonical INP "
            "has not been patched yet; run `aiswmm calibration accept "
            "<run_dir>` to make the candidate effective."
        )
    return items


def _build_cross_references(audit_dir: Path, run_dir: Path) -> dict[str, str]:
    """Map artefact-id -> relative path under ``<run_dir>`` for the markdown.

    Every artefact we *might* link from the report is included if it
    exists on disk. This keeps the markdown navigation reliable across
    partial runs without requiring callers to re-derive paths.
    """
    refs: dict[str, str] = {}
    candidates = [
        ("sensitivity_indices", audit_dir / "sensitivity_indices.json"),
        ("posterior_samples", audit_dir / "posterior_samples.csv"),
        ("chain_convergence", audit_dir / "chain_convergence.json"),
        ("posterior_correlation_plot", audit_dir / "posterior_correlation.png"),
        ("rainfall_ensemble_summary", audit_dir / "rainfall_ensemble_summary.json"),
        ("candidate_calibration", audit_dir / "candidate_calibration.json"),
        ("calibration_summary", audit_dir / "calibration_summary.json"),
        ("uncertainty_summary", audit_dir / "uncertainty_summary.json"),
        ("experiment_provenance", audit_dir / "experiment_provenance.json"),
        ("experiment_note", audit_dir / "experiment_note.md"),
    ]
    for key, path in candidates:
        if path.is_file():
            refs[key] = _rel(path, run_dir)
    # Marginal posterior PNGs are auto-discovered (one per parameter)
    for png in sorted(audit_dir.glob("posterior_*.png")):
        if png.name == "posterior_correlation.png":
            continue
        refs[png.stem] = _rel(png, run_dir)
    return refs


# ---------------------------------------------------------------------------
# Markdown renderer
# ---------------------------------------------------------------------------


_BOUNDARY_LABELS = [
    ("sobol", "Sobol' SA       "),
    ("morris", "Morris SA       "),
    ("dream_zs", "DREAM-ZS        "),
    ("sce_ua", "SCE-UA          "),
    ("rainfall_ensemble", "Rainfall ensemble"),
    ("mc_propagation", "MC propagation  "),
]


def _render_evidence_boundary(evidence: dict[str, dict[str, Any]]) -> str:
    """Render the ``Evidence boundary:`` code block.

    The label column is pre-padded so the ``:`` aligns vertically; the
    "Rainfall ensemble" row is intentionally one column wider than the
    others because the label itself is longer than the rest. Tests
    assert against the literal aligned text, so do not shorten any of
    the padded labels.
    """
    lines = ["```", "Evidence boundary:"]
    for key, label in _BOUNDARY_LABELS:
        row = evidence[key]
        mark = "✓" if row.get("ran") else "✗"
        suffix = ""
        if key == "rainfall_ensemble" and row.get("ran"):
            method = row.get("method")
            if method == "perturbation":
                suffix = " method A only (method B not run)"
            elif method == "idf":
                suffix = " method B only (method A not run)"
            else:
                suffix = " ran"
        elif row.get("ran"):
            source = row.get("source")
            suffix = f" ran ({Path(source).name})" if source else " ran"
        else:
            suffix = " not run"
        lines.append(f"  {label}: {mark}{suffix}")
    lines.append("```")
    return "\n".join(lines)


def _fmt_num(value: Any, digits: int = 4) -> str:
    if value is None:
        return "—"
    try:
        return f"{float(value):.{digits}f}"
    except (TypeError, ValueError):
        return str(value)


def _render_output_envelope(block: dict[str, Any] | None) -> str:
    parts = ["## Output uncertainty envelope", ""]
    if not block:
        parts.append("_No MC propagation or rainfall-driven SWMM ensemble outputs are present._")
        parts.append("")
        return "\n".join(parts)
    mc = block.get("mc_propagation")
    if mc:
        envelope = mc.get("peak_cms_envelope") or {}
        parts.append(f"**Monte Carlo parameter propagation** ({mc.get('samples') or '?'} samples at node `{mc.get('node') or '?'}`).")
        parts.append("")
        parts.append("| Quantile | Peak flow (cms) |")
        parts.append("|---|---|")
        parts.append(f"| p05 | {_fmt_num(envelope.get('p05'))} |")
        parts.append(f"| p50 | {_fmt_num(envelope.get('p50'))} |")
        parts.append(f"| p95 | {_fmt_num(envelope.get('p95'))} |")
        parts.append("")
    rainfall = block.get("rainfall_ensemble")
    if rainfall and rainfall.get("peak_flow"):
        peak = rainfall["peak_flow"]
        parts.append(
            f"**Rainfall ensemble — SWMM-propagated peak flow** "
            f"({rainfall.get('n_realisations') or '?'} realisations, "
            f"method `{rainfall.get('method') or '?'}`)."
        )
        parts.append("")
        parts.append("| Quantile | Peak flow (cms) |")
        parts.append("|---|---|")
        parts.append(f"| p05 | {_fmt_num(peak.get('p05'))} |")
        parts.append(f"| p50 | {_fmt_num(peak.get('p50'))} |")
        parts.append(f"| p95 | {_fmt_num(peak.get('p95'))} |")
        parts.append("")
    return "\n".join(parts)


def _render_parameter_contribution(block: dict[str, Any] | None) -> str:
    parts = ["## Parameter contribution (Sobol' total-effect, sorted)", ""]
    if not block:
        parts.append("_No variance-based sensitivity analysis was run for this case._")
        parts.append("")
        return "\n".join(parts)
    method = block.get("method")
    if method == "sobol":
        parts.append(f"Sample budget: `N*(2k+2)` = {block.get('sample_budget')}.")
        parts.append("")
        parts.append("| Parameter | S_T_i | 95% CI | S_i (first-order) |")
        parts.append("|---|---|---|---|")
        for row in block.get("sobol_total_effect_sorted", []):
            parts.append(
                f"| `{row['parameter']}` | {_fmt_num(row.get('S_T_i'))} | "
                f"±{_fmt_num(row.get('S_T_i_conf'))} | {_fmt_num(row.get('S_i'))} |"
            )
        parts.append("")
    elif method == "morris":
        parts.append(f"Sample budget: `r*(k+1)` = {block.get('sample_budget')}.")
        parts.append("")
        parts.append(
            "_Note: this case ran Morris screening, not Sobol' decomposition. "
            "Rankings are reliable but magnitudes are not variance-quality._"
        )
        parts.append("")
        parts.append("| Parameter | mu_star | sigma | mu_star_conf |")
        parts.append("|---|---|---|---|")
        for row in block.get("morris_mu_star_sorted", []):
            parts.append(
                f"| `{row['parameter']}` | {_fmt_num(row.get('mu_star'))} | "
                f"{_fmt_num(row.get('sigma'))} | ±{_fmt_num(row.get('mu_star_conf'))} |"
            )
        parts.append("")
    else:
        parts.append("_Only OAT screening was run; parameter contribution is unavailable._")
        parts.append("")
    return "\n".join(parts)


def _render_input_contribution(block: dict[str, Any]) -> str:
    parts = ["## Input contribution (rainfall ensemble vs parameter)", ""]
    rainfall = block.get("rainfall_ensemble")
    parameter = block.get("parameter")
    if rainfall:
        peak = rainfall.get("peak_flow_envelope") or {}
        parts.append(
            f"**Rainfall ensemble** — method `{rainfall.get('method')}`, "
            f"{rainfall.get('n_realisations')} realisations."
        )
        if peak:
            parts.append("")
            parts.append(
                f"- peak-flow envelope (cms): "
                f"p05 = {_fmt_num(peak.get('p05'))}, "
                f"p50 = {_fmt_num(peak.get('p50'))}, "
                f"p95 = {_fmt_num(peak.get('p95'))}"
            )
        rainfall_stats = rainfall.get("rainfall_stats") or {}
        intensity = rainfall_stats.get("peak_intensity_mm_per_hr") or {}
        if intensity:
            parts.append(
                f"- rainfall peak-intensity envelope (mm/hr): "
                f"p05 = {_fmt_num(intensity.get('p05'))}, "
                f"p50 = {_fmt_num(intensity.get('p50'))}, "
                f"p95 = {_fmt_num(intensity.get('p95'))}"
            )
        parts.append("")
    if parameter and parameter.get("top"):
        parts.append(f"**Parameter contribution top-3** (method: `{parameter.get('method')}`).")
        parts.append("")
        for row in parameter["top"]:
            label = row.get("parameter")
            if "S_T_i" in row:
                parts.append(f"- `{label}` (S_T_i = {_fmt_num(row.get('S_T_i'))})")
            elif "mu_star" in row:
                parts.append(f"- `{label}` (mu_star = {_fmt_num(row.get('mu_star'))})")
            else:
                parts.append(f"- `{label}`")
        parts.append("")
    note = block.get("comparison_note")
    if note:
        parts.append(f"_{note}_")
        parts.append("")
    return "\n".join(parts)


def _render_structural_assumptions(items: list[str]) -> str:
    parts = ["## Structural assumptions (not quantified)", ""]
    for item in items:
        parts.append(f"- {item}")
    parts.append("")
    return "\n".join(parts)


def _render_cross_references(refs: dict[str, str]) -> str:
    parts = ["## Cross-references", ""]
    if not refs:
        parts.append("_No raw uncertainty artefacts on disk for this run._")
        parts.append("")
        return "\n".join(parts)
    parts.append("| Artefact | Path (relative to run dir) |")
    parts.append("|---|---|")
    for key in sorted(refs):
        parts.append(f"| `{key}` | `{refs[key]}` |")
    parts.append("")
    return "\n".join(parts)


# ---------------------------------------------------------------------------
# Top-level decompose()
# ---------------------------------------------------------------------------


def decompose(run_dir: Path | str) -> DecompositionResult:
    """Build and write the integrated uncertainty source decomposition.

    Always overwrites ``<run_dir>/09_audit/uncertainty_source_summary.md``
    and ``<run_dir>/09_audit/uncertainty_source_decomposition.json``;
    callers that want a single source of truth on disk should run this
    after every relevant artefact is updated.
    """
    run_dir = Path(run_dir)
    audit = run_dir / "09_audit"
    audit.mkdir(parents=True, exist_ok=True)

    sens = _read_json(audit / "sensitivity_indices.json")
    rainfall = _read_json(audit / "rainfall_ensemble_summary.json")
    candidate = _read_json(audit / "candidate_calibration.json")
    mc = _read_json(audit / "uncertainty_summary.json")
    provenance = _read_json(audit / "experiment_provenance.json") or {}

    evidence = _detect_evidence(audit, run_dir)
    output_envelope = _build_output_envelope(mc, rainfall)
    parameter_contribution = _build_parameter_contribution(sens)
    input_contribution = _build_input_contribution(rainfall, parameter_contribution)
    structural = _structural_assumptions(evidence, candidate)
    refs = _build_cross_references(audit, run_dir)

    run_id = provenance.get("run_id") or run_dir.name

    methods_present = [
        label
        for key, label in _BOUNDARY_LABELS
        if evidence[key].get("ran")
    ]
    methods_absent = [
        label
        for key, label in _BOUNDARY_LABELS
        if not evidence[key].get("ran")
    ]

    payload: dict[str, Any] = {
        "schema_version": SCHEMA_VERSION,
        "generated_at_utc": _utc_now_iso(),
        "run_id": run_id,
        "evidence_boundary": evidence,
        "output_envelope": output_envelope,
        "parameter_contribution": parameter_contribution,
        "input_contribution": input_contribution,
        "structural_assumptions": structural,
        "cross_references": refs,
    }

    md_parts = [
        f"# Uncertainty source decomposition — `{run_id}`",
        "",
        "Generated by `skills/swmm-uncertainty/scripts/source_decomposition.py`.",
        "This report integrates the raw uncertainty outputs present in "
        f"`{_rel(audit, run_dir)}` into a single paper-reviewer-facing summary.",
        "",
        _render_evidence_boundary(evidence),
        "",
        _render_output_envelope(output_envelope),
        _render_parameter_contribution(parameter_contribution),
        _render_input_contribution(input_contribution),
        _render_structural_assumptions(structural),
        _render_cross_references(refs),
    ]
    markdown = "\n".join(md_parts).rstrip() + "\n"

    md_path = audit / "uncertainty_source_summary.md"
    json_path = audit / "uncertainty_source_decomposition.json"
    md_path.write_text(markdown, encoding="utf-8")
    json_path.write_text(json.dumps(payload, indent=2, sort_keys=True) + "\n", encoding="utf-8")

    return DecompositionResult(
        markdown_path=md_path,
        json_path=json_path,
        payload=payload,
        methods_present=methods_present,
        methods_absent=methods_absent,
    )


# ---------------------------------------------------------------------------
# Script entry point (for the MCP server and dev invocations)
# ---------------------------------------------------------------------------


def _parse_args(argv: list[str] | None) -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Integrate uncertainty raw outputs in <run_dir>/09_audit/ into "
        "uncertainty_source_summary.md + uncertainty_source_decomposition.json.",
    )
    parser.add_argument("run_dir", type=Path, help="Path to the run directory.")
    return parser.parse_args(argv)


def main(argv: list[str] | None = None) -> int:
    args = _parse_args(argv)
    if not args.run_dir.is_dir():
        print(f"error: run_dir is not a directory: {args.run_dir}", file=sys.stderr)
        return 1
    result = decompose(args.run_dir)
    print(
        json.dumps(
            {
                "ok": True,
                "schema_version": SCHEMA_VERSION,
                "markdown_path": str(result.markdown_path),
                "json_path": str(result.json_path),
                "methods_present": result.methods_present,
                "methods_absent": result.methods_absent,
            },
            indent=2,
        )
    )
    return 0


if __name__ == "__main__":
    raise SystemExit(main())

#!/usr/bin/env python3
from __future__ import annotations

import argparse
import json
import shutil
import subprocess
import sys
import time
from collections import defaultdict
from datetime import datetime, timezone
from pathlib import Path
from typing import Any


SCRIPT_DIR = Path(__file__).resolve().parent
REPO_ROOT = SCRIPT_DIR.parents[2]
CALIBRATION_SCRIPTS = REPO_ROOT / "skills" / "swmm-calibration" / "scripts"
RUNNER_SCRIPT = REPO_ROOT / "skills" / "swmm-runner" / "scripts" / "swmm_runner.py"
if str(CALIBRATION_SCRIPTS) not in sys.path:
    sys.path.insert(0, str(CALIBRATION_SCRIPTS))
if str(SCRIPT_DIR) not in sys.path:
    sys.path.insert(0, str(SCRIPT_DIR))

from fuzzy_membership import build_alpha_intervals, read_baseline_values, resolve_fuzzy_space, write_json  # noqa: E402
from inp_patch import patch_inp_text  # noqa: E402
from sampling import generate_parameter_sets  # noqa: E402


def load_json(path: Path) -> Any:
    return json.loads(path.read_text(encoding="utf-8"))


def utc_now() -> str:
    return datetime.now(timezone.utc).isoformat(timespec="seconds")


def status_counts(results: list[dict[str, Any]]) -> dict[str, int]:
    counts = {"total": len(results), "ok": 0, "failed": 0, "invalid": 0, "dry_run": 0, "other": 0}
    for rec in results:
        status = str(rec.get("status", "other"))
        counts[status if status in counts else "other"] += 1
    return counts


def numeric_envelope(values: list[Any]) -> dict[str, Any]:
    nums = [float(v) for v in values if isinstance(v, (int, float))]
    if not nums:
        return {"count": 0, "min": None, "max": None, "mean": None}
    return {
        "count": len(nums),
        "min": min(nums),
        "max": max(nums),
        "mean": sum(nums) / len(nums),
    }


def run_runner(inp: Path, run_dir: Path, node: str) -> tuple[int, dict[str, Any], str, str]:
    cmd = [
        sys.executable,
        str(RUNNER_SCRIPT),
        "run",
        "--inp",
        str(inp),
        "--run-dir",
        str(run_dir),
        "--node",
        node,
    ]
    proc = subprocess.run(cmd, cwd=REPO_ROOT, capture_output=True, text=True)
    parsed: dict[str, Any] = {}
    if proc.stdout.strip():
        try:
            obj = json.loads(proc.stdout)
            if isinstance(obj, dict):
                parsed = obj
        except json.JSONDecodeError:
            parsed = {}
    return proc.returncode, parsed, proc.stdout, proc.stderr


def evaluate_trial(
    *,
    base_inp_text: str,
    patch_map: dict[str, Any],
    trial: dict[str, Any],
    trials_dir: Path,
    swmm_node: str,
    dry_run: bool,
) -> dict[str, Any]:
    started = time.perf_counter()
    trial_dir = trials_dir / trial["name"]
    trial_dir.mkdir(parents=True, exist_ok=True)
    inp_path = trial_dir / "model.inp"
    result: dict[str, Any] = {
        "trial": trial["name"],
        "params": trial["params"],
        "metadata": trial.get("metadata", {}),
        "run_dir": str(trial_dir),
        "status": "pending",
        "reason_code": None,
        "reason_detail": None,
        "metrics": {},
        "files": {"inp": str(inp_path)},
        "started_at_utc": utc_now(),
        "dry_run": dry_run,
    }

    try:
        inp_path.write_text(patch_inp_text(base_inp_text, patch_map, trial["params"]), encoding="utf-8")
    except Exception as exc:  # noqa: BLE001
        result.update(
            {
                "status": "invalid",
                "reason_code": "patch_failed",
                "reason_detail": str(exc),
                "elapsed_seconds": round(time.perf_counter() - started, 6),
            }
        )
        return result

    if dry_run:
        result.update(
            {
                "status": "dry_run",
                "reason_code": "dry_run_enabled",
                "reason_detail": "Trial INP was generated but SWMM was not executed.",
                "elapsed_seconds": round(time.perf_counter() - started, 6),
            }
        )
        return result

    if not shutil.which("swmm5"):
        result.update(
            {
                "status": "failed",
                "reason_code": "swmm_binary_missing",
                "reason_detail": "swmm5 executable was not found on PATH.",
                "elapsed_seconds": round(time.perf_counter() - started, 6),
            }
        )
        return result

    rc, manifest, stdout, stderr = run_runner(inp_path, trial_dir, swmm_node)
    result["return_code"] = rc
    result["files"].update(
        {
            "rpt": str(trial_dir / "model.rpt"),
            "out": str(trial_dir / "model.out"),
            "manifest": str(trial_dir / "manifest.json"),
            "stdout": str(trial_dir / "stdout.txt"),
            "stderr": str(trial_dir / "stderr.txt"),
        }
    )
    if not (trial_dir / "stdout.txt").exists():
        (trial_dir / "uncertainty_runner_stdout.txt").write_text(stdout, encoding="utf-8", errors="ignore")
    if not (trial_dir / "stderr.txt").exists():
        (trial_dir / "uncertainty_runner_stderr.txt").write_text(stderr, encoding="utf-8", errors="ignore")

    if rc != 0:
        result.update(
            {
                "status": "failed",
                "reason_code": "swmm_execution_failed",
                "reason_detail": f"SWMM runner returned non-zero exit code {rc}.",
                "elapsed_seconds": round(time.perf_counter() - started, 6),
            }
        )
        return result

    metrics = manifest.get("metrics", {}) if isinstance(manifest, dict) else {}
    result["metrics"] = {
        "peak": metrics.get("peak"),
        "continuity": metrics.get("continuity"),
    }
    result.update(
        {
            "status": "ok",
            "reason_code": "ok",
            "reason_detail": None,
            "elapsed_seconds": round(time.perf_counter() - started, 6),
        }
    )
    return result


def summarize_by_alpha(results: list[dict[str, Any]]) -> dict[str, Any]:
    grouped: dict[str, list[dict[str, Any]]] = defaultdict(list)
    for rec in results:
        alpha = (rec.get("metadata") or {}).get("alpha")
        grouped[f"{float(alpha):.2f}" if alpha is not None else "unknown"].append(rec)

    out: dict[str, Any] = {}
    for alpha, records in sorted(grouped.items(), key=lambda item: item[0]):
        ok_records = [rec for rec in records if rec.get("status") == "ok"]
        peak_values = [
            ((rec.get("metrics") or {}).get("peak") or {}).get("peak")
            for rec in ok_records
        ]
        runoff_continuity = [
            (((rec.get("metrics") or {}).get("continuity") or {}).get("continuity_error_percent") or {}).get(
                "runoff_quantity"
            )
            for rec in ok_records
        ]
        routing_continuity = [
            (((rec.get("metrics") or {}).get("continuity") or {}).get("continuity_error_percent") or {}).get(
                "flow_routing"
            )
            for rec in ok_records
        ]
        out[alpha] = {
            "status_counts": status_counts(records),
            "peak_flow": numeric_envelope(peak_values),
            "runoff_continuity_error_percent": numeric_envelope(runoff_continuity),
            "flow_routing_continuity_error_percent": numeric_envelope(routing_continuity),
        }
    return out


def parse_config(config_path: Path) -> dict[str, Any]:
    config = load_json(config_path)
    if not isinstance(config, dict):
        raise ValueError("Uncertainty config must be a JSON object")
    config.setdefault("alpha_levels", [0.0, 0.25, 0.5, 0.75, 1.0])
    config.setdefault("sampling", {})
    config["sampling"].setdefault("method", "lhs")
    config["sampling"].setdefault("samples_per_alpha", 20)
    config["sampling"].setdefault("seed", 42)
    config.setdefault("outputs", {})
    config["outputs"].setdefault("swmm_node", "O1")
    return config


def parse_args() -> argparse.Namespace:
    ap = argparse.ArgumentParser(description="Propagate fuzzy SWMM parameter uncertainty through alpha-cut samples.")
    ap.add_argument("--base-inp", required=True, type=Path)
    ap.add_argument("--patch-map", required=True, type=Path)
    ap.add_argument("--fuzzy-space", required=True, type=Path)
    ap.add_argument("--config", required=True, type=Path)
    ap.add_argument("--run-root", required=True, type=Path)
    ap.add_argument("--summary-json", required=True, type=Path)
    ap.add_argument("--dry-run", action="store_true")
    return ap.parse_args()


def main() -> None:
    args = parse_args()
    config = parse_config(args.config)
    patch_map = load_json(args.patch_map)
    fuzzy_space = load_json(args.fuzzy_space)
    base_inp_text = args.base_inp.read_text(encoding="utf-8", errors="ignore")

    run_root = args.run_root
    run_root.mkdir(parents=True, exist_ok=True)
    trials_dir = run_root / "trials"

    baseline_values = read_baseline_values(args.base_inp, patch_map)
    resolved = resolve_fuzzy_space(fuzzy_space, baseline_values)
    alpha_levels = [float(alpha) for alpha in config["alpha_levels"]]
    alpha_intervals = build_alpha_intervals(resolved, alpha_levels)

    sampling_cfg = config["sampling"]
    trials = generate_parameter_sets(
        alpha_intervals,
        method=str(sampling_cfg["method"]),
        samples_per_alpha=int(sampling_cfg["samples_per_alpha"]),
        seed=int(sampling_cfg["seed"]),
    )

    resolved_json = {"parameters": {name: param.to_dict() for name, param in resolved.items()}}
    parameter_sets_json = {"parameter_sets": trials}

    write_json(run_root / "fuzzy_space.resolved.json", resolved_json)
    write_json(run_root / "alpha_intervals.json", alpha_intervals)
    write_json(run_root / "parameter_sets.json", parameter_sets_json)

    swmm_node = str(config["outputs"].get("swmm_node", "O1"))
    results = [
        evaluate_trial(
            base_inp_text=base_inp_text,
            patch_map=patch_map,
            trial=trial,
            trials_dir=trials_dir,
            swmm_node=swmm_node,
            dry_run=bool(args.dry_run),
        )
        for trial in trials
    ]

    payload = {
        "mode": "fuzzy_uncertainty_propagation",
        "created_at_utc": utc_now(),
        "controls": {
            "base_inp": str(args.base_inp),
            "patch_map": str(args.patch_map),
            "fuzzy_space": str(args.fuzzy_space),
            "config": str(args.config),
            "run_root": str(run_root),
            "dry_run": bool(args.dry_run),
            "swmm_node": swmm_node,
        },
        "sampling": {
            "method": sampling_cfg["method"],
            "samples_per_alpha": sampling_cfg["samples_per_alpha"],
            "seed": sampling_cfg["seed"],
            "trial_count": len(trials),
        },
        "baseline_values": {name: baseline_values[name] for name in resolved},
        "resolved_fuzzy_space": resolved_json,
        "alpha_intervals": alpha_intervals,
        "status_counts": status_counts(results),
        "alpha_summary": summarize_by_alpha(results),
        "results": results,
    }

    write_json(args.summary_json, payload)
    print(json.dumps(payload, indent=2, sort_keys=True))


if __name__ == "__main__":
    main()

#!/usr/bin/env python3
from __future__ import annotations

import sys
import unittest
from pathlib import Path


SCRIPT_DIR = Path(__file__).resolve().parents[1] / "scripts"
sys.path.insert(0, str(SCRIPT_DIR))

from fuzzy_membership import build_alpha_intervals, resolve_fuzzy_space  # noqa: E402


class FuzzyMembershipTests(unittest.TestCase):
    def test_triangular_defaults_to_model_baseline(self) -> None:
        params = resolve_fuzzy_space(
            {
                "parameters": {
                    "pct_imperv_s1": {
                        "type": "triangular",
                        "lower": 15.0,
                        "upper": 40.0,
                        "baseline": "from_model",
                    }
                }
            },
            {"pct_imperv_s1": 25.0},
        )
        interval = params["pct_imperv_s1"].alpha_interval(0.5)
        self.assertEqual(interval, (20.0, 32.5))

    def test_triangular_rejects_baseline_outside_bounds(self) -> None:
        with self.assertRaises(ValueError):
            resolve_fuzzy_space(
                {"parameters": {"p": {"type": "triangular", "lower": 1.0, "upper": 2.0}}},
                {"p": 3.0},
            )

    def test_trapezoidal_core_width_around_baseline(self) -> None:
        params = resolve_fuzzy_space(
            {
                "parameters": {
                    "n": {
                        "type": "trapezoidal",
                        "lower": 0.01,
                        "upper": 0.03,
                        "core_width": 0.004,
                    }
                }
            },
            {"n": 0.02},
        )
        interval = params["n"].alpha_interval(0.5)
        self.assertAlmostEqual(interval[0], 0.014)
        self.assertAlmostEqual(interval[1], 0.026)

    def test_build_alpha_intervals(self) -> None:
        params = resolve_fuzzy_space(
            {"parameters": {"p": {"type": "triangular", "lower": 0, "upper": 10}}},
            {"p": 5},
        )
        intervals = build_alpha_intervals(params, [0.0, 1.0])
        cuts = intervals["parameters"]["p"]["alpha_cuts"]
        self.assertEqual(cuts[0]["lower"], 0.0)
        self.assertEqual(cuts[1]["lower"], 5.0)
        self.assertEqual(cuts[1]["upper"], 5.0)


if __name__ == "__main__":
    unittest.main()

#!/usr/bin/env python3
from __future__ import annotations

import sys
import unittest
from pathlib import Path


SCRIPT_DIR = Path(__file__).resolve().parents[1] / "scripts"
sys.path.insert(0, str(SCRIPT_DIR))

from sampling import generate_parameter_sets  # noqa: E402


class SamplingTests(unittest.TestCase):
    def test_lhs_generates_one_baseline_trial_for_degenerate_alpha(self) -> None:
        alpha_intervals = {
            "alpha_levels": [1.0],
            "parameters": {
                "p": {
                    "alpha_cuts": [
                        {
                            "alpha": 1.0,
                            "lower": 5.0,
                            "upper": 5.0,
                        }
                    ]
                }
            },
        }
        trials = generate_parameter_sets(alpha_intervals, method="lhs", samples_per_alpha=10, seed=42)
        self.assertEqual(len(trials), 1)
        self.assertEqual(trials[0]["params"], {"p": 5.0})

    def test_boundary_sampling_removes_duplicate_corners(self) -> None:
        alpha_intervals = {
            "alpha_levels": [0.0],
            "parameters": {
                "p": {"alpha_cuts": [{"alpha": 0.0, "lower": 1.0, "upper": 1.0}]},
                "q": {"alpha_cuts": [{"alpha": 0.0, "lower": 2.0, "upper": 4.0}]},
            },
        }
        trials = generate_parameter_sets(alpha_intervals, method="boundary", samples_per_alpha=10, seed=42)
        self.assertEqual(len(trials), 2)
        self.assertEqual({trial["params"]["q"] for trial in trials}, {2.0, 4.0})


if __name__ == "__main__":
    unittest.main()