Smart parameter search for VedEf - targeting power and temperature constraints.

import importlib.util
from pathlib import Path

import numpy as np

# Import setup_vedef_problem from the numbered file
spec = importlib.util.spec_from_file_location(
    "problem_module", Path(__file__).parent / "1_setup_vedef_problem.py"
)
problem_module = importlib.util.module_from_spec(spec)
spec.loader.exec_module(problem_module)
setup_vedef_problem = problem_module.setup_vedef_problem


def evaluate_case(V, e, d_t, f, E_p):
    """Evaluate a parameter case."""
    try:
        prob = setup_vedef_problem(target_power=25000.0, max_energy_density=1e9, as_subsystem=False)

        # Fixed parameters
        prob.set_val("pulse_duration", 1e-6, units="s")
        prob.set_val("TP_QM", 160.0 / 86400.0, units="kg/s")
        prob.set_val("ThermalDiameterModel.gas_density", 0.717, units="kg/m**3")
        prob.set_val("ThermalDiameterModel.gas_heat_capacity", 2200.0, units="J/(kg*K)")
        prob.set_val("ThermalDiameterModel.thermal_conductivity", 0.08, units="W/(m*K)")
        prob.set_val("InitialTemperatureModel.gas_density", 0.717, units="kg/m**3")
        prob.set_val("BreakdownModel.gas_properties_pressure", 101325.0, units="Pa")

        # Design variables (using ARC002 naming convention)
        prob.set_val("G_UMAX_OUT", V, units="V")
        prob.set_val("G_e", e, units="m")
        prob.set_val("TP_D_OUT", d_t, units="m")
        prob.set_val("G_F", f, units="Hz")
        prob.set_val("G_Ep", E_p, units="J")

        prob.run_model()

        results = {
            "V": V,
            "e": e,
            "d_t": d_t,
            "f": f,
            "E_p": E_p,
            "T_0": prob.get_val("T_0", units="K")[0],
            "T_max": prob.get_val("T_max", units="K")[0],
            "breakdown_margin": prob.get_val("breakdown_margin", units="V")[0],
            "coverage_margin": prob.get_val("coverage_margin")[0],
            "density_margin": prob.get_val("density_margin", units="J/m**3")[0],
            "power_error": prob.get_val("PowerConstraint.power_error")[0],
        }

        # Check if all values are finite
        if not all(
            np.isfinite(v) for k, v in results.items() if k not in ["V", "e", "d_t", "f", "E_p"]
        ):
            return None

        return results

    except Exception:
        return None


def check_feasibility(results):
    """Check if all constraints are satisfied."""
    if results is None:
        return False

    mobility_ratio = results["T_max"] / (results["T_0"] + 1e-10)

    checks = {
        "breakdown": results["breakdown_margin"] > 0,
        "coverage": results["coverage_margin"] > 0,
        "density": results["density_margin"] > 0,
        "power": results["power_error"] < 0.1,  # 10% tolerance
        "mobility": mobility_ratio > 10.0,
    }

    return all(checks.values()), checks, mobility_ratio


print("=" * 80)
print("VedEf Parameter Search - Targeting Power Constraint")
print("=" * 80)

# Strategy: Fix f and E_p to satisfy power constraint, then vary geometry
target_power = 25000.0  # 25 kW

# Test different f, E_p combinations that give target power
param_sets = [
    # (f_kHz, E_p_mJ)
    (25, 1000),  # 25 kHz, 1000 mJ = 25 kW
    (50, 500),  # 50 kHz, 500 mJ = 25 kW
    (100, 250),  # 100 kHz, 250 mJ = 25 kW
    (125, 200),  # 125 kHz, 200 mJ = 25 kW
    (200, 125),  # 200 kHz, 125 mJ = 25 kW
]

# Geometry variations
V_range = [15e3, 20e3, 30e3, 40e3]  # kV
e_range = [0.005, 0.01, 0.015, 0.02]  # m
d_t_range = [0.015, 0.02, 0.025]  # m

feasible_cases = []
best_case = None
best_violations = 10

print(f"\nSearching {len(param_sets)} power settings x {len(V_range)} voltages x", end="")
print(f" {len(e_range)} gaps x {len(d_t_range)} diameters")
print(f"Total: {len(param_sets) * len(V_range) * len(e_range) * len(d_t_range)} cases\n")

for f_kHz, E_p_mJ in param_sets:
    f = f_kHz * 1e3
    E_p = E_p_mJ * 1e-3

    print(f"\nTesting f={f_kHz}kHz, E_p={E_p_mJ}mJ (Power={f * E_p / 1e3:.1f}kW)")
    print("-" * 80)

    for V in V_range:
        for e in e_range:
            for d_t in d_t_range:
                results = evaluate_case(V, e, d_t, f, E_p)

                if results is None:
                    continue

                feasible, checks, mobility_ratio = check_feasibility(results)

                if feasible:
                    feasible_cases.append(results)
                    print(
                        f"[OK] FEASIBLE: V={V / 1e3:.0f}kV, e={e * 1e3:.0f}mm, d_t={d_t * 1e3:.0f}mm"
                    )
                    print(
                        f"     T_max/T_0={mobility_ratio:.1f}, Power_err={results['power_error'] * 100:.1f}%"
                    )

                # Track best case (minimum violations)
                n_violations = sum(not v for v in checks.values())
                if n_violations < best_violations:
                    best_violations = n_violations
                    best_case = results
                    best_checks = checks

print("\n" + "=" * 80)
print("SUMMARY")
print("=" * 80)

if feasible_cases:
    print(f"\n[SUCCESS] Found {len(feasible_cases)} feasible parameter sets!\n")
    print("Best feasible case:")
    case = feasible_cases[0]
    print(f"  V = {case['V'] / 1e3:.1f} kV")
    print(f"  e = {case['e'] * 1e3:.1f} mm")
    print(f"  d_t = {case['d_t'] * 1e3:.1f} mm")
    print(f"  f = {case['f'] / 1e3:.1f} kHz")
    print(f"  E_p = {case['E_p'] * 1e3:.1f} mJ")
    print(f"\n  T_0 = {case['T_0']:.1f} K")
    print(f"  T_max = {case['T_max']:.1f} K")
    print(f"  T_max/T_0 = {case['T_max'] / (case['T_0'] + 1e-10):.1f}")
    print(f"  Power error = {case['power_error'] * 100:.1f}%")

else:
    print(f"\n[NO FEASIBLE SOLUTION] Closest case had {best_violations} violations:\n")
    if best_case:
        print(f"  V = {best_case['V'] / 1e3:.1f} kV")
        print(f"  e = {best_case['e'] * 1e3:.1f} mm")
        print(f"  d_t = {best_case['d_t'] * 1e3:.1f} mm")
        print(f"  f = {best_case['f'] / 1e3:.1f} kHz")
        print(f"  E_p = {best_case['E_p'] * 1e3:.1f} mJ")
        print(f"\n  T_0 = {best_case['T_0']:.1f} K")
        print(f"  T_max = {best_case['T_max']:.1f} K")
        print(f"  T_max/T_0 = {best_case['T_max'] / (best_case['T_0'] + 1e-10):.1f}")
        print(f"  Power error = {best_case['power_error'] * 100:.1f}%")
        print("\n  Constraint violations:")
        for name, satisfied in best_checks.items():
            status = "[OK]" if satisfied else "[X]"
            print(f"    {status} {name}")

print("\n" + "=" * 80)