cryosim/optimization/design_optimizer.py

"""Hardware dimension optimizer (C4).

Optimizes physical hardware dimensions (bore, port diameter, spring constant,
etc.) that can be handed directly to a machinist, rather than the abstract
calibration knobs used by ``optimizer.py``.

Uses ``scipy.optimize.differential_evolution`` because the physical dimensions
interact nonlinearly and the flow-cliff discontinuity traps gradient methods.
"""

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

import numpy as np
from scipy.optimize import differential_evolution

from cryosim.hardware.design_params import (
    HardwareDesignParams,
    design_to_engine_config,
)
from cryosim.engine.fast import ICV_open
from cryosim.optimization.optimizer import (
    _resolve_target,
    _parse_constraint,
    _check_constraint,
    _extract_metrics,
    _stall_pressure_with_overrides,
    _fill_time_proxy_with_overrides,
)


@dataclass
class DesignOptimizationResult:
    """Result of a hardware-dimension optimization run."""

    optimal_design: HardwareDesignParams
    optimal_metrics: Dict[str, float]
    target: str
    target_value: float
    base_hardware: str
    n_evals: int
    converged: bool
    param_changes: Dict[str, tuple]   # {name: (before, after, pct_change)}
    notes: List[str] = field(default_factory=list)

    def __repr__(self) -> str:
        changed = sum(1 for _, _, pct in self.param_changes.values() if abs(pct) > 1.0)
        return (
            f"DesignOptimizationResult({self.target}={self.target_value:.4f}, "
            f"{changed}/{len(self.param_changes)} params changed, "
            f"{self.n_evals} evals, converged={self.converged})"
        )


# ---------------------------------------------------------------------------
# Manufacturing notes
# ---------------------------------------------------------------------------

def _manufacturing_notes(design: HardwareDesignParams) -> List[str]:
    """Generate practical manufacturing notes for extreme parameter values."""
    notes: List[str] = []
    if design.icv_port_dia_mm < 4.0:
        notes.append(
            f"ICV port {design.icv_port_dia_mm:.1f} mm may be difficult to "
            f"machine and prone to clogging."
        )
    if design.dcv_port_dia_mm < 5.0:
        notes.append(
            f"DCV port {design.dcv_port_dia_mm:.1f} mm is small; verify flow "
            f"capacity at target mass flow."
        )
    if design.bore_mm < 18.0:
        notes.append(
            f"Bore {design.bore_mm:.1f} mm is very small; swept volume will "
            f"limit throughput."
        )
    if design.bore_mm > 45.0:
        notes.append(
            f"Bore {design.bore_mm:.1f} mm is large; verify seal availability "
            f"and structural margin."
        )
    if design.dead_volume_frac > 0.06:
        notes.append(
            f"Dead volume fraction {design.dead_volume_frac:.3f} is high; "
            f"consider tighter piston-to-head clearance."
        )
    if design.icv_spring_force_N > 40.0:
        notes.append(
            f"ICV spring preload {design.icv_spring_force_N:.1f} N is high; "
            f"may delay ICV opening at low pressures."
        )
    if design.dcv_spring_force_N > 15.0:
        notes.append(
            f"DCV spring preload {design.dcv_spring_force_N:.1f} N is high; "
            f"verify DCV closes before backflow at high exit pressures."
        )
    return notes


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

def _build_design_objective(
    base_hardware: str,
    direction: float,
    target_metric: str,
    parsed_constraints: Dict[str, tuple],
    penalty_weight: float,
    speed: float,
    Pexit: float,
    Ptank: float,
    Psat: float,
    all_param_names: List[str],
    optimized_names: List[str],
    fixed_design: HardwareDesignParams,
):
    """Build a closure that maps a design-parameter vector to a scalar objective.

    Parameters
    ----------
    all_param_names : list[str]
        Full ordered list of design parameter names.
    optimized_names : list[str]
        Subset of names being optimized (others stay at ``fixed_design`` values).
    fixed_design : HardwareDesignParams
        Baseline design — non-optimized params are taken from here.
    """
    n_evals = [0]

    def objective(x):
        n_evals[0] += 1

        # Build a design from the vector
        updates = {name: float(val) for name, val in zip(optimized_names, x)}
        design = fixed_design.with_updates(updates)

        try:
            cfg = design_to_engine_config(design, base_hardware=base_hardware)
        except Exception:
            return 1e6

        # --- fill-level objectives ---
        if target_metric == "stall_pressure":
            try:
                stall_p = _stall_pressure_with_overrides(cfg, speed, Ptank, Psat)
            except Exception:
                return 1e6
            obj = direction * stall_p
            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(
                    cfg, speed, Pexit, Ptank, Psat,
                )
            except Exception:
                return 1e6
            return direction * proxy

        # --- single-cycle objectives ---
        args = 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


# ---------------------------------------------------------------------------
# Final evaluation
# ---------------------------------------------------------------------------

def _evaluate_design_final(
    design: HardwareDesignParams,
    base_hardware: str,
    speed: float,
    Pexit: float,
    Ptank: float,
    Psat: float,
    target_metric: str,
) -> tuple:
    """Run a final engine evaluation and return (metrics_dict, target_value)."""
    cfg = design_to_engine_config(design, base_hardware=base_hardware)

    if target_metric == "stall_pressure":
        try:
            stall_p = _stall_pressure_with_overrides(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(cfg, speed, Ptank, Psat)
            proxy = -stall_p
            return {"fill_time_proxy": proxy, "stall_pressure": stall_p}, proxy
        except Exception:
            return {}, 0.0

    args = 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:
        metrics = {}

    return metrics, metrics.get(target_metric, 0.0)


# ---------------------------------------------------------------------------
# Public API
# ---------------------------------------------------------------------------

def design_optimize(
    target: str = "max_stall_pressure",
    speed: float = 0.65,
    Pexit: float = 500.0,
    base_hardware: str = "old_icv",
    constraints: Optional[Dict[str, str]] = None,
    params_to_optimize: Optional[List[str]] = None,
    maxiter: int = 30,
    seed: Optional[int] = None,
    Ptank: float = 7.0,
    Psat: float = 2.0,
    penalty_weight: float = 1000.0,
) -> DesignOptimizationResult:
    """Optimize physical hardware dimensions via Differential Evolution.

    Unlike ``optimize()`` which searches over 7 abstract calibration knobs,
    this function searches over measurable hardware dimensions (bore diameter,
    port diameter, spring constant, etc.) that can be handed directly to a
    machinist.

    Args:
        target: Optimization target — "max_mdot", "max_stall_pressure",
                "min_kWh", "min_fill_time", etc.
        speed: Speed fraction [0-1].
        Pexit: Discharge pressure target [barg] (for single-cycle targets).
        base_hardware: Base config name used to fill in non-optimized params.
        constraints: e.g. {"Tc_peak_K": "<200", "mass_eff": ">0.1"}.
        params_to_optimize: Subset of parameter names to optimize. If None,
                            all 11 parameters are optimized.
        maxiter: Max differential-evolution generations.
        seed: Random seed for reproducibility.
        Ptank: Inlet tank pressure [barg].
        Psat: Saturation pressure [barg].
        penalty_weight: Penalty multiplier for constraint violations.

    Returns:
        DesignOptimizationResult with optimal dimensions, metrics, and
        manufacturing notes.
    """
    constraints = constraints or {}
    baseline = HardwareDesignParams()
    all_names = baseline.param_names

    # Validate params_to_optimize
    if params_to_optimize is not None:
        unknown = set(params_to_optimize) - set(all_names)
        if unknown:
            raise ValueError(
                f"Unknown design parameters: {unknown}. "
                f"Valid names: {all_names}"
            )
        opt_names = list(params_to_optimize)
    else:
        opt_names = list(all_names)

    direction, target_metric = _resolve_target(target)

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

    # Build bounds for the optimized subset only
    opt_bounds = baseline.subset_bounds(opt_names)

    objective, n_evals = _build_design_objective(
        base_hardware=base_hardware,
        direction=direction,
        target_metric=target_metric,
        parsed_constraints=parsed_constraints,
        penalty_weight=penalty_weight,
        speed=speed,
        Pexit=Pexit,
        Ptank=Ptank,
        Psat=Psat,
        all_param_names=all_names,
        optimized_names=opt_names,
        fixed_design=baseline,
    )

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

    # Reconstruct optimal design
    opt_updates = {name: float(val) for name, val in zip(opt_names, result.x)}
    optimal_design = baseline.with_updates(opt_updates)

    # Final evaluation
    final_metrics, target_value = _evaluate_design_final(
        optimal_design, base_hardware, speed, Pexit, Ptank, Psat, target_metric,
    )

    # Compute per-parameter changes
    param_changes: Dict[str, tuple] = {}
    for name in all_names:
        before = getattr(baseline, name)
        after = getattr(optimal_design, name)
        if abs(before) > 1e-12:
            pct = 100.0 * (after - before) / before
        else:
            pct = 0.0 if abs(after) < 1e-12 else float("inf")
        param_changes[name] = (before, after, pct)

    notes = _manufacturing_notes(optimal_design)

    return DesignOptimizationResult(
        optimal_design=optimal_design,
        optimal_metrics=final_metrics,
        target=target,
        target_value=target_value,
        base_hardware=base_hardware,
        n_evals=n_evals[0],
        converged=result.success,
        param_changes=param_changes,
        notes=notes,
    )