cryosim/optimization/optimizer.py

"""Constraint-aware hardware design optimization.

Finds optimal hardware parameters that maximize (or minimize) a target metric
subject to constraints on other metrics.

Provides three optimizers:
    - optimize(): L-BFGS-B single start (fast, may get trapped near flow cliff)
    - multi_start_optimize(): LHS-sampled multi-start L-BFGS-B (better global coverage)
    - differential_evolution_optimize(): DE global optimizer (most robust for cliff regions)
"""

from dataclasses import dataclass
from typing import Dict, List, Optional

import numpy as np
from scipy.optimize import minimize, differential_evolution
from scipy.stats import qmc

from cryosim.calibration.params import (
    PARAM_NAMES, get_bounds, get_nominal_values, apply_overrides,
)
from cryosim.hardware.config import load_config
from cryosim.engine.fast import ICV_open


def _extract_metrics(out, hist):
    return {
        "mdot_kgpm": hist["mdot_kgpm"],
        "mass_eff": float(out[1, 0]),
        "Tc_peak_K": float(np.max(hist["Tc_K"])),
        "kWh_extend": hist["kWh_extend"],
        "pc_peak_barg": float(np.max(hist["pc"])),
    }


def _parse_constraint(s: str):
    s = s.strip()
    if s.startswith(">="):
        return ">=", float(s[2:])
    elif s.startswith("<="):
        return "<=", float(s[2:])
    elif s.startswith(">"):
        return ">", float(s[1:])
    elif s.startswith("<"):
        return "<", float(s[1:])
    raise ValueError(f"Cannot parse constraint: '{s}'. Use '>X' or '<X'.")


def _check_constraint(value, op, threshold):
    if op in (">", ">="):
        return max(0.0, threshold - value)
    return max(0.0, value - threshold)


# ---- Metric-name resolution ---------------------------------------------------

_METRIC_MAP = {
    "mdot": "mdot_kgpm",
    "mass_eff": "mass_eff",
    "kWh": "kWh_extend",
    "Tc_peak": "Tc_peak_K",
    "stall_pressure": "stall_pressure",
    "fill_time": "fill_time",
}


def _resolve_target(target: str):
    """Parse target string into (direction, metric_key).

    Returns:
        (direction, metric_key) where direction is -1.0 for max, +1.0 for min.
    """
    if target.startswith("max_"):
        metric_name = target[4:]
        direction = -1.0
    elif target.startswith("min_"):
        metric_name = target[4:]
        direction = 1.0
    else:
        raise ValueError(f"Target must start with 'max_' or 'min_', got '{target}'")
    target_metric = _METRIC_MAP.get(metric_name, metric_name)
    return direction, target_metric


# ---- Fill-level helpers -------------------------------------------------------

def _stall_pressure_with_overrides(cfg_override, speed, Ptank, Psat):
    """Compute stall pressure using a modified config.

    Evaluates the single-cycle engine at 25 pressures from 50-950 bar and
    interpolates the 0.05 kg/min crossing.
    """
    args = cfg_override.to_engine_args()
    pressures = np.linspace(50, 950, 25)
    flows = np.zeros(len(pressures))
    for i, P in enumerate(pressures):
        try:
            out, hist = ICV_open(
                Pexit_barg=float(P), speed_f=speed,
                Ptank_barg=Ptank, Psat_barg=Psat, **args,
            )
            flows[i] = max(0.0, hist["mdot_kgpm"])
        except Exception:
            flows[i] = 0.0
    # Find where flow drops below 0.05 kg/min
    threshold = 0.05
    for i in range(len(pressures) - 1):
        if flows[i] >= threshold and flows[i + 1] < threshold:
            frac = (threshold - flows[i]) / (flows[i + 1] - flows[i])
            return float(pressures[i] + frac * (pressures[i + 1] - pressures[i]))
    # Never crossed: either always above or always below
    if flows[-1] >= threshold:
        return float(pressures[-1])
    return 0.0


def _fill_time_proxy_with_overrides(cfg_override, speed, target_bar, Ptank, Psat):
    """Estimate fill time using stall pressure as a proxy.

    A proper fill-time estimate would integrate 1/mdot(P) over the pressure
    range, but that requires ~25 engine evaluations per objective call which
    is too expensive inside an optimizer.  Instead we use stall pressure as a
    proxy: higher stall pressure means the pump can sustain flow to higher
    pressures, which dominates fill time.  The proxy value is negated stall
    pressure so that minimizing fill_time is equivalent to maximizing stall
    pressure.  If the pump cannot reach the target pressure, returns a large
    penalty value (1e6).

    Returns:
        Proxy fill time value (lower is better).
    """
    stall_p = _stall_pressure_with_overrides(cfg_override, speed, Ptank, Psat)
    if stall_p < target_bar:
        return 1e6  # unreachable — huge penalty
    return -stall_p  # proxy: higher stall pressure -> lower (better) fill time


# ---- Objective builder --------------------------------------------------------

def _build_objective(
    cfg, direction, target_metric, parsed_constraints, penalty_weight,
    speed, Pexit, Ptank, Psat,
):
    """Build a closure that evaluates the objective for a parameter vector.

    Returns:
        (objective_fn, n_evals_counter) where n_evals_counter is a mutable list [count].
    """
    n_evals = [0]

    def objective(x):
        n_evals[0] += 1
        override_cfg = apply_overrides(cfg, list(x))

        # --- fill-level objectives ---
        if target_metric == "stall_pressure":
            try:
                stall_p = _stall_pressure_with_overrides(override_cfg, speed, Ptank, Psat)
            except Exception:
                return 1e6
            obj = direction * stall_p
            # No per-cycle metrics for constraints when using stall objective
            for cmetric, (op, threshold) in parsed_constraints.items():
                obj += penalty_weight * _check_constraint(0.0, op, threshold) ** 2
            return obj

        if target_metric == "fill_time":
            try:
                proxy = _fill_time_proxy_with_overrides(
                    override_cfg, speed, Pexit, Ptank, Psat,
                )
            except Exception:
                return 1e6
            obj = direction * proxy
            return obj

        # --- single-cycle objectives ---
        args = override_cfg.to_engine_args()
        try:
            out, hist = ICV_open(
                Pexit_barg=Pexit, speed_f=speed,
                Ptank_barg=Ptank, Psat_barg=Psat, **args,
            )
            metrics = _extract_metrics(out, hist)
        except Exception:
            return 1e6

        obj = direction * metrics.get(target_metric, 0.0)
        for cmetric, (op, threshold) in parsed_constraints.items():
            violation = _check_constraint(metrics.get(cmetric, 0.0), op, threshold)
            obj += penalty_weight * violation ** 2
        return obj

    return objective, n_evals


def _evaluate_final(cfg, x, speed, Pexit, Ptank, Psat, target_metric):
    """Run the final evaluation to get metrics and target value for the result."""
    final_cfg = apply_overrides(cfg, list(x))

    if target_metric == "stall_pressure":
        try:
            stall_p = _stall_pressure_with_overrides(final_cfg, speed, Ptank, Psat)
            return {"stall_pressure": stall_p}, stall_p
        except Exception:
            return {}, 0.0

    if target_metric == "fill_time":
        try:
            stall_p = _stall_pressure_with_overrides(final_cfg, speed, Ptank, Psat)
            proxy = -stall_p
            return {"fill_time_proxy": proxy, "stall_pressure": stall_p}, proxy
        except Exception:
            return {}, 0.0

    final_args = final_cfg.to_engine_args()
    try:
        out, hist = ICV_open(
            Pexit_barg=Pexit, speed_f=speed,
            Ptank_barg=Ptank, Psat_barg=Psat, **final_args,
        )
        final_metrics = _extract_metrics(out, hist)
    except Exception:
        final_metrics = {}
    return final_metrics, final_metrics.get(target_metric, 0.0)


@dataclass
class OptimizationResult:
    optimal_values: List[float]
    param_names: List[str]
    optimal_metrics: Dict[str, float]
    target: str
    target_value: float
    constraints: Dict[str, str]
    constraints_satisfied: bool
    base_hardware: str
    n_evals: int
    converged: bool

    def __repr__(self):
        status = "OK" if self.constraints_satisfied else "VIOLATED"
        return (
            f"OptimizationResult({self.target}={self.target_value:.4f}, "
            f"constraints={status}, {self.n_evals} evals)"
        )


def optimize(
    hardware: str = "old_icv",
    target: str = "max_mdot",
    speed: float = 0.65,
    Pexit: float = 500.0,
    constraints: Optional[Dict[str, str]] = None,
    Ptank: float = 7.0,
    Psat: float = 2.0,
    maxiter: int = 50,
    maxfun: Optional[int] = None,
    penalty_weight: float = 1000.0,
) -> OptimizationResult:
    """Find optimal hardware parameters subject to constraints using L-BFGS-B.

    A single-start local optimizer. For problems near the flow cliff
    discontinuity, consider ``multi_start_optimize`` or
    ``differential_evolution_optimize`` which are more robust to local minima.
    Minimum recommended ``maxiter`` is 30.

    Args:
        hardware: Base config name.
        target: "max_mdot", "max_mass_eff", "min_kWh", "min_Tc_peak",
                "max_stall_pressure", "min_fill_time"
        speed, Pexit: Operating conditions.
        constraints: e.g. {"Tc_peak_K": "<200", "mass_eff": ">0.1"}
        maxiter: Max optimizer iterations.
        maxfun: Max function evaluations (default: maxiter * 15).
        penalty_weight: Penalty multiplier for constraint violations.
    """
    constraints = constraints or {}
    cfg = load_config(hardware)
    direction, target_metric = _resolve_target(target)

    parsed_constraints = {}
    for cmetric, cstr in constraints.items():
        op, val = _parse_constraint(cstr)
        parsed_constraints[cmetric] = (op, val)

    effective_maxfun = maxfun or maxiter * 15

    objective, n_evals = _build_objective(
        cfg, direction, target_metric, parsed_constraints, penalty_weight,
        speed, Pexit, Ptank, Psat,
    )

    x0 = get_nominal_values(hardware)
    bounds = get_bounds()

    result = minimize(
        objective, x0=x0, method="L-BFGS-B", bounds=bounds,
        options={"maxiter": maxiter, "maxfun": effective_maxfun, "ftol": 1e-8},
    )

    final_metrics, target_value = _evaluate_final(
        cfg, result.x, speed, Pexit, Ptank, Psat, target_metric,
    )

    all_satisfied = True
    for cmetric, (op, threshold) in parsed_constraints.items():
        if _check_constraint(final_metrics.get(cmetric, 0.0), op, threshold) > 1e-6:
            all_satisfied = False

    return OptimizationResult(
        optimal_values=list(result.x),
        param_names=list(PARAM_NAMES),
        optimal_metrics=final_metrics,
        target=target,
        target_value=target_value,
        constraints=constraints,
        constraints_satisfied=all_satisfied,
        base_hardware=hardware,
        n_evals=n_evals[0],
        converged=result.success,
    )


def multi_start_optimize(
    hardware: str = "old_icv",
    target: str = "max_mdot",
    speed: float = 0.65,
    Pexit: float = 500.0,
    constraints: Optional[Dict[str, str]] = None,
    Ptank: float = 7.0,
    Psat: float = 2.0,
    n_starts: int = 5,
    maxiter: int = 30,
    seed: Optional[int] = None,
) -> OptimizationResult:
    """Multi-start L-BFGS-B optimization with Latin Hypercube Sampling.

    Generates ``n_starts`` starting points spread across the parameter space
    via LHS, runs L-BFGS-B from each, and returns the best result.

    This is more robust than single-start ``optimize()`` near the flow cliff
    discontinuity where L-BFGS-B tends to get trapped.

    Args:
        hardware: Base config name.
        target: "max_mdot", "max_mass_eff", "min_kWh", "min_Tc_peak",
                "max_stall_pressure", "min_fill_time"
        speed, Pexit: Operating conditions.
        constraints: e.g. {"Tc_peak_K": "<200", "mass_eff": ">0.1"}
        n_starts: Number of starting points to sample.
        maxiter: Max L-BFGS-B iterations per start.
        seed: Random seed for reproducibility.
    """
    constraints = constraints or {}
    cfg = load_config(hardware)
    direction, target_metric = _resolve_target(target)

    parsed_constraints = {}
    for cmetric, cstr in constraints.items():
        op, val = _parse_constraint(cstr)
        parsed_constraints[cmetric] = (op, val)

    bounds = get_bounds()
    n_params = len(bounds)
    effective_maxfun = maxiter * 15

    # Generate LHS starting points across parameter bounds
    sampler = qmc.LatinHypercube(d=n_params, seed=seed)
    samples = sampler.random(n=n_starts)
    lower = np.array([b[0] for b in bounds])
    upper = np.array([b[1] for b in bounds])
    start_points = qmc.scale(samples, lower, upper)

    best_result = None
    best_obj = 1e6
    total_evals = 0

    for i in range(n_starts):
        objective, n_evals = _build_objective(
            cfg, direction, target_metric, parsed_constraints, 1000.0,
            speed, Pexit, Ptank, Psat,
        )

        result = minimize(
            objective, x0=start_points[i], method="L-BFGS-B", bounds=bounds,
            options={"maxiter": maxiter, "maxfun": effective_maxfun, "ftol": 1e-8},
        )
        total_evals += n_evals[0]

        if result.fun < best_obj:
            best_obj = result.fun
            best_result = result

    if best_result is None:
        # Fallback: shouldn't happen unless n_starts=0
        best_result = minimize(
            lambda x: 1e6, x0=get_nominal_values(hardware),
            method="L-BFGS-B", bounds=bounds, options={"maxiter": 1},
        )
        total_evals = 0

    final_metrics, target_value = _evaluate_final(
        cfg, best_result.x, speed, Pexit, Ptank, Psat, target_metric,
    )

    all_satisfied = True
    for cmetric, (op, threshold) in parsed_constraints.items():
        if _check_constraint(final_metrics.get(cmetric, 0.0), op, threshold) > 1e-6:
            all_satisfied = False

    return OptimizationResult(
        optimal_values=list(best_result.x),
        param_names=list(PARAM_NAMES),
        optimal_metrics=final_metrics,
        target=target,
        target_value=target_value,
        constraints=constraints,
        constraints_satisfied=all_satisfied,
        base_hardware=hardware,
        n_evals=total_evals,
        converged=best_result.success,
    )


def differential_evolution_optimize(
    hardware: str = "old_icv",
    target: str = "max_mdot",
    speed: float = 0.65,
    Pexit: float = 500.0,
    constraints: Optional[Dict[str, str]] = None,
    Ptank: float = 7.0,
    Psat: float = 2.0,
    maxiter: int = 30,
    seed: Optional[int] = None,
    penalty_weight: float = 1000.0,
) -> OptimizationResult:
    """Global optimization via Differential Evolution.

    Uses ``scipy.optimize.differential_evolution`` which maintains a population
    of candidate solutions and is much less likely to get trapped by the flow
    cliff discontinuity than gradient-based methods.

    Args:
        hardware: Base config name.
        target: "max_mdot", "max_mass_eff", "min_kWh", "min_Tc_peak",
                "max_stall_pressure", "min_fill_time"
        speed, Pexit: Operating conditions.
        constraints: e.g. {"Tc_peak_K": "<200", "mass_eff": ">0.1"}
        maxiter: Max DE generations.
        seed: Random seed for reproducibility.
        penalty_weight: Penalty multiplier for constraint violations.
    """
    constraints = constraints or {}
    cfg = load_config(hardware)
    direction, target_metric = _resolve_target(target)

    parsed_constraints = {}
    for cmetric, cstr in constraints.items():
        op, val = _parse_constraint(cstr)
        parsed_constraints[cmetric] = (op, val)

    objective, n_evals = _build_objective(
        cfg, direction, target_metric, parsed_constraints, penalty_weight,
        speed, Pexit, Ptank, Psat,
    )

    bounds = get_bounds()

    result = differential_evolution(
        objective, bounds=bounds, maxiter=maxiter, seed=seed,
        tol=1e-8, polish=True,
    )

    final_metrics, target_value = _evaluate_final(
        cfg, result.x, speed, Pexit, Ptank, Psat, target_metric,
    )

    all_satisfied = True
    for cmetric, (op, threshold) in parsed_constraints.items():
        if _check_constraint(final_metrics.get(cmetric, 0.0), op, threshold) > 1e-6:
            all_satisfied = False

    return OptimizationResult(
        optimal_values=list(result.x),
        param_names=list(PARAM_NAMES),
        optimal_metrics=final_metrics,
        target=target,
        target_value=target_value,
        constraints=constraints,
        constraints_satisfied=all_satisfied,
        base_hardware=hardware,
        n_evals=n_evals[0],
        converged=result.success,
    )