#!/usr/bin/env python3
"""Candidate-handover artefacts for calibration runs (issue #54).

After any calibration strategy (SCE-UA, DREAM-ZS, random, lhs,
adaptive) finishes, the scaffold must *never* patch the canonical INP
on disk. Instead it writes three artefacts to ``<run_dir>/09_audit/``:

* ``candidate_calibration.json`` — best params + metrics + KGE
  decomposition + secondary metrics + ``evidence_boundary ==
  "candidate_not_accepted_yet"`` + the SHA256 of the patch file (used
  by ``aiswmm calibration accept`` for tamper detection) + (DREAM
  only) a reference to ``posterior_samples.csv``.
* ``candidate_inp_patch.json`` — list of one-line INP edits. Each row
  records ``param``, ``section``, ``object``, ``field_index``,
  ``old_value`` (what is in the canonical INP right now) and
  ``new_value`` (the calibrated value). This is enough for
  ``aiswmm calibration accept`` to re-apply the patch using the
  existing :mod:`inp_patch` machinery.
* ``calibration_report.md`` — human-readable summary with the KGE
  decomposition table, secondary-metrics table, strategy + iteration
  count, and references to the convergence trace / posterior plots
  when applicable.

The agent calls :func:`write_candidate_artefacts` at the end of every
strategy branch; ``aiswmm calibration accept`` is the only path that
turns the candidate into an actual on-disk change to the canonical
INP.
"""

from __future__ import annotations

import hashlib
import json
from datetime import datetime, timezone
from pathlib import Path
from typing import Any, Iterable, Mapping


EVIDENCE_BOUNDARY = "candidate_not_accepted_yet"
CANDIDATE_PATCH_SCHEMA = "1.0"
CANDIDATE_SCHEMA = "1.0"

CANDIDATE_FILENAME = "candidate_calibration.json"
PATCH_FILENAME = "candidate_inp_patch.json"
REPORT_FILENAME = "calibration_report.md"


# ---------------------------------------------------------------------------
# INP patch extraction
# ---------------------------------------------------------------------------


def _strip_inline_comment(line: str) -> str:
    """Return the code portion of an INP line (drop ``;...`` comment tail)."""
    code, *_ = line.split(";", 1)
    return code


def _find_old_value(
    text: str, section: str, obj_name: str, field_index: int
) -> str | None:
    """Walk an INP file and pull the token at ``[section] <object> ... field_index``.

    Mirrors :func:`inp_patch.patch_inp_text` so the diff we emit can be
    applied by the same logic. Returns ``None`` if the section/object
    row is not found or the field index is out of range.
    """
    section_norm = section.upper()
    current_section: str | None = None
    for raw in text.splitlines():
        stripped = raw.strip()
        if not stripped:
            continue
        if stripped.startswith("[") and stripped.endswith("]"):
            current_section = stripped.upper()
            continue
        if stripped.startswith(";"):
            continue
        if current_section != section_norm:
            continue
        code = _strip_inline_comment(raw)
        tokens = code.split()
        if not tokens or tokens[0] != obj_name:
            continue
        if field_index >= len(tokens):
            return None
        return tokens[field_index]
    return None


def build_inp_patch(
    inp_text: str,
    patch_map: Mapping[str, Mapping[str, Any]],
    params: Mapping[str, Any],
) -> dict[str, Any]:
    """Build the ``candidate_inp_patch.json`` payload.

    For each calibrated parameter we look up its row in ``inp_text``
    using the same selector contract as :func:`inp_patch.patch_inp_text`
    (section + object + zero-based field_index). The resulting JSON is
    line-oriented so a human auditor can read it and ``aiswmm
    calibration accept`` can reconstitute a ``{name: value}`` map for
    :func:`inp_patch.patch_inp_text`.

    Raises ``KeyError`` if any parameter is missing from ``patch_map``
    (matches the loud failure mode in :func:`inp_patch.patch_inp_text`).
    """
    missing_keys = sorted(set(params) - set(patch_map))
    if missing_keys:
        raise KeyError(
            "Parameters missing from patch_map: " + ", ".join(missing_keys)
        )
    edits: list[dict[str, Any]] = []
    for name in params:
        spec = patch_map[name]
        section = str(spec["section"])
        obj_name = str(spec["object"])
        field_index = int(spec["field_index"])
        old_value = _find_old_value(inp_text, section, obj_name, field_index)
        edits.append(
            {
                "param": name,
                "section": section,
                "object": obj_name,
                "field_index": field_index,
                "old_value": old_value,
                "new_value": str(params[name]),
            }
        )
    return {"schema_version": CANDIDATE_PATCH_SCHEMA, "edits": edits}


# ---------------------------------------------------------------------------
# SHA helpers (tamper-detection seam for `aiswmm calibration accept`)
# ---------------------------------------------------------------------------


def sha256_of_canonical_json(payload: Any) -> str:
    """SHA256 of ``json.dumps(payload, sort_keys=True, indent=2)``.

    The accept CLI re-computes this against the *on-disk* patch file
    and refuses the operation if it does not match the SHA recorded
    inside ``candidate_calibration.json``. ``sort_keys=True`` lets us
    canonicalise dict ordering across writers; ``indent=2`` matches the
    on-disk format so a human-readable file produces the same SHA as
    the in-memory dict.
    """
    text = json.dumps(payload, sort_keys=True, indent=2)
    return hashlib.sha256(text.encode("utf-8")).hexdigest()


def sha256_of_file(path: Path) -> str:
    """SHA256 of the raw bytes of ``path`` (used for canonical INP)."""
    h = hashlib.sha256()
    with path.open("rb") as fh:
        for chunk in iter(lambda: fh.read(8192), b""):
            h.update(chunk)
    return h.hexdigest()


# ---------------------------------------------------------------------------
# JSON writers
# ---------------------------------------------------------------------------


def _write_json_canonical(path: Path, payload: Any) -> None:
    """Write ``payload`` with the same formatting we hash over.

    Keeping the read-path and write-path text-identical is what makes
    :func:`sha256_of_canonical_json` and a re-read of the file produce
    the same digest.
    """
    path.parent.mkdir(parents=True, exist_ok=True)
    text = json.dumps(payload, sort_keys=True, indent=2)
    path.write_text(text + "\n", encoding="utf-8")


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


# ---------------------------------------------------------------------------
# Markdown report
# ---------------------------------------------------------------------------


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


def _render_kge_decomposition_table(decomp: Mapping[str, Any]) -> str:
    rows = [
        "| Component | Value |",
        "|---|---|",
        f"| r (correlation) | {_format_number(decomp.get('r'))} |",
        f"| alpha (variability ratio) | {_format_number(decomp.get('alpha'))} |",
        f"| beta (bias ratio) | {_format_number(decomp.get('beta'))} |",
    ]
    return "\n".join(rows)


def _render_secondary_table(secondary: Mapping[str, Any]) -> str:
    rows = [
        "| Metric | Value |",
        "|---|---|",
        f"| NSE | {_format_number(secondary.get('nse'))} |",
        f"| PBIAS (%) | {_format_number(secondary.get('pbias_pct'))} |",
        f"| RMSE | {_format_number(secondary.get('rmse'))} |",
        f"| Peak flow error (rel) | {_format_number(secondary.get('peak_error_rel'))} |",
        f"| Peak timing error (min) | {_format_number(secondary.get('peak_timing_min'))} |",
    ]
    return "\n".join(rows)


def _render_best_params_table(params: Mapping[str, Any]) -> str:
    if not params:
        return "_No parameter values reported._"
    rows = ["| Parameter | Value |", "|---|---|"]
    for name in sorted(params):
        rows.append(f"| {name} | {_format_number(params[name])} |")
    return "\n".join(rows)


def _render_posterior_section(
    summary: Mapping[str, Any],
    refs: Mapping[str, str],
) -> str:
    posterior = summary.get("posterior_summary")
    if not isinstance(posterior, Mapping):
        return ""
    lines: list[str] = ["## Posterior (DREAM-ZS)", ""]
    lines.append(
        f"- chains: {posterior.get('n_chains')} (requested: "
        f"{posterior.get('n_chains_requested')})"
    )
    lines.append(
        f"- post-burn-in samples: {posterior.get('n_samples_post_burnin')}"
    )
    lines.append(
        f"- Gelman-Rubin Rhat threshold: {posterior.get('rhat_threshold')}, "
        f"converged: {posterior.get('converged')}"
    )
    rhat = posterior.get("rhat") or {}
    if isinstance(rhat, Mapping) and rhat:
        lines.append("")
        lines.append("| Parameter | Rhat |")
        lines.append("|---|---|")
        for name in sorted(rhat):
            lines.append(f"| {name} | {_format_number(rhat[name])} |")
    posterior_csv = refs.get("posterior_samples_csv")
    if posterior_csv:
        lines.append("")
        lines.append(f"- posterior samples: `{posterior_csv}`")
    correlation_png = refs.get("posterior_correlation_png")
    if correlation_png:
        lines.append(f"- correlation plot: `{correlation_png}`")
    lines.append("")
    return "\n".join(lines)


def render_calibration_report(
    *,
    summary: Mapping[str, Any],
    best_params: Mapping[str, Any],
    candidate_inp_patch_sha256: str,
    refs: Mapping[str, str],
) -> str:
    """Build the markdown body for ``calibration_report.md``.

    Returned as a string so callers can write it (the writer also does
    this, but exposing the pure function keeps it cheap to unit-test).
    """
    strategy = str(summary.get("strategy", "unknown"))
    iterations = summary.get("iterations")
    primary = summary.get("primary_value")
    convergence_ref = refs.get("convergence_csv") or summary.get("convergence_trace_ref")

    parts: list[str] = []
    parts.append("# Calibration candidate report\n")
    parts.append(
        "> **Evidence boundary**: Candidate not accepted yet. The canonical "
        "INP on disk has not been modified. Run `aiswmm calibration accept "
        "<run_dir>` to apply the recorded patch and record a `human_decisions` "
        "entry on this run.\n"
    )
    parts.append("## Summary\n")
    parts.append(f"- strategy: `{strategy}`")
    parts.append(f"- primary objective: `{summary.get('primary_objective', 'kge')}`")
    parts.append(f"- primary value: {_format_number(primary)}")
    if iterations is not None:
        parts.append(f"- iterations: {iterations}")
    parts.append(f"- candidate INP patch SHA256: `{candidate_inp_patch_sha256}`")
    if convergence_ref:
        parts.append(f"- convergence trace: `{convergence_ref}`")
    parts.append("")
    parts.append("## KGE decomposition\n")
    parts.append(_render_kge_decomposition_table(summary.get("kge_decomposition") or {}))
    parts.append("")
    parts.append("## Secondary metrics\n")
    parts.append(_render_secondary_table(summary.get("secondary_metrics") or {}))
    parts.append("")
    parts.append("## Best parameters\n")
    parts.append(_render_best_params_table(best_params))
    parts.append("")
    posterior_block = _render_posterior_section(summary, refs)
    if posterior_block:
        parts.append(posterior_block)
    parts.append(
        "_Generated by the swmm-calibration candidate writer; consult "
        "`candidate_calibration.json` for machine-readable evidence._\n"
    )
    return "\n".join(parts)


# ---------------------------------------------------------------------------
# Top-level orchestrator
# ---------------------------------------------------------------------------


def _build_candidate_payload(
    *,
    summary: Mapping[str, Any],
    best_params: Mapping[str, Any],
    patch_sha256: str,
    canonical_inp: Path,
    canonical_inp_sha256: str,
    extra_refs: Mapping[str, str],
) -> dict[str, Any]:
    """Assemble the dict that gets written to candidate_calibration.json."""
    payload: dict[str, Any] = {
        "schema_version": CANDIDATE_SCHEMA,
        "evidence_boundary": EVIDENCE_BOUNDARY,
        "generated_at_utc": _utc_now_iso(),
        "strategy": summary.get("strategy"),
        "primary_objective": summary.get("primary_objective", "kge"),
        "primary_value": summary.get("primary_value"),
        "iterations": summary.get("iterations"),
        "kge_decomposition": summary.get("kge_decomposition"),
        "secondary_metrics": summary.get("secondary_metrics"),
        "best_params": dict(best_params),
        "candidate_inp_patch_ref": PATCH_FILENAME,
        "candidate_inp_patch_sha256": patch_sha256,
        "canonical_inp_ref": str(canonical_inp),
        "canonical_inp_sha256_at_candidate_time": canonical_inp_sha256,
        "convergence_trace_ref": (
            extra_refs.get("convergence_csv")
            or summary.get("convergence_trace_ref")
        ),
    }
    # DREAM-only extras — only embed if present in the summary so the
    # candidate file stays minimal for SCE-UA.
    if "posterior_summary" in summary:
        payload["posterior_summary"] = summary["posterior_summary"]
    posterior_samples_ref = extra_refs.get("posterior_samples_csv")
    if posterior_samples_ref:
        payload["posterior_samples_ref"] = posterior_samples_ref
    posterior_correlation_ref = extra_refs.get("posterior_correlation_png")
    if posterior_correlation_ref:
        payload["posterior_correlation_ref"] = posterior_correlation_ref
    return payload


def write_candidate_artefacts(
    *,
    run_dir: Path,
    canonical_inp: Path,
    patch_map: Mapping[str, Mapping[str, Any]],
    best_params: Mapping[str, Any],
    summary: Mapping[str, Any],
    extra_refs: Mapping[str, str] | None = None,
) -> dict[str, str]:
    """Emit the three candidate artefacts into ``<run_dir>/09_audit/``.

    The canonical INP at ``canonical_inp`` is **read** to extract the
    old values for the patch diff; this function never writes to it.
    """
    refs: dict[str, str] = dict(extra_refs or {})
    audit_dir = Path(run_dir) / "09_audit"
    audit_dir.mkdir(parents=True, exist_ok=True)

    canonical_inp_text = Path(canonical_inp).read_text(errors="ignore")
    canonical_inp_sha = sha256_of_file(Path(canonical_inp))

    patch_payload = build_inp_patch(canonical_inp_text, patch_map, best_params)
    patch_path = audit_dir / PATCH_FILENAME
    _write_json_canonical(patch_path, patch_payload)
    # Hash the on-disk text (without trailing newline) so the candidate
    # and the recompute by ``aiswmm calibration accept`` line up.
    patch_sha = sha256_of_canonical_json(patch_payload)

    candidate_payload = _build_candidate_payload(
        summary=summary,
        best_params=best_params,
        patch_sha256=patch_sha,
        canonical_inp=canonical_inp,
        canonical_inp_sha256=canonical_inp_sha,
        extra_refs=refs,
    )
    candidate_path = audit_dir / CANDIDATE_FILENAME
    _write_json_canonical(candidate_path, candidate_payload)

    report_text = render_calibration_report(
        summary=summary,
        best_params=best_params,
        candidate_inp_patch_sha256=patch_sha,
        refs=refs,
    )
    report_path = audit_dir / REPORT_FILENAME
    report_path.write_text(report_text, encoding="utf-8")

    return {
        "candidate_path": str(candidate_path),
        "patch_path": str(patch_path),
        "report_path": str(report_path),
        "candidate_inp_patch_sha256": patch_sha,
        "canonical_inp_sha256": canonical_inp_sha,
    }


# ---------------------------------------------------------------------------
# Reader helpers (used by ``aiswmm calibration accept``)
# ---------------------------------------------------------------------------


def read_candidate(run_dir: Path) -> dict[str, Any]:
    """Return the parsed ``candidate_calibration.json`` for ``run_dir``.

    Raises ``FileNotFoundError`` if it is missing — callers must catch
    this and turn it into the accept CLI's "no candidate" refusal.
    """
    path = Path(run_dir) / "09_audit" / CANDIDATE_FILENAME
    if not path.is_file():
        raise FileNotFoundError(str(path))
    return json.loads(path.read_text(encoding="utf-8"))


def read_patch(run_dir: Path) -> dict[str, Any]:
    """Return the parsed ``candidate_inp_patch.json`` for ``run_dir``.

    Raises ``FileNotFoundError`` if it is missing.
    """
    path = Path(run_dir) / "09_audit" / PATCH_FILENAME
    if not path.is_file():
        raise FileNotFoundError(str(path))
    return json.loads(path.read_text(encoding="utf-8"))


def patch_to_params(patch: Mapping[str, Any]) -> dict[str, Any]:
    """Convert ``candidate_inp_patch.json`` into ``{param: new_value}``.

    The result is the input shape :func:`inp_patch.patch_inp_text`
    expects; the accept CLI feeds it straight in.
    """
    edits: Iterable[Mapping[str, Any]] = patch.get("edits") or []
    return {edit["param"]: edit["new_value"] for edit in edits}

#!/usr/bin/env python3
"""DREAM-ZS (DiffeRential Evolution Adaptive Metropolis) Bayesian calibration for SWMM.

Wraps ``spotpy.algorithms.dream`` around the existing patch-and-run pipeline used
by ``swmm_calibrate.py``. Treats DREAM as a posterior sampler over the same
parameter bounds the SCE-UA wrapper consumes, but reports a posterior summary
in addition to the MAP-estimate ``calibration_summary.json``.

Likelihood (issue #53):

    L(theta) = exp(-0.5 * (1 - KGE(theta)) / sigma^2)

spotpy DREAM minimises ``-log L`` via its ``acceptance_test_option=6`` Metropolis
ratio, so we return the log-likelihood ``-0.5 * (1 - KGE) / sigma^2``. Combined
with a uniform prior implied by the parameter bounds, samples are draws from
the posterior.

Outputs (under the audit dir chosen by the caller, defaults to
``<summary-parent>/`` when not provided explicitly):

  * ``posterior_samples.csv``       — all post-burn-in MCMC samples
  * ``best_params.json``            — MAP (highest-likelihood) parameter set
  * ``chain_convergence.json``      — Gelman-Rubin Rhat per parameter
  * ``posterior_<param>.png``       — per-parameter marginal histogram
  * ``posterior_correlation.png``   — parameter correlation matrix
  * ``calibration_summary.json``    — Slice 1 shape + ``posterior_summary``

Why a separate module: keeping spotpy.algorithms.dream off the import path of
the main CLI means existing strategies still work when spotpy is not installed.
``swmm_calibrate.py`` imports this module only when the ``--strategy dream-zs``
branch is taken.
"""

from __future__ import annotations

import csv
import json
import math
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any, Callable, Sequence

import numpy as np
import pandas as pd

from metrics import align_series, compute_metrics, kge
from inp_patch import patch_inp_text
from sceua import (
    REQUIRED_SECONDARY_KEYS,
    build_calibration_summary,
    secondary_metrics_from_bundle,
)


STRATEGY_NAME = "dream-zs"
PRIMARY_OBJECTIVE_NAME = "kge"
DEFAULT_SIGMA = 0.1  # likelihood width on (1 - KGE); 0.1 keeps top-decile fits informative
DEFAULT_RHAT_THRESHOLD = 1.2


# ---------------------------------------------------------------------------
# Config dataclass
# ---------------------------------------------------------------------------


@dataclass(frozen=True)
class DreamZsConfig:
    base_inp: Path
    patch_map: dict
    observed: pd.DataFrame
    run_root: Path
    swmm_node: str
    swmm_attr: str
    aggregate: str
    obs_start: str | None
    obs_end: str | None
    bounds: dict  # name -> ParamBound
    iterations: int
    seed: int
    n_chains: int
    sigma: float
    rhat_threshold: float
    output_dir: Path
    swmm_runner: Callable
    extract_series: Callable
    runs_after_convergence: int = 50


# ---------------------------------------------------------------------------
# Spotpy setup adapter
# ---------------------------------------------------------------------------


class _SwmmDreamSetup:
    """Spotpy setup adapter that returns a Bayesian log-likelihood."""

    def __init__(self, config: DreamZsConfig) -> None:
        self.config = config
        import spotpy  # local import keeps non-dream paths free of the dependency

        self._spotpy = spotpy
        self._param_order = list(config.bounds.keys())
        self._params = [
            spotpy.parameter.Uniform(
                name,
                low=config.bounds[name].min_value,
                high=config.bounds[name].max_value,
                optguess=(config.bounds[name].min_value + config.bounds[name].max_value) / 2.0,
            )
            for name in self._param_order
        ]
        obs_clean = config.observed.copy()
        obs_clean["timestamp"] = pd.to_datetime(obs_clean["timestamp"])
        self._observed = obs_clean.sort_values("timestamp").reset_index(drop=True)
        self._call_count = 0

    def parameters(self):
        return self._spotpy.parameter.generate(self._params)

    def _values_to_named(self, values: Sequence[float]) -> dict[str, float | int]:
        out: dict[str, float | int] = {}
        for i, name in enumerate(self._param_order):
            bound = self.config.bounds[name]
            raw = float(values[i])
            if bound.value_type == "int":
                out[name] = int(round(raw))
            elif bound.precision is not None:
                out[name] = round(raw, bound.precision)
            else:
                out[name] = raw
        return out

    def _run_swmm_for_params(self, params: dict[str, float | int]) -> pd.DataFrame | None:
        cfg = self.config
        self._call_count += 1
        trial_dir = cfg.run_root / f"dream_{self._call_count:04d}"
        trial_dir.mkdir(parents=True, exist_ok=True)
        try:
            patched_text = patch_inp_text(
                cfg.base_inp.read_text(errors="ignore"),
                cfg.patch_map,
                params,
            )
        except Exception:
            return None
        inp = trial_dir / "model.inp"
        inp.write_text(patched_text, encoding="utf-8")
        try:
            rc, _, out_path = cfg.swmm_runner(inp, trial_dir)
        except FileNotFoundError:
            return None
        except Exception:
            return None
        if rc != 0:
            return None
        try:
            sim = cfg.extract_series(out_path)
        except Exception:
            return None
        return sim

    def simulation(self, values) -> np.ndarray:
        params = self._values_to_named(values)
        sim_df = self._run_swmm_for_params(params)
        if sim_df is None or sim_df.empty:
            return np.full(len(self._observed), np.nan, dtype=float)
        aligned = align_series(self._observed, sim_df)
        if aligned.empty:
            return np.full(len(self._observed), np.nan, dtype=float)
        obs_ts = self._observed["timestamp"].to_numpy()
        sim_map = dict(
            zip(aligned["timestamp"].to_numpy(), aligned["flow_sim"].astype(float).to_numpy())
        )
        return np.array([sim_map.get(ts, np.nan) for ts in obs_ts], dtype=float)

    def evaluation(self) -> np.ndarray:
        return self._observed["flow"].astype(float).to_numpy()

    def objectivefunction(self, simulation, evaluation, params=None) -> float:
        """Return the log-likelihood ``-0.5 * (1 - KGE) / sigma^2``.

        DREAM with ``acceptance_test_option=6`` accepts a proposal when the
        returned ``like`` is greater than the chain's current best, so we
        encode "better fit -> larger likelihood" by returning the log-density
        directly. Failed simulations get a very negative likelihood so they
        cannot win acceptance.
        """

        sim = np.asarray(simulation, dtype=float)
        obs = np.asarray(evaluation, dtype=float)
        mask = np.isfinite(sim) & np.isfinite(obs)
        if mask.sum() < 2:
            return -1.0e9
        sim_ok = sim[mask]
        obs_ok = obs[mask]
        ts = self._observed["timestamp"].to_numpy()[mask]
        sim_df = pd.DataFrame({"timestamp": ts, "flow": sim_ok})
        obs_df = pd.DataFrame({"timestamp": ts, "flow": obs_ok})
        kge_result = kge(obs_df, sim_df)
        kge_value = kge_result["kge"]
        if kge_value is None or not math.isfinite(kge_value):
            return -1.0e9
        sigma = self.config.sigma
        return float(-0.5 * (1.0 - kge_value) / (sigma * sigma))

    @property
    def param_order(self) -> list[str]:
        return list(self._param_order)


# ---------------------------------------------------------------------------
# Posterior post-processing
# ---------------------------------------------------------------------------


def _extract_chain_field(results) -> np.ndarray | None:
    """Return the per-row chain index array from spotpy's result table, if present."""

    if results is None or len(results) == 0:
        return None
    if "chain" in results.dtype.names:
        return np.asarray(results["chain"]).astype(int)
    return None


def _split_burn_in(results, n_chains: int) -> tuple[np.ndarray, np.ndarray]:
    """Drop the regular-startpoint initialisation rows and return (post-burnin rows, chain ids).

    spotpy.algorithms.dream writes ``nChains`` initialisation rows (one per chain)
    before the random walk begins. We strip those rows from the posterior.
    """

    chains = _extract_chain_field(results)
    if chains is None:
        # Fall back to row-order interpretation: assume rows interleaved by chain.
        n = len(results)
        idx = np.arange(n)
        if n > n_chains:
            keep = idx >= n_chains
            return idx[keep], (idx[keep] % n_chains)
        return idx, idx % n_chains
    # Burn-in: per chain, drop the first sample.
    keep_mask = np.ones(len(results), dtype=bool)
    for ch in range(n_chains):
        ch_idx = np.where(chains == ch)[0]
        if ch_idx.size:
            keep_mask[ch_idx[0]] = False
    keep = np.where(keep_mask)[0]
    return keep, chains[keep]


def write_posterior_samples_csv(
    results,
    param_order: Sequence[str],
    n_chains: int,
    csv_path: Path,
) -> int:
    """Write the post-burn-in MCMC samples to CSV. Returns number of rows written."""

    keep_idx, chain_ids = _split_burn_in(results, n_chains)
    csv_path.parent.mkdir(parents=True, exist_ok=True)
    header = ["chain", "iteration_in_chain", "likelihood", *param_order]
    chain_seen: dict[int, int] = {}
    with csv_path.open("w", newline="") as fh:
        writer = csv.writer(fh)
        writer.writerow(header)
        for row_pos, idx in enumerate(keep_idx):
            row = results[idx]
            ch = int(chain_ids[row_pos])
            chain_seen[ch] = chain_seen.get(ch, -1) + 1
            iter_in_chain = chain_seen[ch]
            like_val = float(row["like1"]) if "like1" in results.dtype.names else float("nan")
            param_values = []
            for name in param_order:
                col = f"par{name}"
                if col in results.dtype.names:
                    param_values.append(float(row[col]))
                else:
                    param_values.append(float("nan"))
            writer.writerow(
                [
                    ch,
                    iter_in_chain,
                    f"{like_val:.10g}",
                    *[f"{v:.10g}" for v in param_values],
                ]
            )
    return len(keep_idx)


def _gelman_rubin(chains: np.ndarray) -> float:
    """Compute the Gelman-Rubin Rhat for a single parameter.

    ``chains`` has shape (n_chains, n_samples).
    """

    m, n = chains.shape
    if m < 2 or n < 2:
        return float("nan")
    chain_means = chains.mean(axis=1)
    chain_vars = chains.var(axis=1, ddof=1)
    grand_mean = chain_means.mean()
    B = (n / (m - 1)) * np.sum((chain_means - grand_mean) ** 2)
    W = chain_vars.mean()
    if W <= 0.0 or not np.isfinite(W):
        return float("nan")
    var_hat = ((n - 1) / n) * W + (1.0 / n) * B
    rhat = math.sqrt(var_hat / W)
    if not math.isfinite(rhat):
        return float("nan")
    return float(rhat)


def compute_rhat(
    results,
    param_order: Sequence[str],
    n_chains: int,
) -> dict[str, float]:
    """Return Gelman-Rubin Rhat per parameter, using post-burn-in samples."""

    keep_idx, chain_ids = _split_burn_in(results, n_chains)
    rhat: dict[str, float] = {}
    for name in param_order:
        col = f"par{name}"
        if col not in results.dtype.names:
            rhat[name] = float("nan")
            continue
        per_chain: list[list[float]] = [[] for _ in range(n_chains)]
        for row_pos, idx in enumerate(keep_idx):
            ch = int(chain_ids[row_pos])
            if 0 <= ch < n_chains:
                per_chain[ch].append(float(results[idx][col]))
        min_len = min((len(c) for c in per_chain), default=0)
        if min_len < 2:
            rhat[name] = float("nan")
            continue
        arr = np.array([c[:min_len] for c in per_chain], dtype=float)
        rhat[name] = _gelman_rubin(arr)
    return rhat


def write_chain_convergence_json(
    rhat: dict[str, float],
    threshold: float,
    n_chains: int,
    iterations: int,
    json_path: Path,
) -> bool:
    """Write chain_convergence.json. Returns the converged flag.

    Per-parameter Rhat is emitted as numeric where finite and ``null`` where the
    Gelman-Rubin diagnostic could not be computed (e.g. zero between-chain
    variance after a very short run).
    """

    finite_vals = [v for v in rhat.values() if isinstance(v, (int, float)) and math.isfinite(v)]
    converged = bool(finite_vals) and all(v < threshold for v in finite_vals)
    payload = {
        "rhat": {
            name: (float(v) if isinstance(v, (int, float)) and math.isfinite(v) else None)
            for name, v in rhat.items()
        },
        "threshold": float(threshold),
        "converged": converged,
        "n_chains": int(n_chains),
        "iterations": int(iterations),
    }
    json_path.parent.mkdir(parents=True, exist_ok=True)
    json_path.write_text(json.dumps(payload, indent=2), encoding="utf-8")
    return converged


def _matplotlib_backend():
    """Return matplotlib.pyplot with a non-interactive backend forced."""

    import matplotlib

    matplotlib.use("Agg", force=True)
    import matplotlib.pyplot as plt  # noqa: WPS433

    return plt


def write_marginal_histograms(
    results,
    param_order: Sequence[str],
    n_chains: int,
    output_dir: Path,
) -> list[Path]:
    """Plot one marginal histogram per parameter using post-burn-in samples."""

    plt = _matplotlib_backend()
    keep_idx, _ = _split_burn_in(results, n_chains)
    output_dir.mkdir(parents=True, exist_ok=True)
    written: list[Path] = []
    for name in param_order:
        col = f"par{name}"
        if col not in results.dtype.names:
            continue
        vals = np.array([float(results[idx][col]) for idx in keep_idx], dtype=float)
        vals = vals[np.isfinite(vals)]
        if vals.size == 0:
            continue
        fig, ax = plt.subplots(figsize=(5.5, 3.5))
        bins = max(8, min(40, int(math.sqrt(vals.size))))
        ax.hist(vals, bins=bins, color="#4878D0", edgecolor="white", alpha=0.85)
        ax.set_xlabel(name)
        ax.set_ylabel("count")
        ax.set_title(f"Posterior marginal: {name}")
        fig.tight_layout()
        path = output_dir / f"posterior_{name}.png"
        fig.savefig(path, dpi=120)
        plt.close(fig)
        written.append(path)
    return written


def write_correlation_plot(
    results,
    param_order: Sequence[str],
    n_chains: int,
    output_path: Path,
) -> None:
    """Plot the posterior parameter correlation matrix as a heatmap."""

    plt = _matplotlib_backend()
    keep_idx, _ = _split_burn_in(results, n_chains)
    if not param_order:
        return
    columns: list[np.ndarray] = []
    used_names: list[str] = []
    for name in param_order:
        col = f"par{name}"
        if col not in results.dtype.names:
            continue
        vals = np.array([float(results[idx][col]) for idx in keep_idx], dtype=float)
        columns.append(vals)
        used_names.append(name)
    if len(columns) < 1:
        return
    mat = np.vstack(columns)
    if mat.shape[1] < 2:
        # Degenerate: emit a one-cell heatmap so the artefact still exists.
        corr = np.array([[1.0]])
        used_names = used_names[:1] or ["param"]
    else:
        # Replace columns with NaN variance with their mean to avoid all-NaN rows.
        finite_mask = np.isfinite(mat).all(axis=0)
        mat = mat[:, finite_mask]
        if mat.shape[1] < 2:
            corr = np.array([[1.0]])
            used_names = used_names[:1] or ["param"]
        else:
            try:
                corr = np.corrcoef(mat)
            except Exception:
                corr = np.eye(len(used_names))
    fig, ax = plt.subplots(figsize=(1.2 * len(used_names) + 1.5, 1.0 * len(used_names) + 1.2))
    im = ax.imshow(corr, vmin=-1, vmax=1, cmap="RdBu_r", aspect="auto")
    ax.set_xticks(range(len(used_names)))
    ax.set_yticks(range(len(used_names)))
    ax.set_xticklabels(used_names, rotation=45, ha="right")
    ax.set_yticklabels(used_names)
    for i in range(len(used_names)):
        for j in range(len(used_names)):
            val = corr[i, j] if corr.shape == (len(used_names), len(used_names)) else 1.0
            ax.text(j, i, f"{val:.2f}", ha="center", va="center", color="black", fontsize=8)
    ax.set_title("Posterior parameter correlation")
    fig.colorbar(im, ax=ax, fraction=0.046, pad=0.04)
    fig.tight_layout()
    output_path.parent.mkdir(parents=True, exist_ok=True)
    fig.savefig(output_path, dpi=120)
    plt.close(fig)


def _pick_map_row(results) -> tuple[int, float] | None:
    """Pick the highest-likelihood (MAP) row index and its likelihood value."""

    if results is None or len(results) == 0 or "like1" not in results.dtype.names:
        return None
    likes = np.asarray(results["like1"], dtype=float)
    if not np.isfinite(likes).any():
        return None
    # Mask non-finite (failed sims) so they cannot be chosen.
    masked = np.where(np.isfinite(likes), likes, -np.inf)
    best_idx = int(np.argmax(masked))
    return best_idx, float(likes[best_idx])


def _summarise_posterior(results, param_order: Sequence[str], n_chains: int) -> dict[str, Any]:
    keep_idx, _ = _split_burn_in(results, n_chains)
    out: dict[str, Any] = {
        "n_chains": int(n_chains),
        "n_samples_post_burnin": int(len(keep_idx)),
        "per_parameter": {},
    }
    for name in param_order:
        col = f"par{name}"
        if col not in results.dtype.names:
            continue
        vals = np.array([float(results[idx][col]) for idx in keep_idx], dtype=float)
        vals = vals[np.isfinite(vals)]
        if vals.size == 0:
            out["per_parameter"][name] = {
                "mean": None,
                "median": None,
                "std": None,
                "q05": None,
                "q95": None,
            }
            continue
        out["per_parameter"][name] = {
            "mean": float(vals.mean()),
            "median": float(np.median(vals)),
            "std": float(vals.std(ddof=1)) if vals.size > 1 else 0.0,
            "q05": float(np.quantile(vals, 0.05)),
            "q95": float(np.quantile(vals, 0.95)),
        }
    return out


# ---------------------------------------------------------------------------
# Public entry point
# ---------------------------------------------------------------------------


def run_dream_zs(config: DreamZsConfig) -> dict[str, Any]:
    """Run DREAM-ZS, write all 5 posterior artefacts + summary, return a results dict."""

    import spotpy  # noqa: WPS433
    from spotpy.algorithms import dream as dream_algo

    setup = _SwmmDreamSetup(config)
    sampler = dream_algo(
        setup,
        dbname=str(config.run_root / "dream_db"),
        dbformat="ram",
        random_state=config.seed,
        save_sim=False,
    )
    # spotpy.algorithms.dream has two structural constraints we must satisfy:
    #   1. ``nChains >= 2*delta + 1`` so each step has ``delta`` partner chains
    #      to form the differential proposal.
    #   2. ``get_r_hat`` only returns a numeric Rhat when ``nChains > 3``;
    #      with three or fewer chains it returns ``None`` and spotpy's
    #      convergence-limit check raises ``TypeError``.
    # When callers ask for fewer chains we lift the effective count to the
    # smallest value that lets the algorithm physically run; the requested
    # count is preserved in ``posterior_summary.n_chains_requested``.
    requested_chains = int(config.n_chains)
    delta = 3 if requested_chains >= 7 else 1
    min_chains_for_rhat = 4
    effective_chains = max(requested_chains, 2 * delta + 1, min_chains_for_rhat)
    sampler.sample(
        repetitions=config.iterations,
        nChains=effective_chains,
        delta=delta,
        convergence_limit=config.rhat_threshold,
        runs_after_convergence=config.runs_after_convergence,
    )

    results = sampler.getdata()
    if results is None or len(results) == 0:
        raise RuntimeError("DREAM-ZS results database was empty.")

    output_dir = config.output_dir
    output_dir.mkdir(parents=True, exist_ok=True)
    posterior_csv = output_dir / "posterior_samples.csv"
    convergence_json = output_dir / "chain_convergence.json"
    correlation_png = output_dir / "posterior_correlation.png"

    post_rows = write_posterior_samples_csv(
        results=results,
        param_order=setup.param_order,
        n_chains=effective_chains,
        csv_path=posterior_csv,
    )
    rhat = compute_rhat(results, setup.param_order, effective_chains)
    converged = write_chain_convergence_json(
        rhat=rhat,
        threshold=config.rhat_threshold,
        n_chains=effective_chains,
        iterations=config.iterations,
        json_path=convergence_json,
    )
    write_marginal_histograms(
        results=results,
        param_order=setup.param_order,
        n_chains=effective_chains,
        output_dir=output_dir,
    )
    write_correlation_plot(
        results=results,
        param_order=setup.param_order,
        n_chains=effective_chains,
        output_path=correlation_png,
    )

    # MAP estimate: highest-likelihood row in the database.
    best = _pick_map_row(results)
    if best is None:
        raise RuntimeError("DREAM-ZS could not identify a MAP-estimate row (no finite likelihoods).")
    best_idx, _ = best

    best_params: dict[str, float | int] = {}
    for name in setup.param_order:
        col = f"par{name}"
        if col in results.dtype.names:
            value = float(results[best_idx][col])
            bound = config.bounds[name]
            if bound.value_type == "int":
                best_params[name] = int(round(value))
            elif bound.precision is not None:
                best_params[name] = round(value, bound.precision)
            else:
                best_params[name] = value

    # Re-run SWMM at MAP to compute Slice 1 metrics & decomposition.
    cfg = config
    final_trial_dir = cfg.run_root / "dream_map"
    final_trial_dir.mkdir(parents=True, exist_ok=True)
    patched_text = patch_inp_text(
        cfg.base_inp.read_text(errors="ignore"),
        cfg.patch_map,
        best_params,
    )
    final_inp = final_trial_dir / "model.inp"
    final_inp.write_text(patched_text, encoding="utf-8")
    rc, _, out_path = cfg.swmm_runner(final_inp, final_trial_dir)
    if rc != 0:
        raise RuntimeError(f"swmm5 MAP run failed with rc={rc}")
    sim_df = cfg.extract_series(out_path)
    aligned = align_series(setup._observed, sim_df)
    metrics_bundle = compute_metrics(setup._observed, sim_df)
    kge_block = kge(setup._observed, sim_df)
    if kge_block["decomposition"] is None or kge_block["kge"] is None:
        raise RuntimeError("KGE undefined for the MAP parameter set; cannot summarise.")

    obs_for_pbias = (
        aligned["flow_obs"].astype(float)
        if not aligned.empty
        else setup._observed["flow"].astype(float)
    )
    secondary = secondary_metrics_from_bundle(metrics_bundle, obs_for_pbias)

    summary = build_calibration_summary(
        primary_value=kge_block["kge"],
        kge_decomposition=kge_block["decomposition"],
        secondary_metrics=secondary,
        iterations=config.iterations,
        convergence_trace_ref=convergence_json.name,
    )
    # Overwrite strategy name from SCE-UA default to DREAM-ZS.
    summary["strategy"] = STRATEGY_NAME

    # DREAM-specific posterior summary block (additional, non-breaking).
    posterior_summary = _summarise_posterior(results, setup.param_order, effective_chains)
    posterior_summary.update(
        {
            "n_chains": int(effective_chains),
            "n_chains_requested": int(requested_chains),
            "converged": bool(converged),
            "rhat_threshold": float(config.rhat_threshold),
            "rhat": {name: (float(v) if isinstance(v, (int, float)) and math.isfinite(v) else None)
                     for name, v in rhat.items()},
            "sigma": float(config.sigma),
            "n_samples_total": int(len(results)),
            "posterior_csv_ref": posterior_csv.name,
            "correlation_png_ref": correlation_png.name,
        }
    )
    summary["posterior_summary"] = posterior_summary

    return {
        "summary": summary,
        "best_params": best_params,
        "posterior_samples_csv": str(posterior_csv),
        "chain_convergence_json": str(convergence_json),
        "correlation_png": str(correlation_png),
        "total_calls": setup._call_count,
        "metrics_bundle": metrics_bundle.to_dict(),
        "post_burnin_rows": int(post_rows),
        "rhat": rhat,
        "converged": bool(converged),
    }


if __name__ == "__main__":
    # CLI usage is via swmm_calibrate.py search --strategy dream-zs.
    import sys

    print(
        "dream_zs.py is a library module; invoke via 'swmm_calibrate.py search --strategy dream-zs'.",
        file=sys.stderr,
    )
    sys.exit(0)

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

import argparse
import json
from pathlib import Path


def patch_inp_text(text: str, patch_map: dict, params: dict) -> str:
    lines = text.splitlines()
    current_section = None
    touched = set()

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

        code, *comment = raw.split(";", 1)
        tokens = code.split()
        if not tokens:
            continue

        for key, value in params.items():
            spec = patch_map.get(key)
            if not spec:
                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 {key} on line: {raw}")
            tokens[idx] = str(value)
            new_code = "  ".join(tokens)
            if comment:
                new_code += " ;" + comment[0]
            lines[i] = new_code
            touched.add(key)

    missing = sorted(set(params) - touched)
    if missing:
        raise KeyError(f"Did not patch parameter(s): {missing}")
    return "\n".join(lines) + "\n"


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--inp", required=True, type=Path)
    ap.add_argument("--patch-map", required=True, type=Path)
    ap.add_argument("--params", required=True, type=Path, help="JSON object of parameter values")
    ap.add_argument("--out", required=True, type=Path)
    args = ap.parse_args()

    patch_map = json.loads(args.patch_map.read_text())
    params = json.loads(args.params.read_text())
    patched = patch_inp_text(args.inp.read_text(errors="ignore"), patch_map, params)
    args.out.parent.mkdir(parents=True, exist_ok=True)
    args.out.write_text(patched, encoding="utf-8")
    print(json.dumps({"ok": True, "out": str(args.out)}, indent=2))


if __name__ == "__main__":
    main()

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

from dataclasses import dataclass, field
from datetime import datetime
from typing import Iterable

import math
import pandas as pd


@dataclass
class MetricBundle:
    count: int
    nse: float | None
    rmse: float | None
    bias: float | None
    peak_flow_error: float | None
    peak_timing_error_minutes: float | None
    kge: float | None = None
    kge_decomposition: dict | None = None

    def to_dict(self) -> dict:
        return {
            "count": self.count,
            "nse": self.nse,
            "rmse": self.rmse,
            "bias": self.bias,
            "peak_flow_error": self.peak_flow_error,
            "peak_timing_error_minutes": self.peak_timing_error_minutes,
            "kge": self.kge,
            "kge_decomposition": self.kge_decomposition,
        }


def align_series(observed: pd.DataFrame, simulated: pd.DataFrame) -> pd.DataFrame:
    obs = observed[["timestamp", "flow"]].copy()
    sim = simulated[["timestamp", "flow"]].copy()
    obs["timestamp"] = pd.to_datetime(obs["timestamp"])
    sim["timestamp"] = pd.to_datetime(sim["timestamp"])
    out = pd.merge(obs, sim, on="timestamp", how="inner", suffixes=("_obs", "_sim"))
    out = out.dropna(subset=["flow_obs", "flow_sim"]).sort_values("timestamp")
    return out.reset_index(drop=True)


def _safe_div(num: float, den: float) -> float | None:
    if den == 0:
        return None
    return num / den


def _kge_components(sim: pd.Series, obs: pd.Series) -> dict | None:
    """Return Kling-Gupta r/alpha/beta decomposition or None if undefined."""

    if len(sim) < 2 or len(obs) < 2:
        return None
    sim_mean = float(sim.mean())
    obs_mean = float(obs.mean())
    if obs_mean == 0.0:
        return None
    sim_std = float(sim.std(ddof=0))
    obs_std = float(obs.std(ddof=0))
    if obs_std == 0.0:
        return None
    # Pearson correlation; if either series has zero variance, corrcoef gives NaN.
    corr_matrix = sim.to_numpy().reshape(-1)
    obs_array = obs.to_numpy().reshape(-1)
    # numpy.corrcoef returns NaN if either input has zero variance.
    import numpy as np  # local import keeps module load cheap

    if sim_std == 0.0:
        # When the simulated series is constant, correlation is undefined.
        r = 0.0
    else:
        r = float(np.corrcoef(corr_matrix, obs_array)[0, 1])
        if not math.isfinite(r):
            r = 0.0
    alpha = sim_std / obs_std
    beta = sim_mean / obs_mean
    return {"r": r, "alpha": alpha, "beta": beta}


def kge(observed: pd.DataFrame, simulated: pd.DataFrame) -> dict:
    """Compute Kling-Gupta efficiency and its (r, alpha, beta) decomposition.

    Reference: Gupta et al. (2009), Decomposition of the mean squared error
    and NSE performance criteria. KGE = 1 - sqrt((r-1)^2 + (alpha-1)^2 + (beta-1)^2)
    where r is Pearson correlation, alpha = std(sim)/std(obs), beta = mean(sim)/mean(obs).
    A perfect simulation has KGE = 1.
    """

    aligned = align_series(observed, simulated)
    if aligned.empty:
        return {"kge": None, "decomposition": None}
    sim = aligned["flow_sim"].astype(float)
    obs = aligned["flow_obs"].astype(float)
    decomposition = _kge_components(sim, obs)
    if decomposition is None:
        return {"kge": None, "decomposition": None}
    r = decomposition["r"]
    alpha = decomposition["alpha"]
    beta = decomposition["beta"]
    kge_value = 1.0 - math.sqrt((r - 1.0) ** 2 + (alpha - 1.0) ** 2 + (beta - 1.0) ** 2)
    return {"kge": float(kge_value), "decomposition": {"r": r, "alpha": alpha, "beta": beta}}


def compute_metrics(observed: pd.DataFrame, simulated: pd.DataFrame) -> MetricBundle:
    aligned = align_series(observed, simulated)
    n = int(len(aligned))
    if n == 0:
        return MetricBundle(0, None, None, None, None, None, None, None)

    obs = aligned["flow_obs"].astype(float)
    sim = aligned["flow_sim"].astype(float)
    diff = sim - obs

    rmse = float(math.sqrt((diff.pow(2).mean())))
    bias = float(diff.mean())

    den = float(((obs - obs.mean()) ** 2).sum())
    nse = None if den == 0 else float(1.0 - ((diff.pow(2).sum()) / den))

    peak_obs_idx = int(obs.idxmax())
    peak_sim_idx = int(sim.idxmax())
    peak_obs = float(obs.iloc[peak_obs_idx])
    peak_sim = float(sim.iloc[peak_sim_idx])
    peak_flow_error = _safe_div(peak_sim - peak_obs, peak_obs)

    t_obs = pd.Timestamp(aligned.loc[peak_obs_idx, "timestamp"])
    t_sim = pd.Timestamp(aligned.loc[peak_sim_idx, "timestamp"])
    peak_timing_error_minutes = float((t_sim - t_obs).total_seconds() / 60.0)

    decomposition = _kge_components(sim, obs)
    if decomposition is None:
        kge_value: float | None = None
    else:
        r = decomposition["r"]
        alpha = decomposition["alpha"]
        beta = decomposition["beta"]
        kge_value = float(1.0 - math.sqrt((r - 1.0) ** 2 + (alpha - 1.0) ** 2 + (beta - 1.0) ** 2))

    return MetricBundle(
        n,
        nse,
        rmse,
        bias,
        peak_flow_error,
        peak_timing_error_minutes,
        kge_value,
        decomposition,
    )


def score_from_metrics(metrics: MetricBundle, objective: str) -> float:
    obj = objective.lower().strip()
    if obj == "nse":
        return float("-inf") if metrics.nse is None else metrics.nse
    if obj == "kge":
        return float("-inf") if metrics.kge is None else metrics.kge
    if obj == "rmse":
        return float("inf") if metrics.rmse is None else -metrics.rmse
    if obj == "bias":
        return float("inf") if metrics.bias is None else -abs(metrics.bias)
    if obj == "peak_flow_error":
        return float("inf") if metrics.peak_flow_error is None else -abs(metrics.peak_flow_error)
    if obj == "peak_timing_error":
        return float("inf") if metrics.peak_timing_error_minutes is None else -abs(metrics.peak_timing_error_minutes)
    raise ValueError(f"Unsupported objective: {objective}")

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

import argparse
import io
import json
from pathlib import Path

import pandas as pd

TIME_CANDIDATES = ["timestamp", "time", "datetime", "date_time", "date"]
FLOW_CANDIDATES = ["flow", "discharge", "q", "value"]


def _normalize(name: str) -> str:
    return name.strip().lower().replace(" ", "_")


def _non_comment_lines(path: Path) -> list[str]:
    lines: list[str] = []
    for raw in path.read_text(encoding="utf-8", errors="ignore").splitlines():
        stripped = raw.strip()
        if not stripped:
            continue
        if stripped.startswith("#") or stripped.startswith(";;"):
            continue
        lines.append(raw)
    return lines


def detect_delimiter(lines: list[str]) -> str:
    sample = "\n".join(lines[:5])
    if "," in sample:
        return ","
    if "\t" in sample:
        return "\t"
    return r"\s+"


def _looks_headerless_datetime_flow(lines: list[str]) -> bool:
    if not lines:
        return False
    first = lines[0].strip().split()
    if len(first) != 3:
        return False
    dt = pd.to_datetime(f"{first[0]} {first[1]}", errors="coerce")
    flow = pd.to_numeric(first[2], errors="coerce")
    return pd.notna(dt) and pd.notna(flow)


def _read_headerless_datetime_flow(lines: list[str], time_format: str | None = None) -> pd.DataFrame:
    rows = []
    for raw in lines:
        parts = raw.strip().split()
        if len(parts) < 3:
            continue
        dt = pd.to_datetime(f"{parts[0]} {parts[1]}", format=time_format, errors="coerce")
        flow = pd.to_numeric(parts[2], errors="coerce")
        if pd.notna(dt) and pd.notna(flow):
            rows.append((dt, float(flow)))
    if not rows:
        raise ValueError("No valid datetime/flow rows found in headerless observed file")
    out = pd.DataFrame(rows, columns=["timestamp", "flow"])
    return out.sort_values("timestamp").reset_index(drop=True)


def read_series(path: str | Path, timestamp_col: str | None = None, flow_col: str | None = None, time_format: str | None = None) -> pd.DataFrame:
    p = Path(path)
    lines = _non_comment_lines(p)
    if not lines:
        raise ValueError(f"No data rows found in {p}")

    if _looks_headerless_datetime_flow(lines):
        return _read_headerless_datetime_flow(lines, time_format=time_format)

    delim = detect_delimiter(lines)
    df = pd.read_csv(io.StringIO("\n".join(lines)), sep=delim, engine="python")
    if df.empty:
        raise ValueError(f"No data rows found in {p}")

    cols = {c: _normalize(c) for c in df.columns}
    rev = {v: k for k, v in cols.items()}

    if timestamp_col is None:
        # Column inference walks CANDIDATE PRIORITY, not file order:
        # iterating the file's own column order picked whichever
        # candidate appeared first in the header, so a standard
        # hydrometric "Date,Time,Flow" export selected the bare Date
        # column (same-day rows collapsed onto midnight) and
        # "Time,Date,Flow" selected bare times that pandas stamped
        # with TODAY's date (found 2026-08-08, reproduced). A split
        # Date+Time pair is combined into one timestamp — neither
        # half alone is the observation time.
        combined = next(
            (rev[c] for c in ("timestamp", "datetime", "date_time") if c in rev),
            None,
        )
        date_col = rev.get("date")
        time_col = rev.get("time")
        if combined is not None:
            timestamp_col = combined
        elif date_col is not None and time_col is not None:
            combined_name = "__aiswmm_combined_timestamp"
            df[combined_name] = (
                df[date_col].astype(str).str.strip()
                + " "
                + df[time_col].astype(str).str.strip()
            )
            timestamp_col = combined_name
        else:
            timestamp_col = next((rev[c] for c in TIME_CANDIDATES if c in rev), None)
    if flow_col is None:
        flow_col = next((rev[c] for c in FLOW_CANDIDATES if c in rev), None)

    if timestamp_col is None or flow_col is None:
        raise ValueError(f"Could not infer timestamp/flow columns from {list(df.columns)}")

    out = df[[timestamp_col, flow_col]].copy()
    out.columns = ["timestamp", "flow"]
    out["timestamp"] = pd.to_datetime(out["timestamp"], format=time_format, errors="coerce")
    out["flow"] = pd.to_numeric(out["flow"], errors="coerce")
    out = out.dropna(subset=["timestamp", "flow"]).sort_values("timestamp")
    return out.reset_index(drop=True)


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("path")
    ap.add_argument("--timestamp-col", default=None)
    ap.add_argument("--flow-col", default=None)
    ap.add_argument("--time-format", default=None)
    args = ap.parse_args()
    df = read_series(args.path, timestamp_col=args.timestamp_col, flow_col=args.flow_col, time_format=args.time_format)
    print(json.dumps({
        "rows": int(len(df)),
        "start": df.iloc[0]["timestamp"].isoformat() if len(df) else None,
        "end": df.iloc[-1]["timestamp"].isoformat() if len(df) else None,
        "columns": ["timestamp", "flow"],
    }, indent=2))


if __name__ == "__main__":
    main()

#!/usr/bin/env python3
"""SCE-UA (Shuffled Complex Evolution) calibration strategy for SWMM.

Wraps ``spotpy.algorithms.sceua`` around the existing patch-and-run pipeline
used by ``swmm_calibrate.py``. The primary objective is KGE (Kling-Gupta
efficiency); spotpy minimises ``(1 - KGE)``.

Outputs:
  * ``calibration_summary.json`` — shape locked-in by tests/test_calibration_summary_schema.py
  * ``best_params.json`` — JSON object of the best-found parameter values
  * ``convergence.csv`` — per-iteration KGE so reviewers can inspect convergence

Why a separate module: keeping spotpy off the import path of the main CLI
means existing strategies (random / lhs / adaptive) still work when spotpy
is not installed; ``swmm_calibrate.py`` imports this module only when the
``--strategy sceua`` branch is taken.
"""

from __future__ import annotations

import csv
import json
import math
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Callable, Sequence

import numpy as np
import pandas as pd

from metrics import align_series, compute_metrics, kge
from inp_patch import patch_inp_text


# ---------------------------------------------------------------------------
# Schema helpers (pure functions — exercised by test_calibration_summary_schema.py)
# ---------------------------------------------------------------------------

PRIMARY_OBJECTIVE_NAME = "kge"
STRATEGY_NAME = "sceua"

REQUIRED_SECONDARY_KEYS = (
    "nse",
    "pbias_pct",
    "rmse",
    "peak_error_rel",
    "peak_timing_min",
)


def build_calibration_summary(
    primary_value: float,
    kge_decomposition: dict[str, float],
    secondary_metrics: dict[str, float | None],
    iterations: int,
    convergence_trace_ref: str,
) -> dict[str, Any]:
    """Return the calibration_summary.json payload shape defined in issue #48."""

    if not isinstance(primary_value, (int, float)) or not math.isfinite(float(primary_value)):
        raise ValueError(f"primary_value must be a finite number, got {primary_value!r}")
    primary_value_f = float(primary_value)

    decomposition_out = {}
    for key in ("r", "alpha", "beta"):
        if key not in kge_decomposition:
            raise ValueError(f"kge_decomposition missing key {key!r}")
        decomposition_out[key] = float(kge_decomposition[key])

    secondary_out: dict[str, float | None] = {}
    for key in REQUIRED_SECONDARY_KEYS:
        if key not in secondary_metrics:
            raise ValueError(f"secondary_metrics missing key {key!r}")
        value = secondary_metrics[key]
        if value is None:
            secondary_out[key] = None
        else:
            secondary_out[key] = float(value)

    if not isinstance(iterations, int) or iterations < 1:
        raise ValueError(f"iterations must be a positive integer, got {iterations!r}")
    if not isinstance(convergence_trace_ref, str) or not convergence_trace_ref:
        raise ValueError("convergence_trace_ref must be a non-empty string")

    return {
        "primary_objective": PRIMARY_OBJECTIVE_NAME,
        "primary_value": primary_value_f,
        "kge_decomposition": decomposition_out,
        "secondary_metrics": secondary_out,
        "strategy": STRATEGY_NAME,
        "iterations": iterations,
        "convergence_trace_ref": convergence_trace_ref,
    }


def secondary_metrics_from_bundle(metrics_bundle, observed_flow: pd.Series) -> dict[str, float | None]:
    """Pull NSE / PBIAS% / RMSE / peak-flow / peak-timing out of a MetricBundle."""

    nse = metrics_bundle.nse
    rmse = metrics_bundle.rmse
    peak_error_rel = metrics_bundle.peak_flow_error
    peak_timing_min = metrics_bundle.peak_timing_error_minutes
    # PBIAS% = 100 * sum(sim - obs) / sum(obs). Use bias * count / sum(obs).
    obs_sum = float(observed_flow.sum())
    if obs_sum == 0.0 or metrics_bundle.bias is None:
        pbias_pct: float | None = None
    else:
        # bias is mean(sim - obs); total diff = bias * count.
        total_diff = float(metrics_bundle.bias) * float(metrics_bundle.count)
        pbias_pct = float(100.0 * total_diff / obs_sum)
    return {
        "nse": None if nse is None else float(nse),
        "pbias_pct": pbias_pct,
        "rmse": None if rmse is None else float(rmse),
        "peak_error_rel": None if peak_error_rel is None else float(peak_error_rel),
        "peak_timing_min": None if peak_timing_min is None else float(peak_timing_min),
    }


# ---------------------------------------------------------------------------
# spotpy setup wiring
# ---------------------------------------------------------------------------


@dataclass(frozen=True)
class SceuaConfig:
    base_inp: Path
    patch_map: dict
    observed: pd.DataFrame
    run_root: Path
    swmm_node: str
    swmm_attr: str
    aggregate: str
    obs_start: str | None
    obs_end: str | None
    bounds: dict           # name -> ParamBound (from swmm_calibrate.parse_search_space)
    iterations: int
    seed: int
    ngs: int               # number of complexes
    convergence_csv: Path
    swmm_runner: Callable  # signature: (inp_path) -> (rc, rpt, out_path)
    extract_series: Callable  # signature: (out_path) -> pd.DataFrame[timestamp, flow]
    # Optional per-evaluation hook: (call_index, best_kge_so_far, named_params).
    # Injected by callers that want live progress (e.g. the aiswmm calibrate
    # facade, ADR-0005); errors in the callback are swallowed so progress
    # reporting can never break a calibration.
    progress_callback: Callable | None = None


class SwmmSpotSetup:
    """Spotpy setup adapter for the existing SWMM patch-and-run pipeline."""

    def __init__(self, config: SceuaConfig) -> None:
        self.config = config
        # Spotpy expects parameter list via .params attribute or parameters() method.
        import spotpy  # local import keeps non-sceua paths free of the dependency

        self._spotpy = spotpy
        self._param_order = list(config.bounds.keys())
        self._params = [
            spotpy.parameter.Uniform(
                name,
                low=config.bounds[name].min_value,
                high=config.bounds[name].max_value,
                optguess=(config.bounds[name].min_value + config.bounds[name].max_value) / 2.0,
            )
            for name in self._param_order
        ]
        # Observed timestamps cached once so simulation() can align without rereading.
        obs_clean = config.observed.copy()
        obs_clean["timestamp"] = pd.to_datetime(obs_clean["timestamp"])
        self._observed = obs_clean.sort_values("timestamp").reset_index(drop=True)
        # The trace records each rep's primary KGE so we can write convergence.csv.
        self._convergence: list[tuple[int, float]] = []
        self._call_count = 0
        self._last_named_params: dict | None = None

    # spotpy hook: return the parameter generator list.
    def parameters(self):
        return self._spotpy.parameter.generate(self._params)

    def _values_to_named(self, values: Sequence[float]) -> dict[str, float | int]:
        out: dict[str, float | int] = {}
        for i, name in enumerate(self._param_order):
            bound = self.config.bounds[name]
            raw = float(values[i])
            if bound.value_type == "int":
                out[name] = int(round(raw))
            elif bound.precision is not None:
                out[name] = round(raw, bound.precision)
            else:
                out[name] = raw
        return out

    def _run_swmm_for_params(self, params: dict[str, float | int]) -> pd.DataFrame | None:
        cfg = self.config
        self._call_count += 1
        self._last_named_params = dict(params)
        trial_dir = cfg.run_root / f"sceua_{self._call_count:04d}"
        trial_dir.mkdir(parents=True, exist_ok=True)
        try:
            patched_text = patch_inp_text(
                cfg.base_inp.read_text(errors="ignore"),
                cfg.patch_map,
                params,
            )
        except Exception:
            return None
        inp = trial_dir / "model.inp"
        inp.write_text(patched_text, encoding="utf-8")
        try:
            rc, _, out_path = cfg.swmm_runner(inp, trial_dir)
        except FileNotFoundError:
            return None
        except Exception:
            return None
        if rc != 0:
            return None
        try:
            sim = cfg.extract_series(out_path)
        except Exception:
            return None
        return sim

    def simulation(self, values) -> np.ndarray:
        params = self._values_to_named(values)
        sim_df = self._run_swmm_for_params(params)
        if sim_df is None or sim_df.empty:
            # Return a series of NaNs aligned to observed length; objectivefunction
            # will collapse this to a large penalty.
            return np.full(len(self._observed), np.nan, dtype=float)
        aligned = align_series(self._observed, sim_df)
        if aligned.empty:
            return np.full(len(self._observed), np.nan, dtype=float)
        # Spotpy expects simulation() and evaluation() to share length.
        # Project the aligned simulated flow onto the full observed length, NaN-padding
        # any observed timestamps without a matching simulation point.
        obs_ts = self._observed["timestamp"].to_numpy()
        sim_map = dict(zip(aligned["timestamp"].to_numpy(), aligned["flow_sim"].astype(float).to_numpy()))
        sim_array = np.array([sim_map.get(ts, np.nan) for ts in obs_ts], dtype=float)
        return sim_array

    def evaluation(self) -> np.ndarray:
        return self._observed["flow"].astype(float).to_numpy()

    def _record_evaluation(self, kge_value: float) -> None:
        """Append to the convergence trace and fire the optional progress
        hook (ADR-0005): best-KGE-so-far plus the params of THIS call.
        Callback errors are swallowed: progress can never break a run."""
        self._convergence.append((self._call_count, kge_value))
        callback = getattr(self.config, "progress_callback", None)
        if callback is None:
            return
        finite = [k for _, k in self._convergence if math.isfinite(k)]
        best = max(finite) if finite else float("nan")
        try:
            callback(self._call_count, best, dict(self._last_named_params or {}))
        except Exception:
            pass

    def objectivefunction(self, simulation, evaluation, params=None) -> float:
        sim = np.asarray(simulation, dtype=float)
        obs = np.asarray(evaluation, dtype=float)
        mask = np.isfinite(sim) & np.isfinite(obs)
        if mask.sum() < 2:
            self._record_evaluation(float("nan"))
            return 1.0e6  # large penalty: minimisation drives this down
        sim_ok = sim[mask]
        obs_ok = obs[mask]
        ts = self._observed["timestamp"].to_numpy()[mask]
        sim_df = pd.DataFrame({"timestamp": ts, "flow": sim_ok})
        obs_df = pd.DataFrame({"timestamp": ts, "flow": obs_ok})
        kge_result = kge(obs_df, sim_df)
        kge_value = kge_result["kge"]
        if kge_value is None or not math.isfinite(kge_value):
            self._record_evaluation(float("nan"))
            return 1.0e6
        self._record_evaluation(float(kge_value))
        return float(1.0 - kge_value)  # spotpy sceua minimises this

    # Accessors used by run_sceua after sampling completes.
    @property
    def convergence_trace(self) -> list[tuple[int, float]]:
        return list(self._convergence)

    @property
    def param_order(self) -> list[str]:
        return list(self._param_order)


def write_convergence_csv(trace: list[tuple[int, float]], path: Path) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    with path.open("w", newline="") as fh:
        writer = csv.writer(fh)
        writer.writerow(["iteration", "kge"])
        for iteration, kge_value in trace:
            writer.writerow([iteration, "" if math.isnan(kge_value) else f"{kge_value:.8f}"])


def _best_iteration(trace: list[tuple[int, float]]) -> tuple[int, float] | None:
    valid = [(it, v) for it, v in trace if not math.isnan(v)]
    if not valid:
        return None
    return max(valid, key=lambda x: x[1])


def run_sceua(config: SceuaConfig) -> dict[str, Any]:
    """Run SCE-UA, write convergence.csv, return a results dict with best_params + summary."""

    import spotpy  # noqa: WPS433

    setup = SwmmSpotSetup(config)
    sampler = spotpy.algorithms.sceua(
        setup,
        dbname=str(config.run_root / "sceua_db"),
        dbformat="ram",
        random_state=config.seed,
        save_sim=False,
    )
    sampler.sample(config.iterations, ngs=config.ngs)
    trace = setup.convergence_trace
    write_convergence_csv(trace, config.convergence_csv)

    best = _best_iteration(trace)
    if best is None:
        raise RuntimeError("SCE-UA produced no valid iterations; cannot build summary.")

    best_iter, best_kge = best
    # Re-run SWMM at the best parameters to capture full metrics + decomposition.
    # The trace iteration -> params mapping lives implicitly in the per-call trial
    # directories; we re-fetch using spotpy's results database.
    results = sampler.getdata()
    # Spotpy 1.6 ramfs returns a numpy structured array: 'like1' is the objective
    # (we returned 1 - KGE), 'par<name>' the parameters.
    if results is None or len(results) == 0:
        raise RuntimeError("SCE-UA results database was empty.")

    # Pick the best row by smallest 'like1' (= smallest 1 - KGE, i.e. largest KGE).
    like_field = None
    for candidate in ("like1", "like_kge", "like"):
        if candidate in results.dtype.names:
            like_field = candidate
            break
    if like_field is None:
        # Fall back to first numeric field.
        like_field = results.dtype.names[0]
    best_row_idx = int(np.argmin(results[like_field]))
    best_row = results[best_row_idx]
    best_params: dict[str, float | int] = {}
    for name in setup.param_order:
        col = f"par{name}"
        if col in results.dtype.names:
            value = float(best_row[col])
            bound = config.bounds[name]
            if bound.value_type == "int":
                best_params[name] = int(round(value))
            elif bound.precision is not None:
                best_params[name] = round(value, bound.precision)
            else:
                best_params[name] = value

    # Re-run SWMM with best_params to compute full metrics on the aligned series.
    cfg = config
    final_trial_dir = cfg.run_root / "sceua_best"
    final_trial_dir.mkdir(parents=True, exist_ok=True)
    patched_text = patch_inp_text(
        cfg.base_inp.read_text(errors="ignore"),
        cfg.patch_map,
        best_params,
    )
    final_inp = final_trial_dir / "model.inp"
    final_inp.write_text(patched_text, encoding="utf-8")
    rc, _, out_path = cfg.swmm_runner(final_inp, final_trial_dir)
    if rc != 0:
        raise RuntimeError(f"swmm5 final-best run failed with rc={rc}")
    sim_df = cfg.extract_series(out_path)

    aligned = align_series(setup._observed, sim_df)
    metrics_bundle = compute_metrics(setup._observed, sim_df)
    kge_block = kge(setup._observed, sim_df)
    if kge_block["decomposition"] is None or kge_block["kge"] is None:
        raise RuntimeError("KGE undefined for the best parameter set; cannot summarise.")

    # Observed flow on the aligned overlap, for PBIAS%.
    obs_for_pbias = aligned["flow_obs"].astype(float) if not aligned.empty else setup._observed["flow"].astype(float)
    secondary = secondary_metrics_from_bundle(metrics_bundle, obs_for_pbias)

    summary = build_calibration_summary(
        primary_value=kge_block["kge"],
        kge_decomposition=kge_block["decomposition"],
        secondary_metrics=secondary,
        iterations=config.iterations,
        convergence_trace_ref=config.convergence_csv.name,
    )

    return {
        "summary": summary,
        "best_params": best_params,
        "best_iteration": best_iter,
        "best_kge_from_trace": best_kge,
        "convergence_csv": str(config.convergence_csv),
        "total_calls": setup._call_count,
        "metrics_bundle": metrics_bundle.to_dict(),
    }


if __name__ == "__main__":
    # CLI usage is via swmm_calibrate.py search --strategy sceua; this module is
    # intentionally library-style.
    import sys

    print("sceua.py is a library module; invoke via 'swmm_calibrate.py search --strategy sceua'.", file=sys.stderr)
    sys.exit(0)

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

import argparse
import json
import math
import random
import re
import subprocess
import sys
import time
import uuid
from dataclasses import dataclass
from datetime import datetime, timezone
from pathlib import Path
from typing import Any

import pandas as pd
from swmmtoolbox import swmmtoolbox

from candidate_writer import write_candidate_artefacts
from inp_patch import patch_inp_text
from metrics import align_series, compute_metrics, score_from_metrics
from obs_reader import read_series


# PRD-GF-CORE: gap-fill emission helpers.
#
# The agent runtime intercepts ``{"ok": False, "gap_signal": ...}``
# results and routes them through the proposer/UI/recorder. When a
# Python caller (e.g. a test or an MCP wrapper) invokes
# :func:`prepare_calibration_inputs` and the observed-flow file or
# the calibration target field is missing, we emit a structured gap
# signal instead of raising. The legacy CLI still raises
# ``SystemExit`` / ``FileNotFoundError`` for backwards compatibility
# with the long-form ``aiswmm`` recipes that drive this script via
# subprocess; only the new in-process entry point speaks gap-fill.

_VALID_CALIBRATION_TARGETS = {"flow", "depth", "head", "volume"}


def _new_gap_id() -> str:
    """Return a short ``gap-<hex>`` identifier (mirrors gap_fill.protocol)."""
    return f"gap-{uuid.uuid4().hex[:12]}"


def emit_observed_flow_gap_signal(observed_path: str | None) -> dict[str, Any]:
    """Build the L1 ``gap_signal`` result for a missing observed-flow file.

    Used when ``observed_path`` is ``None`` or points at a path that
    does not exist. The shape matches the runtime's interception
    contract.
    """
    return {
        "tool": "swmm_calibrate",
        "args": {"observed": observed_path},
        "ok": False,
        "summary": "missing observed flow file",
        "gap_signal": {
            "gap_id": _new_gap_id(),
            "severity": "L1",
            "kind": "file_path",
            "field": "observed",
            "context": {
                "tool": "swmm_calibrate",
                "step": "load_observed_series",
                "provided_path": observed_path,
            },
        },
    }


def emit_calibration_target_gap_signal(provided: str | None) -> dict[str, Any]:
    """Build the L3 ``gap_signal`` result for a missing calibration target.

    Used when the caller did not specify which series the calibration
    objective should target (flow / depth / head / volume). The
    proposer's registry layer is unlikely to know — most calibration
    tasks default to ``flow`` — so the agent UI prompts the user to
    pick.
    """
    return {
        "tool": "swmm_calibrate",
        "args": {"calibration_target": provided},
        "ok": False,
        "summary": "missing calibration target field",
        "gap_signal": {
            "gap_id": _new_gap_id(),
            "severity": "L3",
            "kind": "param_value",
            "field": "calibration_target",
            "context": {
                "tool": "swmm_calibrate",
                "step": "select_target_series",
                "allowed_values": sorted(_VALID_CALIBRATION_TARGETS),
            },
            "suggestion": {"default": "flow"},
        },
    }


def prepare_calibration_inputs(
    *,
    observed: str | Path | None,
    calibration_target: str | None,
) -> dict[str, Any]:
    """Validate the two PRD-GF-CORE-tracked calibration inputs.

    Returns one of:

    - ``{"ok": True, ...}`` when both inputs are present and the file
      exists on disk.
    - An L1 ``gap_signal`` result when the observed file is missing.
    - An L3 ``gap_signal`` result when the calibration target is
      missing or invalid.

    The function does **not** load the file or run a calibration — it
    is a thin validation gate. The agent runtime calls this before
    invoking the full CLI; on a gap signal the runtime routes through
    the gap-fill state machine and re-invokes with merged args.

    Two emit points (per PRD-GF-CORE):

    1. **L1** — ``observed`` is ``None``, empty string, or points at
       a path that does not exist.
    2. **L3** — ``calibration_target`` is ``None``, empty string, or
       not in ``{"flow", "depth", "head", "volume"}``.

    Order of checks: L1 first, then L3. The runtime's batching
    contract accepts both at once if a single call carries both gaps.
    """
    obs_str = str(observed) if observed is not None else None
    if not obs_str or not obs_str.strip():
        return emit_observed_flow_gap_signal(None)
    if not Path(obs_str).is_file():
        return emit_observed_flow_gap_signal(obs_str)

    if not calibration_target or not str(calibration_target).strip():
        return emit_calibration_target_gap_signal(None)
    if str(calibration_target) not in _VALID_CALIBRATION_TARGETS:
        return emit_calibration_target_gap_signal(str(calibration_target))

    return {
        "tool": "swmm_calibrate",
        "args": {
            "observed": obs_str,
            "calibration_target": str(calibration_target),
        },
        "ok": True,
        "summary": "calibration inputs valid",
    }


@dataclass(frozen=True)
class ParamBound:
    name: str
    min_value: float
    max_value: float
    value_type: str = "float"
    precision: int | None = None

    def from_unit(self, unit_value: float) -> float | int:
        clamped = min(1.0, max(0.0, unit_value))
        value = self.min_value + (self.max_value - self.min_value) * clamped
        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]:
        return {
            "min": self.min_value,
            "max": self.max_value,
            "type": self.value_type,
            "precision": self.precision,
        }


def utc_now_iso() -> str:
    return datetime.now(timezone.utc).isoformat()


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


def ensure_param_sets(obj: Any) -> list[dict]:
    if isinstance(obj, list):
        return obj
    if isinstance(obj, dict) and "parameter_sets" in obj:
        return obj["parameter_sets"]
    raise ValueError("Parameter sets JSON must be a list or contain 'parameter_sets'")


def ensure_named_trial(item: dict, idx: int) -> dict:
    out = dict(item)
    out.setdefault("name", f"trial_{idx:03d}")
    if "params" not in out or not isinstance(out["params"], dict):
        raise ValueError(f"Trial {out['name']} missing params object")
    return out


def is_finite_number(value: Any) -> bool:
    return isinstance(value, (int, float)) and math.isfinite(float(value))


def parse_search_space(obj: Any) -> dict[str, ParamBound]:
    if not isinstance(obj, dict) or not obj:
        raise ValueError("Search-space JSON must be a non-empty object")

    out: dict[str, ParamBound] = {}
    for name, spec in obj.items():
        if isinstance(spec, dict):
            if "min" not in spec or "max" not in spec:
                raise ValueError(f"Search-space parameter '{name}' requires 'min' and 'max'")
            min_value = float(spec["min"])
            max_value = float(spec["max"])
            raw_type = str(spec.get("type", "float")).lower().strip()
            if raw_type in {"int", "integer"}:
                value_type = "int"
            elif raw_type in {"float", "number"}:
                value_type = "float"
            else:
                raise ValueError(f"Unsupported search-space type for '{name}': {raw_type}")
            precision = spec.get("precision")
            if precision is not None:
                precision = int(precision)
                if precision < 0:
                    raise ValueError(f"Precision for '{name}' must be >= 0")
            if value_type == "int":
                precision = None
        elif isinstance(spec, (list, tuple)) and len(spec) == 2:
            min_value = float(spec[0])
            max_value = float(spec[1])
            value_type = "float"
            precision = None
        else:
            raise ValueError(
                f"Search-space parameter '{name}' must be either [min, max] or an object with min/max/type"
            )

        if not math.isfinite(min_value) or not math.isfinite(max_value):
            raise ValueError(f"Search-space parameter '{name}' has non-finite bounds")
        if min_value > max_value:
            raise ValueError(f"Search-space parameter '{name}' has min > max")

        out[name] = ParamBound(
            name=name,
            min_value=min_value,
            max_value=max_value,
            value_type=value_type,
            precision=precision,
        )
    return out


def serialize_bounds(bounds: dict[str, ParamBound]) -> dict[str, dict[str, Any]]:
    return {name: bound.to_dict() for name, bound in bounds.items()}


def sample_random_params(bounds: dict[str, ParamBound], count: int, rng: random.Random) -> list[dict[str, float | int]]:
    out: list[dict[str, float | int]] = []
    for _ in range(count):
        params: dict[str, float | int] = {}
        for name, bound in bounds.items():
            params[name] = bound.from_unit(rng.random())
        out.append(params)
    return out


def sample_lhs_params(bounds: dict[str, ParamBound], count: int, rng: random.Random) -> list[dict[str, float | int]]:
    if count <= 0:
        raise ValueError("LHS sample count must be >= 1")

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

    out: list[dict[str, float | int]] = []
    for idx in range(count):
        params: dict[str, float | int] = {}
        for name, bound in bounds.items():
            params[name] = bound.from_unit(unit_vectors[name][idx])
        out.append(params)
    return out


def refine_bounds_from_elite(
    current_bounds: dict[str, ParamBound],
    global_bounds: dict[str, ParamBound],
    elite_results: list[dict],
    margin_fraction: float,
    min_span_fraction: float,
) -> dict[str, ParamBound]:
    refined: dict[str, ParamBound] = {}
    for name, current in current_bounds.items():
        global_bound = global_bounds[name]
        values = [
            float(rec["params"][name])
            for rec in elite_results
            if isinstance(rec.get("params"), dict) and name in rec["params"]
        ]
        if not values:
            refined[name] = current
            continue

        lo = min(values)
        hi = max(values)
        spread = hi - lo

        global_span = global_bound.max_value - global_bound.min_value
        min_span = max(0.0, global_span * min_span_fraction)

        if spread <= 0:
            center = sum(values) / len(values)
            current_span = current.max_value - current.min_value
            span = max(min_span, current_span * 0.25)
            lo = center - (span / 2.0)
            hi = center + (span / 2.0)
        else:
            lo = lo - (spread * margin_fraction)
            hi = hi + (spread * margin_fraction)
            if (hi - lo) < min_span:
                center = (hi + lo) / 2.0
                lo = center - (min_span / 2.0)
                hi = center + (min_span / 2.0)

        lo = max(global_bound.min_value, lo)
        hi = min(global_bound.max_value, hi)
        if lo >= hi:
            lo = global_bound.min_value
            hi = global_bound.max_value

        refined[name] = ParamBound(
            name=name,
            min_value=float(lo),
            max_value=float(hi),
            value_type=global_bound.value_type,
            precision=global_bound.precision,
        )
    return refined


def build_search_trials(
    samples: list[dict[str, float | int]],
    trial_prefix: str,
    start_index: int,
    strategy: str,
    round_index: int,
) -> list[dict]:
    trials: list[dict] = []
    for local_idx, params in enumerate(samples, start=1):
        global_idx = start_index + local_idx - 1
        trials.append(
            {
                "name": f"{trial_prefix}_{global_idx:04d}",
                "params": params,
                "metadata": {
                    "search_strategy": strategy,
                    "search_round": round_index,
                    "search_sample_index": local_idx,
                },
            }
        )
    return trials


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)
    if isinstance(series, pd.Series):
        df = series.reset_index()
        df.columns = ["timestamp", "flow"]
    elif isinstance(series, pd.DataFrame):
        df = series.reset_index()
        df.columns = ["timestamp", "flow"]
        df = df[["timestamp", "flow"]]
    else:
        raise TypeError(f"Unexpected series type from swmmtoolbox: {type(series)}")

    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


# Canonical SWMM ``ERROR <digits>:`` marker. SWMM 5.x writes these into the
# ``.rpt`` whenever it refuses to advance, and it can still exit 0 while doing
# so. Mirrors ``agentic_swmm/agent/honesty.py`` (skill scripts stay import-free
# from the package per ADR-0006, so the scan is duplicated locally).
_RPT_ERROR_RE = re.compile(r"^\s*(ERROR\s+\d+:.*)$")


def rpt_error_lines(rpt: Path) -> list[str]:
    """Return the verbatim ``ERROR <digits>:`` lines in a SWMM ``.rpt``.

    Empty list when the file is missing, unreadable, or clean.
    """
    try:
        text = rpt.read_text(encoding="utf-8", errors="replace")
    except OSError:
        return []
    return [m.group(1).rstrip() for m in (_RPT_ERROR_RE.match(r) for r in text.splitlines()) if m]


def run_swmm(inp: Path, run_dir: Path, rpt_name: str = "model.rpt", out_name: str = "model.out") -> tuple[int, Path, Path]:
    run_dir.mkdir(parents=True, exist_ok=True)
    rpt = run_dir / rpt_name
    out = run_dir / out_name
    # Honesty (review P1-10): never score a stale ``.rpt``/``.out`` left in a
    # reused trial directory. Clear both before the run so a failed or errored
    # run cannot inherit the previous trial's outputs and get its score.
    for stale in (rpt, out):
        try:
            stale.unlink()
        except FileNotFoundError:
            pass
    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")
    # Honesty: swmm5 can exit 0 while writing ``ERROR <digits>:`` lines and no
    # usable ``.out``. Drop the ``.out`` so no caller scores an invalid run;
    # every caller either checks ``rc``/output existence or fails to extract.
    if proc.returncode == 0 and rpt_error_lines(rpt):
        try:
            out.unlink()
        except FileNotFoundError:
            pass
    return proc.returncode, rpt, out


def describe_series(df: pd.DataFrame) -> dict[str, Any]:
    if df.empty:
        return {"count": 0, "start": None, "end": None}
    ts = pd.to_datetime(df["timestamp"])
    return {
        "count": int(len(df)),
        "start": ts.min().isoformat(),
        "end": ts.max().isoformat(),
    }


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 load_observed_series(
    observed_path: Path,
    timestamp_col: str | None,
    flow_col: str | None,
    time_format: str | None,
    obs_start: str | None,
    obs_end: str | None,
) -> pd.DataFrame:
    observed = read_series(
        observed_path,
        timestamp_col=timestamp_col,
        flow_col=flow_col,
        time_format=time_format,
    )
    return filter_series_window(observed, obs_start, obs_end)


def finalize_trial(
    result: dict,
    status: str,
    reason_code: str,
    reason_detail: str | None,
    started_perf: float,
) -> dict:
    diagnostics = result.setdefault("diagnostics", {})
    diagnostics["finished_at_utc"] = utc_now_iso()
    diagnostics["elapsed_seconds"] = round(time.perf_counter() - started_perf, 6)
    result["status"] = status
    result["reason_code"] = reason_code
    result["reason_detail"] = reason_detail
    if status in {"failed", "invalid"}:
        result["error"] = reason_detail or reason_code
    return result


def evaluate_trial(
    base_inp: Path,
    patch_map: dict,
    trial: dict,
    observed: pd.DataFrame,
    run_root: Path,
    swmm_node: str,
    swmm_attr: str,
    objective: str,
    aggregate: str,
    obs_start: str | None,
    obs_end: str | None,
    dry_run: bool = False,
) -> dict:
    started_perf = time.perf_counter()
    trial_name = trial["name"]
    trial_dir = run_root / trial_name
    trial_dir.mkdir(parents=True, exist_ok=True)
    patched_inp = trial_dir / "model.inp"

    result: dict[str, Any] = {
        "trial": trial_name,
        "params": trial["params"],
        "run_dir": str(trial_dir),
        "dry_run": dry_run,
        "metrics": None,
        "objective": None,
        "status": "pending",
        "reason_code": None,
        "reason_detail": None,
        "observed_series": describe_series(observed),
        "diagnostics": {
            "started_at_utc": utc_now_iso(),
            "swmm_node": swmm_node,
            "swmm_attr": swmm_attr,
            "aggregate": aggregate,
            "observed_window": {"start": obs_start, "end": obs_end},
            "observed_points": int(len(observed)),
        },
    }
    if "metadata" in trial:
        result["metadata"] = trial["metadata"]

    try:
        patched_text = patch_inp_text(base_inp.read_text(errors="ignore"), patch_map, trial["params"])
        patched_inp.write_text(patched_text, encoding="utf-8")
        result["files"] = {"inp": str(patched_inp)}
    except Exception as exc:  # noqa: BLE001
        return finalize_trial(
            result,
            status="invalid",
            reason_code="patch_failed",
            reason_detail=f"Failed to apply parameters to INP: {exc}",
            started_perf=started_perf,
        )

    if dry_run:
        return finalize_trial(
            result,
            status="dry_run",
            reason_code="dry_run_enabled",
            reason_detail="Trial not executed because --dry-run was set",
            started_perf=started_perf,
        )

    try:
        rc, rpt, out = run_swmm(patched_inp, trial_dir)
    except FileNotFoundError as exc:
        return finalize_trial(
            result,
            status="failed",
            reason_code="swmm_binary_missing",
            reason_detail=f"swmm5 executable not found: {exc}",
            started_perf=started_perf,
        )
    except Exception as exc:  # noqa: BLE001
        return finalize_trial(
            result,
            status="failed",
            reason_code="swmm_execution_error",
            reason_detail=f"Failed to execute swmm5: {exc}",
            started_perf=started_perf,
        )

    result["return_code"] = rc
    result["files"].update({"rpt": str(rpt), "out": str(out)})
    # Honesty (review P1-10): a run that wrote SWMM ERROR lines, or produced no
    # fresh .out, is not scorable no matter what the process exit code was.
    rpt_errors = rpt_error_lines(rpt)
    if rpt_errors:
        return finalize_trial(
            result,
            status="failed",
            reason_code="swmm_reported_error",
            reason_detail=f"SWMM reported {len(rpt_errors)} error line(s); first: {rpt_errors[0]}",
            started_perf=started_perf,
        )
    if rc != 0:
        return finalize_trial(
            result,
            status="failed",
            reason_code="swmm_execution_failed",
            reason_detail=f"swmm5 returned non-zero exit code {rc}",
            started_perf=started_perf,
        )
    if not out.exists():
        return finalize_trial(
            result,
            status="failed",
            reason_code="swmm_output_missing",
            reason_detail="SWMM produced no .out for this trial; nothing to score",
            started_perf=started_perf,
        )

    try:
        simulated = extract_simulated_series(out, swmm_node=swmm_node, swmm_attr=swmm_attr, aggregate=aggregate)
    except Exception as exc:  # noqa: BLE001
        return finalize_trial(
            result,
            status="invalid",
            reason_code="simulation_extract_failed",
            reason_detail=f"Failed to extract simulated series from model.out: {exc}",
            started_perf=started_perf,
        )

    try:
        aligned = align_series(observed, simulated)
        metrics = compute_metrics(observed, simulated)
    except Exception as exc:  # noqa: BLE001
        return finalize_trial(
            result,
            status="invalid",
            reason_code="metric_computation_failed",
            reason_detail=f"Failed while computing metrics: {exc}",
            started_perf=started_perf,
        )

    aligned_view = aligned.rename(columns={"flow_obs": "flow"})
    aligned_view = aligned_view[["timestamp", "flow"]] if not aligned_view.empty else pd.DataFrame(columns=["timestamp", "flow"])

    result["simulated_series"] = describe_series(simulated)
    result["aligned_series"] = describe_series(aligned_view)
    result["metrics"] = metrics.to_dict()

    observed_count = int(len(observed))
    overlap_fraction = None if observed_count == 0 else float(metrics.count / observed_count)
    result["diagnostics"]["simulated_points"] = int(len(simulated))
    result["diagnostics"]["aligned_points"] = int(metrics.count)
    result["diagnostics"]["overlap_fraction_of_observed"] = overlap_fraction

    if metrics.count == 0:
        return finalize_trial(
            result,
            status="invalid",
            reason_code="no_overlap",
            reason_detail="Observed and simulated series do not overlap on timestamps",
            started_perf=started_perf,
        )

    try:
        objective_value = score_from_metrics(metrics, objective)
    except Exception as exc:  # noqa: BLE001
        return finalize_trial(
            result,
            status="invalid",
            reason_code="objective_scoring_failed",
            reason_detail=f"Objective scoring failed: {exc}",
            started_perf=started_perf,
        )

    if not is_finite_number(objective_value):
        return finalize_trial(
            result,
            status="invalid",
            reason_code="objective_unavailable",
            reason_detail=f"Objective '{objective}' was not available for this trial",
            started_perf=started_perf,
        )

    result["objective"] = float(objective_value)
    if overlap_fraction is not None and overlap_fraction < 0.75:
        result["warning"] = (
            "Low overlap between observed and simulated timestamps. "
            "Check simulation window, observed window, and aggregation choices."
        )

    return finalize_trial(
        result,
        status="ok",
        reason_code="ok",
        reason_detail=None,
        started_perf=started_perf,
    )


def evaluate_trials(
    base_inp: Path,
    patch_map: dict,
    trials: list[dict],
    observed: pd.DataFrame,
    run_root: Path,
    swmm_node: str,
    swmm_attr: str,
    objective: str,
    aggregate: str,
    obs_start: str | None,
    obs_end: str | None,
    dry_run: bool,
) -> list[dict]:
    return [
        evaluate_trial(
            base_inp,
            patch_map,
            trial,
            observed,
            run_root,
            swmm_node,
            swmm_attr,
            objective,
            aggregate,
            obs_start,
            obs_end,
            dry_run=dry_run,
        )
        for trial in trials
    ]


def rank_results(results: list[dict]) -> list[dict]:
    def key(rec: dict) -> tuple[int, float, str]:
        objective = rec.get("objective")
        has_objective = is_finite_number(objective)
        status_rank = 0 if rec.get("status") == "ok" else 1
        score = float(objective) if has_objective else float("-inf")
        return (status_rank, -score, rec.get("trial", ""))

    return sorted(results, key=key)


def pick_best_result(ranked_results: list[dict]) -> dict | None:
    for rec in ranked_results:
        if rec.get("status") == "ok" and is_finite_number(rec.get("objective")):
            return rec
    return ranked_results[0] if ranked_results else None


def summarize_status_counts(results: list[dict]) -> 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"))
        if status in counts:
            counts[status] += 1
        else:
            counts["other"] += 1
    return counts


def build_ranking_table(ranked_results: list[dict]) -> list[dict]:
    table: list[dict[str, Any]] = []
    for idx, rec in enumerate(ranked_results, start=1):
        metrics = rec.get("metrics") or {}
        metadata = rec.get("metadata") or {}
        table.append(
            {
                "rank": idx,
                "trial": rec.get("trial"),
                "status": rec.get("status"),
                "reason_code": rec.get("reason_code"),
                "objective": rec.get("objective"),
                "metrics_count": metrics.get("count"),
                "nse": metrics.get("nse"),
                "rmse": metrics.get("rmse"),
                "bias": metrics.get("bias"),
                "peak_flow_error": metrics.get("peak_flow_error"),
                "peak_timing_error_minutes": metrics.get("peak_timing_error_minutes"),
                "overlap_fraction_of_observed": (rec.get("diagnostics") or {}).get("overlap_fraction_of_observed"),
                "search_round": metadata.get("search_round"),
                "run_dir": rec.get("run_dir"),
            }
        )
    return table


def _fmt_num(value: Any, digits: int = 4) -> str:
    if value is None:
        return "-"
    if isinstance(value, (int, float)) and math.isfinite(float(value)):
        return f"{float(value):.{digits}f}"
    return str(value)


def format_ranking_text(ranking_table: list[dict], objective: str, top_n: int) -> str:
    if not ranking_table:
        return "No ranking rows available."
    shown = ranking_table[: max(1, top_n)]
    lines = [
        f"Ranking summary ({objective}, showing {len(shown)} of {len(ranking_table)})",
        "rank | trial | status | reason | objective | nse | rmse | count",
        "-----+-------+--------+--------+-----------+-----+------+------",
    ]
    for row in shown:
        lines.append(
            " | ".join(
                [
                    str(row.get("rank", "-")),
                    str(row.get("trial", "-")),
                    str(row.get("status", "-")),
                    str(row.get("reason_code", "-")),
                    _fmt_num(row.get("objective")),
                    _fmt_num(row.get("nse")),
                    _fmt_num(row.get("rmse")),
                    str(row.get("metrics_count", "-")),
                ]
            )
        )
    return "\n".join(lines)


def _pbias_pct_from_bundle_dict(metrics_dict: dict[str, Any], observed: pd.DataFrame) -> float | None:
    """Return PBIAS% in the calibration_summary shape, or ``None``.

    The legacy random/lhs/adaptive trial pipeline emits a
    :class:`MetricBundle`-as-dict (``metrics_dict``) that has ``bias``
    (mean diff) and ``count`` (overlap length); the candidate writer
    expects PBIAS%, defined as ``100 * sum(sim - obs) / sum(obs)``.
    Reconstructing it from ``bias * count`` keeps the conversion small
    and consistent with the SCE-UA path in :mod:`sceua`.
    """
    bias = metrics_dict.get("bias")
    count = metrics_dict.get("count")
    if bias is None or count is None or count == 0:
        return None
    try:
        obs_sum = float(observed["flow"].astype(float).sum())
    except Exception:
        return None
    if obs_sum == 0.0:
        return None
    return float(100.0 * float(bias) * int(count) / obs_sum)


def build_candidate_summary_from_best(
    best: dict[str, Any] | None,
    *,
    strategy: str,
    iterations: int,
    observed: pd.DataFrame,
    convergence_trace_ref: str | None = None,
) -> dict[str, Any] | None:
    """Project a best-result dict into the candidate writer's summary shape.

    The candidate writer is strategy-agnostic and expects the
    ``calibration_summary.json`` shape (KGE primary + decomposition +
    secondary metrics + strategy + iterations + convergence_trace_ref).
    The legacy random/lhs/adaptive payload exposes the same metric
    fields via the best trial's :class:`MetricBundle`-as-dict, so we
    can construct an equivalent summary without re-running SWMM.

    Returns ``None`` when there is no usable best (e.g. all trials
    failed or KGE is undefined) — the candidate writer is then skipped
    by the caller. We choose to skip rather than emit a partial
    candidate so the on-disk evidence boundary stays sharp.
    """
    if not best:
        return None
    metrics = best.get("metrics") or {}
    kge_value = metrics.get("kge")
    decomposition = metrics.get("kge_decomposition")
    if kge_value is None or not isinstance(decomposition, dict):
        return None
    secondary = {
        "nse": metrics.get("nse"),
        "pbias_pct": _pbias_pct_from_bundle_dict(metrics, observed),
        "rmse": metrics.get("rmse"),
        "peak_error_rel": metrics.get("peak_flow_error"),
        "peak_timing_min": metrics.get("peak_timing_error_minutes"),
    }
    return {
        "primary_objective": "kge",
        "primary_value": float(kge_value),
        "kge_decomposition": {
            "r": float(decomposition.get("r", 0.0)),
            "alpha": float(decomposition.get("alpha", 0.0)),
            "beta": float(decomposition.get("beta", 0.0)),
        },
        "secondary_metrics": secondary,
        "strategy": strategy,
        "iterations": int(iterations),
        "convergence_trace_ref": convergence_trace_ref,
    }


def emit_candidate_artefacts(
    args: argparse.Namespace,
    *,
    summary: dict[str, Any] | None,
    best_params: dict[str, Any] | None,
    extra_refs: dict[str, str] | None = None,
) -> None:
    """Write 3-artefact candidate handover when ``--candidate-run-dir`` is set.

    A no-op when the caller did not request a candidate (back-compat
    with existing flows that do not yet route through ``09_audit/``).
    Likewise a no-op when there is no usable best result — the
    canonical INP stays untouched and there is nothing to hand over.
    """
    run_dir = getattr(args, "candidate_run_dir", None)
    if run_dir is None:
        return
    if summary is None or not best_params:
        return
    write_candidate_artefacts(
        run_dir=Path(run_dir),
        canonical_inp=Path(args.base_inp),
        patch_map=load_json(args.patch_map),
        best_params=best_params,
        summary=summary,
        extra_refs=extra_refs or {},
    )


def build_common_controls(args: argparse.Namespace) -> dict[str, Any]:
    return {
        "base_inp": str(args.base_inp),
        "patch_map": str(args.patch_map),
        "observed": str(args.observed),
        "run_root": str(args.run_root),
        "swmm_node": args.swmm_node,
        "swmm_attr": args.swmm_attr,
        "objective": args.objective,
        "aggregate": args.aggregate,
        "obs_start": args.obs_start,
        "obs_end": args.obs_end,
        "dry_run": bool(args.dry_run),
    }


def emit_payload(args: argparse.Namespace, payload: dict) -> None:
    args.summary_json.parent.mkdir(parents=True, exist_ok=True)
    args.summary_json.write_text(json.dumps(payload, indent=2), encoding="utf-8")

    ranking_table = payload.get("ranking_table")
    if args.ranking_json and isinstance(ranking_table, list):
        args.ranking_json.parent.mkdir(parents=True, exist_ok=True)
        args.ranking_json.write_text(json.dumps(ranking_table, indent=2), encoding="utf-8")

    if args.print_ranking and isinstance(ranking_table, list):
        print(
            format_ranking_text(ranking_table, objective=args.objective, top_n=args.ranking_top),
            file=sys.stderr,
        )

    print(json.dumps(payload, indent=2))


def cmd_sensitivity(args: argparse.Namespace) -> None:
    patch_map = load_json(args.patch_map)
    observed = load_observed_series(
        args.observed,
        args.timestamp_col,
        args.flow_col,
        args.time_format,
        args.obs_start,
        args.obs_end,
    )

    trials = [ensure_named_trial(t, i + 1) for i, t in enumerate(ensure_param_sets(load_json(args.parameter_sets)))]
    results = evaluate_trials(
        args.base_inp,
        patch_map,
        trials,
        observed,
        args.run_root,
        args.swmm_node,
        args.swmm_attr,
        args.objective,
        args.aggregate,
        args.obs_start,
        args.obs_end,
        args.dry_run,
    )
    ranked = rank_results(results)
    ranking_table = build_ranking_table(ranked)

    payload = {
        "mode": "sensitivity",
        "objective": args.objective,
        "controls": build_common_controls(args),
        "status_counts": summarize_status_counts(ranked),
        "ranking_table": ranking_table,
        "results": ranked,
    }
    emit_payload(args, payload)


def cmd_calibrate(args: argparse.Namespace) -> None:
    patch_map = load_json(args.patch_map)
    observed = load_observed_series(
        args.observed,
        args.timestamp_col,
        args.flow_col,
        args.time_format,
        args.obs_start,
        args.obs_end,
    )

    trials = [ensure_named_trial(t, i + 1) for i, t in enumerate(ensure_param_sets(load_json(args.parameter_sets)))]
    results = evaluate_trials(
        args.base_inp,
        patch_map,
        trials,
        observed,
        args.run_root,
        args.swmm_node,
        args.swmm_attr,
        args.objective,
        args.aggregate,
        args.obs_start,
        args.obs_end,
        args.dry_run,
    )
    ranked = rank_results(results)
    ranking_table = build_ranking_table(ranked)
    best = pick_best_result(ranked)

    payload = {
        "mode": "calibrate",
        "objective": args.objective,
        "controls": build_common_controls(args),
        "status_counts": summarize_status_counts(ranked),
        "ranking_table": ranking_table,
        "best": best,
        "results": ranked,
    }

    if args.best_params_out and best:
        args.best_params_out.parent.mkdir(parents=True, exist_ok=True)
        args.best_params_out.write_text(json.dumps(best["params"], indent=2), encoding="utf-8")

    summary_for_candidate = build_candidate_summary_from_best(
        best,
        strategy="calibrate",
        iterations=len(trials),
        observed=observed,
    )
    emit_candidate_artefacts(
        args,
        summary=summary_for_candidate,
        best_params=(best["params"] if best else None),
    )
    emit_payload(args, payload)


def cmd_validate(args: argparse.Namespace) -> None:
    patch_map = load_json(args.patch_map)
    params_obj = load_json(args.best_params)
    observed = load_observed_series(
        args.observed,
        args.timestamp_col,
        args.flow_col,
        args.time_format,
        args.obs_start,
        args.obs_end,
    )

    trial = {"name": args.trial_name, "params": params_obj}
    result = evaluate_trial(
        args.base_inp,
        patch_map,
        trial,
        observed,
        args.run_root,
        args.swmm_node,
        args.swmm_attr,
        args.objective,
        args.aggregate,
        args.obs_start,
        args.obs_end,
        dry_run=args.dry_run,
    )

    ranking_table = build_ranking_table([result])
    payload = {
        "mode": "validate",
        "objective": args.objective,
        "controls": build_common_controls(args),
        "status_counts": summarize_status_counts([result]),
        "ranking_table": ranking_table,
        "result": result,
    }
    emit_payload(args, payload)


def _cmd_search_dream_zs(
    args: argparse.Namespace,
    patch_map: dict,
    bounds: dict[str, ParamBound],
    observed: pd.DataFrame,
) -> None:
    """DREAM-ZS branch of search; depends on the optional `spotpy` package."""

    try:
        from dream_zs import DreamZsConfig, run_dream_zs  # local import: spotpy only needed here
    except ImportError as exc:  # pragma: no cover - defensive
        raise SystemExit(
            "DREAM-ZS strategy requires the optional 'spotpy' dependency. "
            "Install it with `pip install spotpy`.\n"
            f"Underlying error: {exc}"
        ) from exc

    if args.objective != "kge":
        raise SystemExit(
            "--strategy dream-zs currently requires --objective kge "
            f"(got --objective {args.objective}). "
            "The DREAM-ZS likelihood is defined on (1 - KGE)."
        )
    if args.dream_chains < 2:
        raise SystemExit("--dream-chains must be >= 2 for a Gelman-Rubin Rhat check.")

    run_root = Path(args.run_root)
    run_root.mkdir(parents=True, exist_ok=True)
    summary_path = Path(args.summary_json)
    summary_path.parent.mkdir(parents=True, exist_ok=True)

    output_dir = (
        Path(args.dream_output_dir)
        if args.dream_output_dir is not None
        else summary_path.parent
    )

    def _runner(inp: Path, trial_dir: Path):
        return run_swmm(inp, trial_dir)

    def _extract(out_path: Path) -> pd.DataFrame:
        return extract_simulated_series(
            out_path,
            swmm_node=args.swmm_node,
            swmm_attr=args.swmm_attr,
            aggregate=args.aggregate,
        )

    config = DreamZsConfig(
        base_inp=Path(args.base_inp),
        patch_map=patch_map,
        observed=observed,
        run_root=run_root,
        swmm_node=args.swmm_node,
        swmm_attr=args.swmm_attr,
        aggregate=args.aggregate,
        obs_start=args.obs_start,
        obs_end=args.obs_end,
        bounds=bounds,
        iterations=int(args.iterations),
        seed=int(args.seed),
        n_chains=int(args.dream_chains),
        sigma=float(args.dream_sigma),
        rhat_threshold=float(args.dream_rhat_threshold),
        output_dir=output_dir,
        swmm_runner=_runner,
        extract_series=_extract,
        runs_after_convergence=int(args.dream_runs_after_convergence),
    )

    result = run_dream_zs(config)
    summary = result["summary"]
    summary["controls"] = {
        **build_common_controls(args),
        "search_space": str(args.search_space),
        "search_strategy": args.strategy,
        "seed": args.seed,
        "iterations": args.iterations,
        "dream_chains": args.dream_chains,
        "dream_sigma": args.dream_sigma,
        "dream_rhat_threshold": args.dream_rhat_threshold,
        "dream_output_dir": str(output_dir),
        "parsed_search_space": serialize_bounds(bounds),
    }

    summary_path.write_text(json.dumps(summary, indent=2), encoding="utf-8")
    if args.best_params_out:
        Path(args.best_params_out).parent.mkdir(parents=True, exist_ok=True)
        Path(args.best_params_out).write_text(
            json.dumps(result["best_params"], indent=2),
            encoding="utf-8",
        )
    else:
        # The acceptance criteria call for best_params.json under the audit
        # directory regardless of --best-params-out; mirror it there when no
        # explicit path was supplied.
        (output_dir / "best_params.json").write_text(
            json.dumps(result["best_params"], indent=2),
            encoding="utf-8",
        )
    dream_extra_refs = {
        "convergence_csv": Path(summary["convergence_trace_ref"]).name,
        "posterior_samples_csv": Path(result["posterior_samples_csv"]).name,
        "posterior_correlation_png": Path(result["correlation_png"]).name,
    }
    emit_candidate_artefacts(
        args,
        summary=summary,
        best_params=result.get("best_params"),
        extra_refs=dream_extra_refs,
    )
    print(json.dumps(summary, indent=2))


def _cmd_search_sceua(
    args: argparse.Namespace,
    patch_map: dict,
    bounds: dict[str, ParamBound],
    observed: pd.DataFrame,
) -> None:
    """SCE-UA branch of search; depends on the optional `spotpy` package."""

    try:
        from sceua import SceuaConfig, run_sceua  # local import: spotpy only needed here
    except ImportError as exc:  # pragma: no cover - defensive
        raise SystemExit(
            "SCE-UA strategy requires the optional 'spotpy' dependency. "
            "Install it with `pip install spotpy`.\n"
            f"Underlying error: {exc}"
        ) from exc

    if args.objective != "kge":
        # SCE-UA is wired to minimise (1 - KGE); make this explicit at the CLI to
        # avoid silent objective drift. Users who want NSE / RMSE optimisation
        # should keep using the random / lhs / adaptive strategies for now.
        raise SystemExit(
            "--strategy sceua currently requires --objective kge "
            f"(got --objective {args.objective}). "
            "Other objectives are tracked in issue #53 (DREAM-ZS) and follow-ups."
        )

    run_root = Path(args.run_root)
    run_root.mkdir(parents=True, exist_ok=True)
    summary_path = Path(args.summary_json)
    summary_path.parent.mkdir(parents=True, exist_ok=True)
    convergence_csv = (
        Path(args.convergence_csv)
        if args.convergence_csv is not None
        else summary_path.parent / "convergence.csv"
    )

    def _runner(inp: Path, trial_dir: Path):
        return run_swmm(inp, trial_dir)

    def _extract(out_path: Path) -> pd.DataFrame:
        return extract_simulated_series(
            out_path,
            swmm_node=args.swmm_node,
            swmm_attr=args.swmm_attr,
            aggregate=args.aggregate,
        )

    config = SceuaConfig(
        base_inp=Path(args.base_inp),
        patch_map=patch_map,
        observed=observed,
        run_root=run_root,
        swmm_node=args.swmm_node,
        swmm_attr=args.swmm_attr,
        aggregate=args.aggregate,
        obs_start=args.obs_start,
        obs_end=args.obs_end,
        bounds=bounds,
        iterations=int(args.iterations),
        seed=int(args.seed),
        ngs=int(args.sceua_ngs),
        convergence_csv=convergence_csv,
        swmm_runner=_runner,
        extract_series=_extract,
    )

    result = run_sceua(config)
    summary = result["summary"]
    # Add `controls` block so the SCE-UA summary remains comparable to the other
    # strategies' top-level CLI payloads.
    summary["controls"] = {
        **build_common_controls(args),
        "search_space": str(args.search_space),
        "search_strategy": args.strategy,
        "seed": args.seed,
        "iterations": args.iterations,
        "sceua_ngs": args.sceua_ngs,
        "parsed_search_space": serialize_bounds(bounds),
    }

    summary_path.write_text(json.dumps(summary, indent=2), encoding="utf-8")
    if args.best_params_out:
        Path(args.best_params_out).parent.mkdir(parents=True, exist_ok=True)
        Path(args.best_params_out).write_text(
            json.dumps(result["best_params"], indent=2),
            encoding="utf-8",
        )
    emit_candidate_artefacts(
        args,
        summary=summary,
        best_params=result.get("best_params"),
        extra_refs={"convergence_csv": Path(summary["convergence_trace_ref"]).name},
    )
    print(json.dumps(summary, indent=2))


def cmd_search(args: argparse.Namespace) -> None:
    if args.iterations < 1:
        raise ValueError("--iterations must be >= 1")
    if args.strategy == "adaptive" and args.rounds < 2:
        raise ValueError("--rounds must be >= 2 for adaptive strategy")
    if args.strategy not in {"adaptive"} and args.rounds != 1:
        raise ValueError("--rounds can only be >1 when --strategy adaptive")
    if not (0 < args.elite_fraction <= 1):
        raise ValueError("--elite-fraction must be in (0, 1]")
    if not (0 <= args.refine_margin <= 1):
        raise ValueError("--refine-margin must be in [0, 1]")
    if not (0 < args.min_span_fraction <= 1):
        raise ValueError("--min-span-fraction must be in (0, 1]")

    patch_map = load_json(args.patch_map)
    bounds = parse_search_space(load_json(args.search_space))
    observed = load_observed_series(
        args.observed,
        args.timestamp_col,
        args.flow_col,
        args.time_format,
        args.obs_start,
        args.obs_end,
    )

    if args.strategy == "sceua":
        _cmd_search_sceua(args, patch_map, bounds, observed)
        return
    if args.strategy == "dream-zs":
        _cmd_search_dream_zs(args, patch_map, bounds, observed)
        return

    rng = random.Random(args.seed)
    trial_counter = 1
    all_results: list[dict] = []
    round_summaries: list[dict] = []

    def sample_once(active_bounds: dict[str, ParamBound], strategy: str) -> list[dict[str, float | int]]:
        if strategy == "random":
            return sample_random_params(active_bounds, args.iterations, rng)
        return sample_lhs_params(active_bounds, args.iterations, rng)

    if args.strategy in {"random", "lhs"}:
        samples = sample_once(bounds, args.strategy)
        trials = build_search_trials(
            samples=samples,
            trial_prefix="search",
            start_index=trial_counter,
            strategy=args.strategy,
            round_index=1,
        )
        trial_counter += len(trials)

        round_results = evaluate_trials(
            args.base_inp,
            patch_map,
            trials,
            observed,
            args.run_root,
            args.swmm_node,
            args.swmm_attr,
            args.objective,
            args.aggregate,
            args.obs_start,
            args.obs_end,
            args.dry_run,
        )
        all_results.extend(round_results)

        round_ranked = rank_results(round_results)
        round_best = pick_best_result(round_ranked)
        round_summaries.append(
            {
                "round": 1,
                "sampling_strategy": args.strategy,
                "trial_count": len(round_results),
                "status_counts": summarize_status_counts(round_results),
                "best_trial": round_best["trial"] if round_best else None,
                "best_objective": round_best["objective"] if round_best else None,
                "bounds_before": serialize_bounds(bounds),
                "bounds_after": serialize_bounds(bounds),
            }
        )

    else:
        active_bounds = dict(bounds)
        for round_idx in range(1, args.rounds + 1):
            bounds_before = serialize_bounds(active_bounds)
            samples = sample_once(active_bounds, "lhs")
            trials = build_search_trials(
                samples=samples,
                trial_prefix="search",
                start_index=trial_counter,
                strategy="adaptive_lhs",
                round_index=round_idx,
            )
            trial_counter += len(trials)

            round_results = evaluate_trials(
                args.base_inp,
                patch_map,
                trials,
                observed,
                args.run_root,
                args.swmm_node,
                args.swmm_attr,
                args.objective,
                args.aggregate,
                args.obs_start,
                args.obs_end,
                args.dry_run,
            )
            all_results.extend(round_results)

            round_ranked = rank_results(round_results)
            valid_round = [
                rec
                for rec in round_ranked
                if rec.get("status") == "ok" and is_finite_number(rec.get("objective"))
            ]
            elite_count = int(math.ceil(len(valid_round) * args.elite_fraction)) if valid_round else 0
            elite_count = max(1, elite_count) if valid_round else 0
            elite = valid_round[:elite_count]

            if elite and round_idx < args.rounds:
                active_bounds = refine_bounds_from_elite(
                    current_bounds=active_bounds,
                    global_bounds=bounds,
                    elite_results=elite,
                    margin_fraction=args.refine_margin,
                    min_span_fraction=args.min_span_fraction,
                )

            round_best = pick_best_result(round_ranked)
            round_summaries.append(
                {
                    "round": round_idx,
                    "sampling_strategy": "lhs",
                    "trial_count": len(round_results),
                    "status_counts": summarize_status_counts(round_results),
                    "elite_count": elite_count,
                    "best_trial": round_best["trial"] if round_best else None,
                    "best_objective": round_best["objective"] if round_best else None,
                    "bounds_before": bounds_before,
                    "bounds_after": serialize_bounds(active_bounds),
                }
            )

    ranked = rank_results(all_results)
    best = pick_best_result(ranked)
    ranking_table = build_ranking_table(ranked)

    payload = {
        "mode": "search",
        "objective": args.objective,
        "controls": {
            **build_common_controls(args),
            "search_space": str(args.search_space),
            "search_strategy": args.strategy,
            "seed": args.seed,
            "iterations": args.iterations,
            "rounds": args.rounds,
            "elite_fraction": args.elite_fraction,
            "refine_margin": args.refine_margin,
            "min_span_fraction": args.min_span_fraction,
            "parsed_search_space": serialize_bounds(bounds),
        },
        "status_counts": summarize_status_counts(ranked),
        "rounds": round_summaries,
        "ranking_table": ranking_table,
        "best": best,
        "results": ranked,
    }

    if args.best_params_out and best:
        args.best_params_out.parent.mkdir(parents=True, exist_ok=True)
        args.best_params_out.write_text(json.dumps(best["params"], indent=2), encoding="utf-8")

    summary_for_candidate = build_candidate_summary_from_best(
        best,
        strategy=args.strategy,
        iterations=int(args.iterations) * max(1, int(args.rounds)),
        observed=observed,
    )
    emit_candidate_artefacts(
        args,
        summary=summary_for_candidate,
        best_params=(best["params"] if best else None),
    )
    emit_payload(args, payload)


def build_parser() -> argparse.ArgumentParser:
    ap = argparse.ArgumentParser()
    sub = ap.add_subparsers(dest="cmd", required=True)

    def add_common(sp: argparse.ArgumentParser, include_param_sets: bool = True) -> None:
        sp.add_argument("--base-inp", required=True, type=Path)
        sp.add_argument("--patch-map", required=True, type=Path)
        if include_param_sets:
            sp.add_argument("--parameter-sets", required=True, type=Path)
        sp.add_argument("--observed", required=True, type=Path)
        sp.add_argument("--run-root", required=True, type=Path)
        sp.add_argument("--swmm-node", default="O1")
        sp.add_argument("--swmm-attr", default="Total_inflow")
        sp.add_argument(
            "--objective",
            default="nse",
            choices=["nse", "kge", "rmse", "bias", "peak_flow_error", "peak_timing_error"],
        )
        sp.add_argument("--aggregate", choices=["none", "daily_mean"], default="none")
        sp.add_argument("--timestamp-col", default=None)
        sp.add_argument("--flow-col", default=None)
        sp.add_argument("--time-format", default=None)
        sp.add_argument("--obs-start", default=None, help="Inclusive observed-series window start, e.g. 1984-05-23")
        sp.add_argument("--obs-end", default=None, help="Inclusive observed-series window end, e.g. 1984-05-28")
        sp.add_argument("--summary-json", required=True, type=Path)
        sp.add_argument("--ranking-json", default=None, type=Path)
        sp.add_argument("--print-ranking", action="store_true")
        sp.add_argument("--ranking-top", default=10, type=int)
        sp.add_argument("--dry-run", action="store_true")

    sp_s = sub.add_parser("sensitivity")
    add_common(sp_s, include_param_sets=True)
    sp_s.set_defaults(func=cmd_sensitivity)

    sp_c = sub.add_parser("calibrate")
    add_common(sp_c, include_param_sets=True)
    sp_c.add_argument("--best-params-out", default=None, type=Path)
    sp_c.add_argument(
        "--candidate-run-dir",
        type=Path,
        default=None,
        help=(
            "Run directory to receive the candidate-handover artefacts "
            "(candidate_calibration.json, candidate_inp_patch.json, "
            "calibration_report.md) in <run>/09_audit/. Required by "
            "`aiswmm calibration accept` (PRD-Z, issue #54)."
        ),
    )
    sp_c.set_defaults(func=cmd_calibrate)

    sp_v = sub.add_parser("validate")
    add_common(sp_v, include_param_sets=False)
    sp_v.add_argument("--best-params", required=True, type=Path)
    sp_v.add_argument("--trial-name", default="validation")
    sp_v.set_defaults(func=cmd_validate)

    sp_search = sub.add_parser("search")
    add_common(sp_search, include_param_sets=False)
    sp_search.add_argument("--search-space", required=True, type=Path)
    sp_search.add_argument(
        "--strategy",
        choices=["random", "lhs", "adaptive", "sceua", "dream-zs"],
        default="lhs",
    )
    sp_search.add_argument("--iterations", type=int, default=12, help="Trial count per round")
    sp_search.add_argument("--rounds", type=int, default=1, help="Number of rounds (adaptive requires >=2)")
    sp_search.add_argument("--seed", type=int, default=42)
    sp_search.add_argument("--elite-fraction", type=float, default=0.3)
    sp_search.add_argument("--refine-margin", type=float, default=0.1)
    sp_search.add_argument("--min-span-fraction", type=float, default=0.1)
    sp_search.add_argument("--best-params-out", default=None, type=Path)
    sp_search.add_argument(
        "--convergence-csv",
        default=None,
        type=Path,
        help="Where SCE-UA writes the per-iteration KGE trace (default: alongside summary).",
    )
    sp_search.add_argument(
        "--sceua-ngs",
        type=int,
        default=4,
        help="Number of complexes for SCE-UA (default 4). Spotpy recommends 2*p+1 minimum.",
    )
    sp_search.add_argument(
        "--dream-chains",
        type=int,
        default=4,
        help="Number of MCMC chains for DREAM-ZS (default 4). >=2 required for Rhat.",
    )
    sp_search.add_argument(
        "--dream-sigma",
        type=float,
        default=0.1,
        help="Likelihood width sigma on (1-KGE) for DREAM-ZS (default 0.1).",
    )
    sp_search.add_argument(
        "--dream-rhat-threshold",
        type=float,
        default=1.2,
        help="Gelman-Rubin Rhat convergence threshold for DREAM-ZS (default 1.2).",
    )
    sp_search.add_argument(
        "--dream-output-dir",
        type=Path,
        default=None,
        help=(
            "Audit directory for DREAM-ZS artefacts (posterior_samples.csv, "
            "chain_convergence.json, posterior_<param>.png, posterior_correlation.png). "
            "Defaults to the parent of --summary-json."
        ),
    )
    sp_search.add_argument(
        "--dream-runs-after-convergence",
        type=int,
        default=50,
        help="Extra DREAM-ZS samples after Gelman-Rubin convergence (default 50).",
    )
    sp_search.add_argument(
        "--candidate-run-dir",
        type=Path,
        default=None,
        help=(
            "Run directory to receive the candidate-handover artefacts "
            "(candidate_calibration.json, candidate_inp_patch.json, "
            "calibration_report.md) in <run>/09_audit/. Required by "
            "`aiswmm calibration accept` (PRD-Z, issue #54)."
        ),
    )
    sp_search.set_defaults(func=cmd_search)

    return ap


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


if __name__ == "__main__":
    main()