optimize.py

"""Multi-objective mix-design optimisation over the trained concrete GNN.

Function 3 (after forward "mix -> strength" and inverse "strength -> mixes"):
given a *target compressive strength*, search the mix-design space for the
batches that minimise **embodied carbon** and **cost** while still reaching the
target. The GNN is used only as the strength oracle/constraint; carbon and cost
are deterministic linear functions of the ingredient masses, so they need no
model and no retraining (see ``embodied_carbon`` in inference.py and
``mix_cost`` below).

Design space is parameterised by the *absolute-volume* method so every candidate
is a physically valid 1 m^3 batch: the binder, water, admixture and (optional)
fibre masses are free decision variables; the coarse/fine aggregate masses are
*derived* to close the 1 m^3 yield (given a coarse/total aggregate split that is
itself a decision variable). Constraints: predicted f_c >= target, water/binder
within a realistic band, and a minimum aggregate volume.

The optimiser is a compact, dependency-free NSGA-II (numpy only) so the app
stays deployable on the existing requirements (no pymoo). It returns the Pareto
front of (carbon, cost); a target sweep gives the carbon-vs-strength frontier.

CLI (verification):

    # Pareto front of carbon vs cost for C40/50 at 28 d
    python app/optimize.py front --target 50 --pop 48 --gen 30

    # min-carbon mix for several strengths -> carbon-vs-strength curve
    python app/optimize.py sweep --targets 30,40,50,60,80

    # binder-only / UHPC search (no coarse-aggregate floor, allow fibre)
    python app/optimize.py front --target 120 --no-coarse-floor --fibre
"""

from __future__ import annotations

import argparse
import os
from dataclasses import dataclass, field
from pathlib import Path
from typing import Dict, List, Optional, Sequence, Tuple

import numpy as np
import pandas as pd

from inference import (  # local module (sets up the concrete_gnn import path)
    AGE_COL,
    DEFAULT_AGE,
    Predictor,
    _HERE,
    build_input_frame,
    embodied_carbon,
)

# ---------------------------------------------------------------------------
# Ingredient cost inventory.
# Indicative gate prices, GBP per kg of material. Concrete-ingredient prices are
# regional and volatile (bulk cement, SCM availability, fibre type, sand
# haulage) -- these are order-of-magnitude defaults for demonstration only.
# Substitute your own supplier quotes; the app/CLI accept an override dict.
# ---------------------------------------------------------------------------
COST_FACTORS_GBP_PER_KG: Dict[str, float] = {
    "cement_kg_m3": 0.105,
    "slag_kg_m3": 0.060,
    "fly_ash_kg_m3": 0.045,
    "silica_fume_kg_m3": 0.550,
    "metakaolin_kg_m3": 0.480,
    "limestone_powder_kg_m3": 0.055,
    "other_scm_kg_m3": 0.120,
    "water_kg_m3": 0.0010,
    "superplasticizer_kg_m3": 2.100,
    "coarse_aggregate_kg_m3": 0.014,
    "fine_aggregate_kg_m3": 0.012,
    "fibre_content_kg_m3": 2.600,
}

# Particle densities (kg/m^3, SSD) for the absolute-volume yield calculation.
DENSITY: Dict[str, float] = {
    "cement_kg_m3": 3150.0,
    "slag_kg_m3": 2900.0,
    "fly_ash_kg_m3": 2300.0,
    "silica_fume_kg_m3": 2200.0,
    "metakaolin_kg_m3": 2500.0,
    "limestone_powder_kg_m3": 2700.0,
    "other_scm_kg_m3": 2600.0,
    "water_kg_m3": 1000.0,
    "superplasticizer_kg_m3": 1070.0,
    "coarse_aggregate_kg_m3": 2650.0,
    "fine_aggregate_kg_m3": 2630.0,
    "fibre_content_kg_m3": 7850.0,  # steel; PP/PVA users should lower this
}
AIR_CONTENT = 0.02  # entrapped air volume fraction

BINDER_COLS = (
    "cement_kg_m3", "slag_kg_m3", "fly_ash_kg_m3", "silica_fume_kg_m3",
    "metakaolin_kg_m3", "limestone_powder_kg_m3", "other_scm_kg_m3",
)

# Fallback design-space box bounds (lo, hi) for the free decision variables,
# used when the checkpoint ships no ``feature_bounds`` (model_config.json). These
# span ordinary structural concrete; widen for UHPC. Overridden per-variable by
# the checkpoint's p01..p99 when available.
DEFAULT_BOUNDS: Dict[str, Tuple[float, float]] = {
    "cement_kg_m3": (150.0, 700.0),
    "slag_kg_m3": (0.0, 320.0),
    "fly_ash_kg_m3": (0.0, 260.0),
    "silica_fume_kg_m3": (0.0, 80.0),
    "water_kg_m3": (120.0, 250.0),
    "superplasticizer_kg_m3": (0.0, 25.0),
    "fibre_content_kg_m3": (0.0, 120.0),
}
COARSE_FRACTION_BOUNDS = (0.45, 0.70)  # coarse / total aggregate, by volume


def mix_cost(
    mix: Dict[str, float],
    factors: Optional[Dict[str, float]] = None,
) -> Tuple[float, pd.DataFrame]:
    """Ingredient cost (GBP per m^3 of concrete) for one mix.

    Mirrors ``inference.embodied_carbon``: sums amount_i * price_i over the
    ingredient columns. Returns ``(total, breakdown)``; missing/NaN amounts
    count as 0.
    """
    factors = COST_FACTORS_GBP_PER_KG if factors is None else factors
    rows, total = [], 0.0
    for col, fac in factors.items():
        amt = mix.get(col, 0.0)
        if amt is None or (isinstance(amt, float) and np.isnan(amt)):
            amt = 0.0
        amt, fac = float(amt), float(fac)
        cost = amt * fac
        total += cost
        rows.append({
            "ingredient": col,
            "amount_kg_m3": round(amt, 2),
            "price_gbp_per_kg": fac,
            "cost_gbp_m3": round(cost, 2),
        })
    return total, pd.DataFrame(rows)


# ---------------------------------------------------------------------------
# Problem definition: decode a decision vector -> a physically valid mix.
# ---------------------------------------------------------------------------
@dataclass
class MixProblem:
    """Maps the NSGA-II decision vector to mixes, objectives and constraints."""

    predictor: Predictor
    target: float
    age: float = DEFAULT_AGE
    allow_fibre: bool = False
    require_coarse: bool = True            # enforce a minimum aggregate volume
    # w/b floor low enough to reach high-strength / UHPC mixes; ceiling for lean
    # low-strength mixes. min aggregate kept modest so paste-rich high-strength
    # mixes stay feasible when coarse aggregate is required.
    wb_bounds: Tuple[float, float] = (0.18, 0.65)
    min_aggregate_volume: float = 0.35     # m^3/m^3 when require_coarse
    coarse_fraction_bounds: Tuple[float, float] = COARSE_FRACTION_BOUNDS
    # Binder components the user can't use (forced to 0, dropped as decision vars).
    exclude_scms: Tuple[str, ...] = ()
    # Curing regime applied to every design (keys CURING_CARBON_FACTORS, so heated
    # curing shows its energy carbon). None -> standard/ambient.
    curing_regime: Optional[str] = None
    # The GNN strength oracle is stochastic (the mesoscale graph is sampled per
    # call; ~3% CoV). ``n_eval`` graph samples are averaged per candidate. During
    # the SEARCH a low n_eval is fine (NSGA-II tolerates noisy objectives and most
    # of the cost is graph construction); the final reported front is re-evaluated
    # at ``final_n_eval`` for accurate strengths + a stable feasibility re-check.
    n_eval: int = 1
    final_n_eval: int = 3
    strength_margin: float = 0.0
    # Which two objectives define the Pareto front:
    #   "carbon_strength" -> minimise carbon, maximise strength (wide trade-off;
    #                        strength ranges from the target upward).
    #   "carbon_cost"     -> minimise carbon and cost (thin/collinear at a fixed
    #                        strength, but a genuine green-vs-cheap front).
    objective: str = "carbon_strength"
    carbon_factors: Optional[Dict[str, float]] = None
    curing_factors: Optional[Dict[str, float]] = None
    cost_factors: Optional[Dict[str, float]] = None
    bounds_override: Optional[Dict[str, Tuple[float, float]]] = None

    # populated in __post_init__
    free_vars: List[str] = field(default_factory=list)
    lows: np.ndarray = field(default_factory=lambda: np.zeros(0))
    highs: np.ndarray = field(default_factory=lambda: np.zeros(0))

    def __post_init__(self) -> None:
        # Realistic paste/aggregate split. An aggregate-volume floor ALWAYS applies
        # so the search can't return paste-only mixes (cement+SCM > aggregate).
        #  * with coarse: aggregate-rich, like normal/HPC concrete (≥0.55 vol),
        #    of which 40–70% is coarse.
        #  * UHPC (no coarse): a fine-aggregate (sand) floor (~0.30 vol) and no
        #    coarse aggregate — UHPC is paste-rich but still ~⅓ sand by volume.
        if self.require_coarse:
            self.min_aggregate_volume = 0.55
            self.coarse_fraction_bounds = (0.40, 0.70)
        else:
            self.min_aggregate_volume = 0.30
            self.coarse_fraction_bounds = (0.0, 0.0)  # sand only
        free = ["cement_kg_m3", "slag_kg_m3", "fly_ash_kg_m3",
                "silica_fume_kg_m3", "water_kg_m3", "superplasticizer_kg_m3"]
        if self.allow_fibre:
            free.append("fibre_content_kg_m3")
        # Drop excluded SCMs (stay absent at 0); cement/water/sp always remain.
        free = [v for v in free if v not in set(self.exclude_scms)]
        self.free_vars = free
        lo, hi = [], []
        for v in free:
            b = self._var_bounds(v)
            lo.append(b[0]); hi.append(b[1])
        # last decision variable: coarse/total aggregate volume split
        lo.append(self.coarse_fraction_bounds[0])
        hi.append(self.coarse_fraction_bounds[1])
        self.lows = np.asarray(lo, dtype=float)
        self.highs = np.asarray(hi, dtype=float)

    def _var_bounds(self, col: str) -> Tuple[float, float]:
        if self.bounds_override and col in self.bounds_override:
            return self.bounds_override[col]
        cb = self.predictor.bounds.get(col) if self.predictor.bounds else None
        if cb:
            lo = max(0.0, float(cb.get("p01", cb.get("min", 0.0))))
            hi = float(cb.get("p99", cb.get("max", lo)))
            if hi > lo:
                return lo, hi
        return DEFAULT_BOUNDS.get(col, (0.0, 0.0))

    @property
    def n_var(self) -> int:
        return len(self.free_vars) + 1

    # ---- decode ---- #
    def decode(self, x: np.ndarray) -> Tuple[Dict[str, float], float]:
        """Decision vector -> (full mix dict, aggregate volume m^3)."""
        mix: Dict[str, float] = {AGE_COL: float(self.age)}
        vol_used = AIR_CONTENT
        for i, v in enumerate(self.free_vars):
            m = float(max(0.0, x[i]))
            mix[v] = m
            vol_used += m / DENSITY[v]
        agg_vol = 1.0 - vol_used
        cf = float(np.clip(x[-1], 0.0, 1.0))
        agg_clamped = max(0.0, agg_vol)
        mix["coarse_aggregate_kg_m3"] = agg_clamped * cf * DENSITY["coarse_aggregate_kg_m3"]
        mix["fine_aggregate_kg_m3"] = agg_clamped * (1.0 - cf) * DENSITY["fine_aggregate_kg_m3"]
        if self.curing_regime:
            mix["curing_regime_norm"] = self.curing_regime
        return mix, agg_vol

    def binder(self, mix: Dict[str, float]) -> float:
        return max(1e-6, sum(float(mix.get(c, 0.0) or 0.0) for c in BINDER_COLS))

    # ---- batch evaluation ---- #
    def evaluate(self, X: np.ndarray, n_eval: Optional[int] = None
                 ) -> Tuple[np.ndarray, np.ndarray, List[Dict[str, float]]]:
        """Population -> (objectives F[n,2], constraint-violation CV[n], mixes).

        Strength is the mean over ``n_eval`` graph samples (the oracle is noisy),
        evaluated in a single batched GNN call. ``n_eval`` overrides the search
        default (used for the higher-fidelity final re-evaluation).
        """
        mixes, aggs = [], []
        for row in X:
            mix, agg = self.decode(row)
            mixes.append(mix); aggs.append(agg)
        n, k = len(mixes), max(1, n_eval if n_eval is not None else self.n_eval)
        preds = self.predictor.predict_df(
            build_input_frame([m for m in mixes for _ in range(k)]),
            which=("gnn",))["gnn"].reshape(n, k)
        fc_mean = preds.mean(axis=1)
        fc_std = preds.std(axis=1)
        floor = self.target * (1.0 + self.strength_margin)
        F = np.zeros((n, 2), dtype=float)
        CV = np.zeros(n, dtype=float)
        for i, (mix, agg) in enumerate(zip(mixes, aggs)):
            carbon = embodied_carbon(mix, self.carbon_factors, self.curing_factors)[0]
            cost = mix_cost(mix, self.cost_factors)[0]
            # Second objective: maximise strength (store -strength) for the wide
            # carbon-vs-strength front, else minimise cost.
            F[i, 0] = carbon
            F[i, 1] = (-float(fc_mean[i]) if self.objective == "carbon_strength"
                       else cost)
            wb = float(mix["water_kg_m3"]) / self.binder(mix)
            g = [
                (floor - float(fc_mean[i])) / max(self.target, 1.0),   # strength
                self.min_aggregate_volume - agg,                       # min aggregate
                self.wb_bounds[0] - wb,                                # w/b floor
                wb - self.wb_bounds[1],                                # w/b ceiling
            ]
            CV[i] = float(sum(max(0.0, gi) for gi in g))
            mix["pred_gnn"] = float(fc_mean[i])
            mix["pred_gnn_std"] = float(fc_std[i])
            mix["embodied_carbon_kgco2e_m3"] = float(carbon)
            mix["cost_gbp_m3"] = float(cost)
            mix["water_binder_ratio"] = float(wb)
        return F, CV, mixes


# ---------------------------------------------------------------------------
# Compact constrained NSGA-II (numpy only).
# ---------------------------------------------------------------------------
def _constrained_dominates(i: int, j: int, F: np.ndarray, CV: np.ndarray) -> bool:
    """Deb's constraint-domination: feasibility first, then Pareto on F."""
    if CV[i] <= 0 and CV[j] <= 0:
        le = np.all(F[i] <= F[j])
        lt = np.any(F[i] < F[j])
        return bool(le and lt)
    if CV[i] <= 0 and CV[j] > 0:
        return True
    if CV[i] > 0 and CV[j] > 0:
        return CV[i] < CV[j]
    return False


def _fast_non_dominated_sort(F: np.ndarray, CV: np.ndarray) -> List[np.ndarray]:
    n = len(F)
    S: List[List[int]] = [[] for _ in range(n)]
    ndom = np.zeros(n, dtype=int)
    fronts: List[List[int]] = [[]]
    for p in range(n):
        for q in range(n):
            if p == q:
                continue
            if _constrained_dominates(p, q, F, CV):
                S[p].append(q)
            elif _constrained_dominates(q, p, F, CV):
                ndom[p] += 1
        if ndom[p] == 0:
            fronts[0].append(p)
    i = 0
    while fronts[i]:
        nxt: List[int] = []
        for p in fronts[i]:
            for q in S[p]:
                ndom[q] -= 1
                if ndom[q] == 0:
                    nxt.append(q)
        i += 1
        fronts.append(nxt)
    return [np.asarray(f, dtype=int) for f in fronts[:-1]]


def _crowding_distance(F: np.ndarray, front: np.ndarray) -> np.ndarray:
    m = F.shape[1]
    d = np.zeros(len(front), dtype=float)
    for k in range(m):
        vals = F[front, k]
        order = np.argsort(vals)
        d[order[0]] = d[order[-1]] = np.inf
        lo, hi = vals[order[0]], vals[order[-1]]
        span = hi - lo
        if span <= 0:
            continue
        for t in range(1, len(front) - 1):
            d[order[t]] += (vals[order[t + 1]] - vals[order[t - 1]]) / span
    return d


def _tournament(pop_idx: np.ndarray, rank: np.ndarray, crowd: np.ndarray,
                rng: np.random.Generator, n: int) -> np.ndarray:
    a = rng.integers(0, len(pop_idx), size=n)
    b = rng.integers(0, len(pop_idx), size=n)
    out = np.empty(n, dtype=int)
    for t in range(n):
        ia, ib = pop_idx[a[t]], pop_idx[b[t]]
        if rank[ia] < rank[ib] or (rank[ia] == rank[ib] and crowd[ia] > crowd[ib]):
            out[t] = ia
        else:
            out[t] = ib
    return out


def _sbx(p1: np.ndarray, p2: np.ndarray, lo: np.ndarray, hi: np.ndarray,
         rng: np.random.Generator, eta: float = 15.0, pc: float = 0.9):
    c1, c2 = p1.copy(), p2.copy()
    if rng.random() > pc:
        return c1, c2
    for i in range(len(p1)):
        if rng.random() > 0.5 or abs(p1[i] - p2[i]) < 1e-12:
            continue
        x1, x2 = min(p1[i], p2[i]), max(p1[i], p2[i])
        u = rng.random()
        beta = 1.0 + 2.0 * (x1 - lo[i]) / (x2 - x1)
        alpha = 2.0 - beta ** -(eta + 1)
        bq = (u * alpha) ** (1 / (eta + 1)) if u <= 1 / alpha \
            else (1 / (2 - u * alpha)) ** (1 / (eta + 1))
        c1[i] = 0.5 * ((x1 + x2) - bq * (x2 - x1))
        beta = 1.0 + 2.0 * (hi[i] - x2) / (x2 - x1)
        alpha = 2.0 - beta ** -(eta + 1)
        bq = (u * alpha) ** (1 / (eta + 1)) if u <= 1 / alpha \
            else (1 / (2 - u * alpha)) ** (1 / (eta + 1))
        c2[i] = 0.5 * ((x1 + x2) + bq * (x2 - x1))
    return np.clip(c1, lo, hi), np.clip(c2, lo, hi)


def _poly_mutation(x: np.ndarray, lo: np.ndarray, hi: np.ndarray,
                   rng: np.random.Generator, eta: float = 20.0,
                   pm: Optional[float] = None) -> np.ndarray:
    pm = (1.0 / len(x)) if pm is None else pm
    y = x.copy()
    for i in range(len(x)):
        if rng.random() > pm or hi[i] <= lo[i]:
            continue
        u = rng.random()
        delta1 = (y[i] - lo[i]) / (hi[i] - lo[i])
        delta2 = (hi[i] - y[i]) / (hi[i] - lo[i])
        if u < 0.5:
            dq = (2 * u + (1 - 2 * u) * (1 - delta1) ** (eta + 1)) ** (1 / (eta + 1)) - 1
        else:
            dq = 1 - (2 * (1 - u) + 2 * (u - 0.5) * (1 - delta2) ** (eta + 1)) ** (1 / (eta + 1))
        y[i] = np.clip(y[i] + dq * (hi[i] - lo[i]), lo[i], hi[i])
    return y


def nsga2(problem: MixProblem, pop_size: int = 48, n_gen: int = 30,
          seed: int = 17, verbose: bool = True, progress_cb=None):
    """Run NSGA-II on a MixProblem. Returns (X, F, CV) of the final population.

    ``progress_cb(fraction)`` is called after each generation (0..1) for UIs.
    """
    rng = np.random.default_rng(seed)
    lo, hi = problem.lows, problem.highs
    X = lo + rng.random((pop_size, problem.n_var)) * (hi - lo)
    F, CV, _ = problem.evaluate(X)

    for gen in range(n_gen):
        fronts = _fast_non_dominated_sort(F, CV)
        rank = np.zeros(len(X), dtype=int)
        crowd = np.zeros(len(X), dtype=float)
        for r, fr in enumerate(fronts):
            rank[fr] = r
            crowd[fr] = _crowding_distance(F, fr)
        # offspring
        parents = _tournament(np.arange(len(X)), rank, crowd, rng, pop_size)
        kids = []
        for k in range(0, pop_size, 2):
            p1, p2 = X[parents[k]], X[parents[(k + 1) % pop_size]]
            c1, c2 = _sbx(p1, p2, lo, hi, rng)
            kids.append(_poly_mutation(c1, lo, hi, rng))
            kids.append(_poly_mutation(c2, lo, hi, rng))
        Xc = np.asarray(kids[:pop_size])
        Fc, CVc, _ = problem.evaluate(Xc)
        # environmental selection from combined parent+child pool
        X = np.vstack([X, Xc]); F = np.vstack([F, Fc]); CV = np.concatenate([CV, CVc])
        fronts = _fast_non_dominated_sort(F, CV)
        keep: List[int] = []
        for fr in fronts:
            if len(keep) + len(fr) <= pop_size:
                keep.extend(fr.tolist())
            else:
                d = _crowding_distance(F, fr)
                order = fr[np.argsort(-d)]
                keep.extend(order[: pop_size - len(keep)].tolist())
                break
        keep = np.asarray(keep, dtype=int)
        X, F, CV = X[keep], F[keep], CV[keep]
        if verbose:
            feas = int(np.sum(CV <= 0))
            best = F[CV <= 0, 0].min() if feas else np.nan
            print(f"  gen {gen + 1:>3}/{n_gen}  feasible {feas:>3}/{pop_size}"
                  f"  min-carbon {best:8.1f}", flush=True)
        if progress_cb is not None:
            progress_cb((gen + 1) / n_gen)
    return X, F, CV


# ---------------------------------------------------------------------------
# Public entry points
# ---------------------------------------------------------------------------
_INGREDIENT_COLS = list(DENSITY.keys())


def _select_spread(F: np.ndarray, n: int) -> np.ndarray:
    """Pick up to ``n`` well-spread indices: best fronts first, then by crowding.

    Reuses the NSGA-II environmental-selection rule, so the chosen points favour
    the true Pareto front and, when it is thinner than ``n``, fill from the next
    fronts -- giving a visible spread of options even when carbon and cost are
    nearly collinear at a fixed strength.
    """
    if len(F) <= n:
        return np.arange(len(F))
    fronts = _fast_non_dominated_sort(F, np.zeros(len(F)))
    chosen: List[int] = []
    for fr in fronts:
        if len(chosen) + len(fr) <= n:
            chosen.extend(fr.tolist())
        else:
            d = _crowding_distance(F, fr)
            order = fr[np.argsort(-d)]
            chosen.extend(order[: n - len(chosen)].tolist())
            break
        if len(chosen) >= n:
            break
    return np.asarray(chosen, dtype=int)


def _pareto_table(problem: MixProblem, X: np.ndarray, F: np.ndarray,
                  CV: np.ndarray, n_points: int = 8) -> pd.DataFrame:
    """Up to ``n_points`` feasible low-carbon/low-cost mixes, sorted by carbon."""
    feas = CV <= 0
    if not np.any(feas):
        return pd.DataFrame()
    # Re-evaluate the search-feasible set once with a fresh (independent) set of
    # graph samples and keep only those still feasible -- the strength oracle is
    # noisy, so a borderline mix can flip. Deterministic constraints (yield, w/b)
    # cannot flip, so re-using CV here is sound.
    Xf = X[feas]
    Ff, CVf, mixes = problem.evaluate(Xf, n_eval=problem.final_n_eval)
    sel = np.where(CVf <= 0)[0]
    if len(sel) == 0:
        return pd.DataFrame()
    mixes = [mixes[i] for i in sel]
    df = pd.DataFrame(mixes)
    for c in _INGREDIENT_COLS:           # inactive components are absent -> 0
        if c not in df.columns:
            df[c] = 0.0
    df = df.drop_duplicates(subset=_INGREDIENT_COLS).reset_index(drop=True)
    # Choose a spread of up to n_points distinct mixes over the two objectives
    # actually optimised, so the points span the real trade-off.
    if problem.objective == "carbon_strength":
        obj = np.column_stack([df["embodied_carbon_kgco2e_m3"].to_numpy(float),
                               -df["pred_gnn"].to_numpy(float)])
    else:
        obj = df[["embodied_carbon_kgco2e_m3", "cost_gbp_m3"]].to_numpy(dtype=float)
    df = df.iloc[_select_spread(obj, n_points)].reset_index(drop=True)
    # Drop ingredient columns that are zero across every mix (e.g. unused SCMs)
    # to keep the table readable; always keep the structural staples.
    core = {"cement_kg_m3", "water_kg_m3", "coarse_aggregate_kg_m3",
            "fine_aggregate_kg_m3"}
    show_ingredients = [c for c in _INGREDIENT_COLS
                        if c in core or df[c].abs().to_numpy().sum() > 0]
    keep = ["pred_gnn", "pred_gnn_std", "embodied_carbon_kgco2e_m3",
            "cost_gbp_m3", "water_binder_ratio"] + show_ingredients + [AGE_COL]
    cols = [c for c in keep if c in df.columns]
    df = df[cols].copy()
    for c in cols:
        df[c] = pd.to_numeric(df[c], errors="coerce").round(
            3 if c == "water_binder_ratio" else 1)
    df["target_mpa"] = float(problem.target)
    return df.sort_values("embodied_carbon_kgco2e_m3").reset_index(drop=True)


def optimize_mix(predictor: Predictor, target: float, *, pop_size: int = 48,
                 n_gen: int = 30, seed: int = 17, allow_fibre: bool = False,
                 require_coarse: bool = True, age: float = DEFAULT_AGE,
                 wb_bounds: Tuple[float, float] = (0.18, 0.65),
                 n_eval: int = 1, final_n_eval: int = 3,
                 strength_margin: float = 0.0, n_points: int = 8,
                 objective: str = "carbon_strength",
                 exclude_scms: Tuple[str, ...] = (),
                 curing_regime: Optional[str] = None,
                 cost_factors: Optional[Dict[str, float]] = None,
                 carbon_factors: Optional[Dict[str, float]] = None,
                 curing_factors: Optional[Dict[str, float]] = None,
                 bounds_override: Optional[Dict[str, Tuple[float, float]]] = None,
                 verbose: bool = True, progress_cb=None) -> pd.DataFrame:
    """Pareto-optimal (carbon vs cost) mixes that reach >= ``target`` MPa.

    Returns a DataFrame (one row per non-dominated mix) sorted by embodied
    carbon, with predicted strength (mean +/- std over ``n_eval`` graph samples),
    carbon, cost, w/b and full ingredient amounts (kg/m^3). Empty if no feasible
    mix was found (try widening bounds, raising ``n_gen``, or relaxing
    constraints). ``strength_margin`` adds fractional headroom over the target
    (e.g. 0.1 to design for the mean exceeding the target by 10%).
    """
    problem = MixProblem(
        predictor=predictor, target=float(target), age=age,
        allow_fibre=allow_fibre, require_coarse=require_coarse,
        wb_bounds=wb_bounds, n_eval=n_eval, final_n_eval=final_n_eval,
        strength_margin=strength_margin,
        objective=objective, exclude_scms=exclude_scms,
        curing_regime=curing_regime, cost_factors=cost_factors,
        carbon_factors=carbon_factors, curing_factors=curing_factors,
        bounds_override=bounds_override,
    )
    X, F, CV = nsga2(problem, pop_size=pop_size, n_gen=n_gen, seed=seed,
                     verbose=verbose, progress_cb=progress_cb)
    return _pareto_table(problem, X, F, CV, n_points=n_points)


def carbon_strength_frontier(predictor: Predictor, targets: Sequence[float],
                             **kw) -> pd.DataFrame:
    """Min-carbon feasible mix for each target -> the carbon-vs-strength curve.

    For each target strength runs ``optimize_mix`` and keeps the lowest-carbon
    Pareto point. Returns one row per target (target, achieved strength, carbon,
    cost, w/b, ingredients). Targets with no feasible mix are skipped.
    """
    rows = []
    for t in targets:
        df = optimize_mix(predictor, t, **kw)
        if df.empty:
            continue
        rows.append(df.iloc[0])
    return pd.DataFrame(rows).reset_index(drop=True)


# ---------------------------------------------------------------------------
# Checkpoint discovery (auto-tracks the most recently retrained model)
# ---------------------------------------------------------------------------
def discover_checkpoint() -> Path:
    """Pick the checkpoint dir: $CONCRETE_CKPT_DIR, else newest hierarchical.pt."""
    env = os.environ.get("CONCRETE_CKPT_DIR")
    if env and (Path(env) / "hierarchical.pt").exists():
        return Path(env)
    roots = [_HERE, _HERE / "checkpoints_full_rich",
             _HERE.parent / "Hybrid" / "outputs"]
    cands: List[Path] = []
    for r in roots:
        if not r.exists():
            continue
        cands.extend(p.parent for p in r.glob("*/hierarchical.pt"))
        if (r / "hierarchical.pt").exists():
            cands.append(r)
    if not cands:
        raise FileNotFoundError(
            "No hierarchical.pt found. Set CONCRETE_CKPT_DIR or pass "
            "--checkpoint-dir.")
    return max(cands, key=lambda d: (d / "hierarchical.pt").stat().st_mtime)


# ---------------------------------------------------------------------------
# CLI
# ---------------------------------------------------------------------------
def _print_df(df: pd.DataFrame) -> None:
    if df.empty:
        print("No feasible mix found — widen bounds / raise --gen / relax constraints.")
        return
    pd.set_option("display.width", 220, "display.max_columns", 40)
    print(df.to_string(index=False))


def main() -> None:
    ap = argparse.ArgumentParser(description=__doc__)
    sub = ap.add_subparsers(dest="cmd", required=True)

    common = argparse.ArgumentParser(add_help=False)
    common.add_argument("--checkpoint-dir", default=None,
                        help="defaults to the most recently trained checkpoint")
    common.add_argument("--pop", type=int, default=48)
    common.add_argument("--gen", type=int, default=30)
    common.add_argument("--seed", type=int, default=17)
    common.add_argument("--fibre", action="store_true", help="allow fibre dosing")
    common.add_argument("--no-coarse-floor", action="store_true",
                        help="drop the minimum-aggregate constraint (UHPC/mortar)")
    common.add_argument("--age", type=float, default=DEFAULT_AGE)
    common.add_argument("--n-eval", type=int, default=3,
                        help="graph samples averaged per candidate (noisy oracle)")
    common.add_argument("--margin", type=float, default=0.0,
                        help="fractional strength headroom over target, e.g. 0.1")

    f = sub.add_parser("front", parents=[common],
                       help="Pareto front (carbon vs cost) at one target")
    f.add_argument("--target", type=float, required=True)

    s = sub.add_parser("sweep", parents=[common],
                       help="min-carbon mix across several targets")
    s.add_argument("--targets", type=str, required=True,
                   help="comma-separated MPa values, e.g. 30,40,50,60,80")

    args = ap.parse_args()
    ckpt = Path(args.checkpoint_dir) if args.checkpoint_dir else discover_checkpoint()
    print(f"checkpoint: {ckpt}")
    pred = Predictor(ckpt)
    kw = dict(pop_size=args.pop, n_gen=args.gen, seed=args.seed,
              allow_fibre=args.fibre, require_coarse=not args.no_coarse_floor,
              age=args.age, n_eval=args.n_eval, strength_margin=args.margin)

    if args.cmd == "front":
        print(f"Optimising carbon & cost at f_c >= {args.target:.0f} MPa ...")
        _print_df(optimize_mix(pred, args.target, **kw))
    else:
        targets = [float(t) for t in args.targets.split(",") if t.strip()]
        print(f"Carbon-vs-strength frontier for {targets} MPa ...")
        _print_df(carbon_strength_frontier(pred, targets, verbose=False, **kw))


if __name__ == "__main__":
    main()