sweep_analysis.py

"""Deterministic sweep analysis for cryogenic pump simulator results.

Provides sensitivity ranking, Pareto frontier computation, and text-based
insights for 1D and 2D parameter sweeps. No LLM needed -- pure numerical
analysis using finite differences, dominance checks, and feasibility filters.

Designed to consume sweep result dicts produced by SweepAnalyzer in
scenario_analyzer.py.
"""

import numpy as np
from typing import Dict, List, Optional


# ---------------------------------------------------------------------------
# 1. Sensitivity ranking
# ---------------------------------------------------------------------------

def sensitivity_ranking(
    sweep_results: List[dict],
    param_key: str = 'param_value',
    metric: str = 'mdot_kgpm',
) -> Dict:
    """Compute finite-difference sensitivity of *metric* w.r.t. *param_key*.

    Parameters
    ----------
    sweep_results : list of dict
        Each dict must contain at least *param_key* and *metric* keys.
        Only rows where ``success`` is truthy (if present) and *metric*
        is finite are used.
    param_key : str
        Name of the independent variable field (default ``'param_value'``).
    metric : str
        Name of the dependent variable field (default ``'mdot_kgpm'``).

    Returns
    -------
    dict with keys:
        mean_gradient   – average |dy/dx| over all consecutive pairs
        max_gradient    – maximum |dy/dx|
        direction       – 'increasing', 'decreasing', or 'non-monotonic'
        elasticity      – median (dy/y) / (dx/x), or NaN if undefined
        n_points        – number of valid points used
    """
    # Filter to successful, finite-metric rows
    filtered = []
    for r in sweep_results:
        if not r.get('success', True):
            continue
        val = r.get(metric)
        if val is None or (isinstance(val, float) and np.isnan(val)):
            continue
        pv = r.get(param_key)
        if pv is None or (isinstance(pv, float) and np.isnan(pv)):
            continue
        filtered.append(r)

    n = len(filtered)
    if n < 2:
        return {
            'mean_gradient': float('nan'),
            'max_gradient': float('nan'),
            'direction': 'insufficient_data',
            'elasticity': float('nan'),
            'n_points': n,
        }

    # Sort by param value
    filtered.sort(key=lambda r: r[param_key])

    xs = np.array([r[param_key] for r in filtered], dtype=np.float64)
    ys = np.array([r[metric] for r in filtered], dtype=np.float64)

    dx = np.diff(xs)
    dy = np.diff(ys)

    # Gradients (skip zero-dx pairs)
    valid = np.abs(dx) > 0.0
    if not np.any(valid):
        return {
            'mean_gradient': float('nan'),
            'max_gradient': float('nan'),
            'direction': 'constant_param',
            'elasticity': float('nan'),
            'n_points': n,
        }

    grads = dy[valid] / dx[valid]

    mean_gradient = float(np.mean(np.abs(grads)))
    max_gradient = float(np.max(np.abs(grads)))

    # Direction
    if np.all(grads >= 0):
        direction = 'increasing'
    elif np.all(grads <= 0):
        direction = 'decreasing'
    else:
        direction = 'non-monotonic'

    # Elasticity: (dy/y_mid) / (dx/x_mid) for each consecutive pair
    x_mid = 0.5 * (xs[:-1] + xs[1:])
    y_mid = 0.5 * (ys[:-1] + ys[1:])

    # Only compute where both midpoints are nonzero and dx is nonzero
    e_valid = valid & (np.abs(x_mid) > 0.0) & (np.abs(y_mid) > 0.0)
    if np.any(e_valid):
        rel_dy = dy[e_valid] / y_mid[e_valid]
        rel_dx = dx[e_valid] / x_mid[e_valid]
        # Guard against zero rel_dx (shouldn't happen given valid filter, but be safe)
        e_mask = np.abs(rel_dx) > 0.0
        if np.any(e_mask):
            elasticities = rel_dy[e_mask] / rel_dx[e_mask]
            elasticity = float(np.median(elasticities))
        else:
            elasticity = float('nan')
    else:
        elasticity = float('nan')

    return {
        'mean_gradient': mean_gradient,
        'max_gradient': max_gradient,
        'direction': direction,
        'elasticity': elasticity,
        'n_points': n,
    }


# ---------------------------------------------------------------------------
# 2. Pareto frontier
# ---------------------------------------------------------------------------

def pareto_frontier(
    results: List[dict],
    maximize: Optional[List[str]] = None,
    minimize: Optional[List[str]] = None,
) -> List[dict]:
    """Find Pareto-optimal (non-dominated) points from sweep results.

    Parameters
    ----------
    results : list of dict
        Each dict should contain the objective keys. Only rows with finite
        values for all objectives are considered.
    maximize : list of str or None
        Objective keys to maximize. Defaults to ``['mdot_kgpm']``.
    minimize : list of str or None
        Objective keys to minimize. Defaults to ``[]``.

    Returns
    -------
    list of dict
        Subset of *results* that are Pareto-optimal. Order is preserved.

    Notes
    -----
    A point *a* dominates *b* iff *a* is at least as good in ALL objectives
    and strictly better in at least one. O(n^2) pairwise comparison -- fine
    for typical sweep sizes (10-100 points).
    """
    if maximize is None:
        maximize = ['mdot_kgpm']
    if minimize is None:
        minimize = []

    all_keys = list(maximize) + list(minimize)
    if not all_keys:
        return list(results)

    # Filter to rows with finite values for all objectives
    candidates = []
    for r in results:
        ok = True
        for k in all_keys:
            v = r.get(k)
            if v is None or (isinstance(v, float) and np.isnan(v)):
                ok = False
                break
        if ok:
            candidates.append(r)

    if not candidates:
        return []

    def _dominates(a: dict, b: dict) -> bool:
        """Return True if *a* dominates *b*."""
        at_least_as_good = True
        strictly_better = False
        for k in maximize:
            av, bv = a[k], b[k]
            if av < bv:
                at_least_as_good = False
                return False
            if av > bv:
                strictly_better = True
        for k in minimize:
            av, bv = a[k], b[k]
            if av > bv:
                at_least_as_good = False
                return False
            if av < bv:
                strictly_better = True
        return at_least_as_good and strictly_better

    frontier = []
    for i, candidate in enumerate(candidates):
        dominated = False
        for j, other in enumerate(candidates):
            if i == j:
                continue
            if _dominates(other, candidate):
                dominated = True
                break
        if not dominated:
            frontier.append(candidate)

    return frontier


# ---------------------------------------------------------------------------
# 3. 1D sweep insights
# ---------------------------------------------------------------------------

def sweep_insights(sweep_result: dict) -> str:
    """Generate a plain-text summary of a 1D parameter sweep.

    Parameters
    ----------
    sweep_result : dict
        Expected keys:
            param_name  – str, e.g. 'ICVport_mm'
            results     – list of dicts, each with at least:
                param_value, mdot_kgpm, mass_eff, pc_peak_barg, success

    Returns
    -------
    str
        Multi-line plain text report.
    """
    param_name = sweep_result.get('param_name', 'unknown')
    results = sweep_result.get('results', [])

    if not results:
        return f"Sweep of {param_name}: no results."

    # Separate successful runs
    ok = [r for r in results if r.get('success', False)
          and not np.isnan(r.get('mdot_kgpm', float('nan')))]

    lines = [f"=== Sweep: {param_name} ({len(results)} points, {len(ok)} successful) ==="]

    if not ok:
        lines.append("No successful runs -- all points failed or returned NaN.")
        return '\n'.join(lines)

    # Sensitivity
    sens_mdot = sensitivity_ranking(ok, metric='mdot_kgpm')
    sens_eff = sensitivity_ranking(ok, metric='mass_eff')

    lines.append("")
    lines.append(f"Sensitivity (mdot): direction={sens_mdot['direction']}, "
                 f"mean |grad|={sens_mdot['mean_gradient']:.4g}, "
                 f"elasticity={sens_mdot['elasticity']:.3g}")
    lines.append(f"Sensitivity (eff):  direction={sens_eff['direction']}, "
                 f"mean |grad|={sens_eff['mean_gradient']:.4g}, "
                 f"elasticity={sens_eff['elasticity']:.3g}")

    # Best flow
    best_flow = max(ok, key=lambda r: r['mdot_kgpm'])
    lines.append("")
    lines.append(f"Best flow:      {best_flow['mdot_kgpm']:.4f} kg/min "
                 f"at {param_name}={best_flow['param_value']:.4g}")

    # Best efficiency
    best_eff = max(ok, key=lambda r: r['mass_eff'])
    lines.append(f"Best efficiency: {best_eff['mass_eff']:.4f} "
                 f"at {param_name}={best_eff['param_value']:.4g}")

    # Pareto frontier (maximize mdot and efficiency simultaneously)
    pf = pareto_frontier(ok, maximize=['mdot_kgpm', 'mass_eff'])
    lines.append(f"Pareto frontier: {len(pf)} non-dominated point(s)")

    # MAWP constraint violations
    mawp_violations = [r for r in ok if r.get('pc_peak_barg', 0) >= 960]
    if mawp_violations:
        pct = 100.0 * len(mawp_violations) / len(ok)
        worst = max(mawp_violations, key=lambda r: r['pc_peak_barg'])
        lines.append(f"MAWP violations: {len(mawp_violations)}/{len(ok)} points "
                     f"({pct:.0f}%) exceed 960 barg "
                     f"(worst: {worst['pc_peak_barg']:.1f} barg "
                     f"at {param_name}={worst['param_value']:.4g})")
    else:
        lines.append("MAWP violations: none (all points below 960 barg)")

    return '\n'.join(lines)


# ---------------------------------------------------------------------------
# 4. 2D sweep insights
# ---------------------------------------------------------------------------

def sweep_2d_insights(sr: dict) -> str:
    """Generate a plain-text summary of a 2D parameter sweep.

    Parameters
    ----------
    sr : dict
        Expected keys:
            param_x  – str, name of X-axis parameter
            param_y  – str, name of Y-axis parameter
            vals_x   – 1D array-like, X parameter values
            vals_y   – 1D array-like, Y parameter values
            grid     – dict of 2D numpy arrays keyed by metric name.
                       Required: 'mdot_kgpm', 'mass_eff', 'pc_peak_barg'.
                       Optional: 'Tc_peak_K'.
                       Shape: (len(vals_y), len(vals_x)).

    Returns
    -------
    str
        Multi-line plain text report.
    """
    param_x = sr.get('param_x', 'x')
    param_y = sr.get('param_y', 'y')
    vals_x = np.asarray(sr.get('vals_x', []))
    vals_y = np.asarray(sr.get('vals_y', []))
    grid = sr.get('grid', {})

    lines = [f"=== 2D Sweep: {param_x} x {param_y} "
             f"({len(vals_x)} x {len(vals_y)} = {len(vals_x) * len(vals_y)} points) ==="]

    mdot_grid = grid.get('mdot_kgpm')
    if mdot_grid is None:
        lines.append("No mdot_kgpm grid found -- cannot analyse.")
        return '\n'.join(lines)

    mdot_grid = np.asarray(mdot_grid, dtype=np.float64)
    finite_mask = np.isfinite(mdot_grid)
    n_finite = int(np.sum(finite_mask))
    n_total = mdot_grid.size

    lines.append(f"Valid points: {n_finite}/{n_total}")

    if n_finite == 0:
        lines.append("All points failed or returned NaN.")
        return '\n'.join(lines)

    # Global best mdot
    best_idx = np.unravel_index(np.nanargmax(mdot_grid), mdot_grid.shape)
    best_mdot = float(mdot_grid[best_idx])
    best_y = float(vals_y[best_idx[0]]) if best_idx[0] < len(vals_y) else float('nan')
    best_x = float(vals_x[best_idx[1]]) if best_idx[1] < len(vals_x) else float('nan')

    lines.append("")
    lines.append(f"Global best mdot: {best_mdot:.4f} kg/min "
                 f"at {param_x}={best_x:.4g}, {param_y}={best_y:.4g}")

    # Efficiency at best mdot
    eff_grid = grid.get('mass_eff')
    if eff_grid is not None:
        eff_grid = np.asarray(eff_grid, dtype=np.float64)
        eff_at_best = float(eff_grid[best_idx])
        lines.append(f"Efficiency at best mdot: {eff_at_best:.4f}")

    # Feasible region (pc_peak_barg < 960)
    pc_grid = grid.get('pc_peak_barg')
    if pc_grid is not None:
        pc_grid = np.asarray(pc_grid, dtype=np.float64)
        feasible = finite_mask & np.isfinite(pc_grid) & (pc_grid < 960.0)
        n_feasible = int(np.sum(feasible))
        lines.append("")
        lines.append(f"Feasible region (pc < 960 barg): "
                     f"{n_feasible}/{n_finite} valid points "
                     f"({100.0 * n_feasible / max(n_finite, 1):.0f}%)")

        if n_feasible > 0:
            # Best feasible mdot
            mdot_feasible = np.where(feasible, mdot_grid, np.nan)
            feas_idx = np.unravel_index(np.nanargmax(mdot_feasible), mdot_feasible.shape)
            feas_mdot = float(mdot_feasible[feas_idx])
            feas_x = float(vals_x[feas_idx[1]]) if feas_idx[1] < len(vals_x) else float('nan')
            feas_y = float(vals_y[feas_idx[0]]) if feas_idx[0] < len(vals_y) else float('nan')
            lines.append(f"Best feasible mdot: {feas_mdot:.4f} kg/min "
                         f"at {param_x}={feas_x:.4g}, {param_y}={feas_y:.4g}")
        else:
            lines.append("No feasible points -- all exceed MAWP constraint.")
    else:
        lines.append("(pc_peak_barg grid not provided -- skipping feasibility check)")

    # Dominant axis: which parameter has more impact on mdot?
    lines.append("")
    _dominant_axis(lines, mdot_grid, vals_x, vals_y, param_x, param_y, finite_mask)

    return '\n'.join(lines)


def _dominant_axis(
    lines: list,
    mdot_grid: np.ndarray,
    vals_x: np.ndarray,
    vals_y: np.ndarray,
    param_x: str,
    param_y: str,
    finite_mask: np.ndarray,
) -> None:
    """Append dominant-axis analysis to *lines* (in-place).

    Compares the mean absolute range of mdot along each axis to determine
    which parameter has more influence on flow rate.
    """
    ny, nx = mdot_grid.shape

    # Range of mdot along X for each Y row (holding Y constant, varying X)
    range_along_x = []
    for j in range(ny):
        row = mdot_grid[j, :]
        row_valid = row[finite_mask[j, :]]
        if len(row_valid) >= 2:
            range_along_x.append(float(np.max(row_valid) - np.min(row_valid)))

    # Range of mdot along Y for each X column (holding X constant, varying Y)
    range_along_y = []
    for i in range(nx):
        col = mdot_grid[:, i]
        col_valid = col[finite_mask[:, i]]
        if len(col_valid) >= 2:
            range_along_y.append(float(np.max(col_valid) - np.min(col_valid)))

    mean_range_x = float(np.mean(range_along_x)) if range_along_x else 0.0
    mean_range_y = float(np.mean(range_along_y)) if range_along_y else 0.0

    lines.append(f"Mean mdot range varying {param_x} (hold {param_y}): "
                 f"{mean_range_x:.4f} kg/min")
    lines.append(f"Mean mdot range varying {param_y} (hold {param_x}): "
                 f"{mean_range_y:.4f} kg/min")

    if mean_range_x > 0 or mean_range_y > 0:
        if mean_range_x > mean_range_y * 1.1:
            lines.append(f"Dominant axis: {param_x} "
                         f"({mean_range_x / max(mean_range_y, 1e-12):.1f}x more impact)")
        elif mean_range_y > mean_range_x * 1.1:
            lines.append(f"Dominant axis: {param_y} "
                         f"({mean_range_y / max(mean_range_x, 1e-12):.1f}x more impact)")
        else:
            lines.append("Dominant axis: roughly equal sensitivity to both parameters")
    else:
        lines.append("Dominant axis: insufficient data to determine")