"""Phaxor — Spring Design Engine (Python port)"""
import math

# Material Data
# Sut functions return MPa given diameter d in mm
WIRE_MATERIALS = {
    'music': {
        'label': 'Music Wire (ASTM A228)',
        'G': 81.7,
        'Sut': lambda d: 2211 / (d ** 0.145),
        'density': 7850,
    },
    'chrome-vanadium': {
        'label': 'Chrome-Vanadium (ASTM A232)',
        'G': 77.2,
        'Sut': lambda d: 2005 / (d ** 0.168),
        'density': 7850,
    },
    'chrome-silicon': {
        'label': 'Chrome-Silicon (ASTM A401)',
        'G': 77.2,
        'Sut': lambda d: 1974 / (d ** 0.108),
        'density': 7850,
    },
    'stainless-302': {
        'label': 'Stainless Steel 302 (ASTM A313)',
        'G': 69.0,
        'Sut': lambda d: 1867 / (d ** 0.146),
        'density': 7920,
    },
    'phosphor-bronze': {
        'label': 'Phosphor Bronze (ASTM B159)',
        'G': 41.4,
        'Sut': lambda d: 1000 / (d ** 0.028),
        'density': 8860,
    },
    'custom': {
        'label': 'Custom',
        'G': 79.3,
        'Sut': lambda d: 1500,
        'density': 7850,
    },
}

STANDARD_WIRES = [0.5, 0.6, 0.7, 0.8, 0.9, 1.0, 1.2, 1.4, 1.6, 1.8, 2.0, 2.3, 2.5, 2.8, 3.0, 3.2, 3.5, 4.0, 4.5, 5.0, 5.5, 6.0, 6.3, 7.0, 8.0, 9.0, 10.0, 11.0, 12.0, 13.0, 14.0, 16.0]

def solve_spring_design(inputs: dict) -> dict | None:
    """Helical Compression Spring Design."""
    load = float(inputs.get('load', 0))
    deflection = float(inputs.get('deflection', 0))
    spring_index_c = float(inputs.get('springIndexC', 0))
    mat_key = inputs.get('matKey', 'music')
    end_type = inputs.get('endType', 'closed-ground')
    fos = float(inputs.get('fos', 0))

    if load <= 0 or deflection <= 0 or spring_index_c <= 0 or fos <= 0:
        return None

    mat = WIRE_MATERIALS.get(mat_key, WIRE_MATERIALS['music'])
    g_modulus = mat['G'] * 1000 # MPa

    # Spring rate
    k = load / deflection

    c = spring_index_c
    kw = (4 * c - 1) / (4 * c - 4) + 0.615 / c

    # Iterative solver for wire diameter
    d = 3.0
    for _ in range(20):
        sut = mat['Sut'](d)
        tau_allow = 0.45 * sut / fos
        try:
            d_new = math.sqrt((kw * 8 * load * c) / (math.pi * tau_allow))
        except ValueError:
            d_new = d
        if abs(d_new - d) < 0.001:
            break
        d = d_new

    # Snap to nearest standard wire
    d_std = min(STANDARD_WIRES, key=lambda x: abs(x - d))
    d = d_std

    coil_d = c * d
    try:
        na = (g_modulus * d) / (8 * (c ** 3) * k)
    except ZeroDivisionError:
        na = 0

    nt = 0
    ls = 0
    if end_type == 'closed-ground':
        nt = na + 2
        ls = (nt + 1) * d
    elif end_type == 'closed':
        nt = na + 2
        ls = (nt + 1) * d
    elif end_type == 'open':
        nt = na
        ls = (nt + 1) * d
    elif end_type == 'open-ground':
        nt = na + 1
        ls = nt * d
    else:
        nt = na + 2
        ls = (nt + 1) * d
    
    lf = ls + deflection * 1.15
    pitch = (lf - 2 * d) / na if na > 0 else 0

    tau_actual = kw * (8 * load * coil_d) / (math.pi * d ** 3)
    tau_uncorrected = (8 * load * coil_d) / (math.pi * d ** 3)

    rho = mat['density']
    freq = 0
    if na > 0 and coil_d > 0:
        term1 = (d / 1000) / (2 * math.pi * na * (coil_d / 1000) ** 2)
        term2 = math.sqrt((g_modulus * 1e6) / (2 * rho))
        freq = term1 * term2
    
    buckling_ratio = lf / coil_d if coil_d > 0 else 0
    buckling_ok = buckling_ratio < 4

    wire_length = math.pi * coil_d * nt
    volume = (math.pi * d ** 2 / 4) * wire_length
    weight = volume * rho * 1e-9 * 1000

    return {
        'wireD': d,
        'coilD': float(f"{coil_d:.2f}"),
        'springIndex': c,
        'wahlFactor': float(f"{kw:.3f}"),
        'shearStress': float(f"{tau_actual:.2f}"),
        'shearStressUncorrected': float(f"{tau_uncorrected:.2f}"),
        'springRate': float(f"{k:.2f}"),
        'activeCoils': float(f"{na:.1f}"),
        'totalCoils': float(f"{nt:.1f}"),
        'freeLength': float(f"{lf:.2f}"),
        'solidLength': float(f"{ls:.2f}"),
        'deflectionAtLoad': deflection,
        'naturalFreq': float(f"{freq:.1f}"),
        'bucklingRatio': float(f"{buckling_ratio:.2f}"),
        'bucklingOK': buckling_ok,
        'weight': float(f"{weight:.2f}"),
        'outerD': float(f"{coil_d + d:.2f}"),
        'innerD': float(f"{coil_d - d:.2f}"),
        'pitch': float(f"{pitch:.2f}")
    }
