"""Валидация GPU-батч-интегратора против scipy-эталона.

Батч-путь (numpy/cupy, фикс. шаг RK4) обязан давать те же числа, что
проверенный scipy solve_ivp на той же физике — иначе ускорение бессмысленно.
Здесь backend всегда numpy (GPU для теста не нужен: код один и тот же).
"""

import numpy as np
from scipy.integrate import solve_ivp

from gausse.gpu.backend import get_backend
from gausse.gpu.batch_integrator import integrate_batch_discharge, params_from_models
from gausse.physics import circuit
from gausse.physics.circuit import StageCircuitParams
from gausse.physics.inductance import (
    CoilInductanceModel,
    air_core_inductance_wheeler,
    demagnetizing_factor_prolate,
    effective_permeability,
    winding_geometry,
)


def _setup(turns_per_layer, layers, tube_od, slug_d, slug_len, mu_r, b_sat, cap_uf, esr, charge_v, wire_rho, wire_gauge):
    """Строит (model, circuit_params, q0, x_fire, v_fire) как это делает run_stage."""
    wire_od = (wire_gauge + 0.05) / 1000
    geo = winding_geometry(tube_od, wire_od, turns_per_layer, layers)
    area = np.pi * (wire_gauge / 1000 / 2) ** 2
    r_wire = wire_rho * geo.total_wire_length_m / area
    l_air = air_core_inductance_wheeler(geo.mean_radius_m, geo.coil_length_m, geo.radial_depth_m, geo.total_turns)
    aspect = slug_len / slug_d
    mu_eff = effective_permeability(mu_r, demagnetizing_factor_prolate(aspect))
    mass = np.pi * (slug_d / 2) ** 2 * slug_len * 7850.0
    model = CoilInductanceModel(
        l_air_h=l_air, coil_length_m=geo.coil_length_m, slug_length_m=slug_len,
        mu_eff=mu_eff, smoothing_width_m=wire_od, total_turns=geo.total_turns, b_sat_tesla=b_sat,
    )
    cap_f = cap_uf * 1e-6
    cp = StageCircuitParams(capacitance_f=cap_f, r_total_ohm=r_wire + esr, mass_kg=mass)
    q0 = cap_f * charge_v
    return model, cp, q0, -0.02, 5.0  # x_fire, v_fire


def _scipy_discharge(model, cp, q0, x_fire, v_fire):
    sol = solve_ivp(
        circuit.derivatives, (0.0, 0.03), [q0, 0.0, x_fire, v_fire],
        args=(model, cp), events=circuit.zero_current_crossing_event, rtol=1e-9, atol=1e-12,
    )
    if len(sol.t_events[0]) == 0:
        return None
    q_f, i_f, x_f, v_f = sol.y_events[0][0]
    energy_diss = float(np.trapezoid(sol.y[1] ** 2 * cp.r_total_ohm, sol.t))
    return {"exit_v": v_f, "exit_x": x_f, "peak_current": float(np.max(np.abs(sol.y[1]))), "energy_diss": energy_diss}


CONFIGS = [
    dict(turns_per_layer=20, layers=4, tube_od=0.01, slug_d=0.008, slug_len=0.02, mu_r=200, b_sat=1.9, cap_uf=100, esr=0.05, charge_v=300, wire_rho=1.68e-8, wire_gauge=1.0),
    dict(turns_per_layer=30, layers=5, tube_od=0.012, slug_d=0.006, slug_len=0.018, mu_r=500, b_sat=1.8, cap_uf=220, esr=0.08, charge_v=400, wire_rho=1.68e-8, wire_gauge=0.8),
    dict(turns_per_layer=15, layers=3, tube_od=0.009, slug_d=0.005, slug_len=0.015, mu_r=300, b_sat=2.0, cap_uf=68, esr=0.1, charge_v=250, wire_rho=2.82e-8, wire_gauge=1.4),
    dict(turns_per_layer=25, layers=6, tube_od=0.011, slug_d=0.007, slug_len=0.025, mu_r=400, b_sat=1.9, cap_uf=150, esr=0.06, charge_v=350, wire_rho=1.68e-8, wire_gauge=0.71),
]


def test_batch_discharge_matches_scipy_reference():
    xp, _ = get_backend(prefer_gpu=False)  # тест на CPU/numpy
    models, cps, q0s, x0s, v0s = [], [], [], [], []
    refs = []
    for c in CONFIGS:
        model, cp, q0, x0, v0 = _setup(**c)
        ref = _scipy_discharge(model, cp, q0, x0, v0)
        assert ref is not None, "scipy-эталон должен скоммутироваться"
        refs.append(ref)
        models.append(model); cps.append(cp); q0s.append(q0); x0s.append(x0); v0s.append(v0)

    params = params_from_models(xp, models, cps)
    out = integrate_batch_discharge(
        xp, xp.asarray(q0s), xp.asarray(x0s), xp.asarray(v0s), params, dt=1e-6, max_steps=40000
    )

    assert bool(xp.all(out["feasible"])), "батч должен скоммутировать все конфигурации"
    for k, ref in enumerate(refs):
        # фикс. шаг RK4 против адаптивного solve_ivp — совпадение в пределах ~1%
        assert out["exit_v"][k] == np.float64(ref["exit_v"]) or abs(out["exit_v"][k] - ref["exit_v"]) <= 0.02 * abs(ref["exit_v"]) + 0.05, \
            f"конфиг {k}: exit_v батч {out['exit_v'][k]:.3f} vs scipy {ref['exit_v']:.3f}"
        assert abs(out["peak_current"][k] - ref["peak_current"]) <= 0.02 * ref["peak_current"] + 0.5, \
            f"конфиг {k}: пик тока батч {out['peak_current'][k]:.1f} vs scipy {ref['peak_current']:.1f}"
        assert abs(out["energy_dissipated_j"][k] - ref["energy_diss"]) <= 0.03 * ref["energy_diss"] + 0.01, \
            f"конфиг {k}: потери батч {out['energy_dissipated_j'][k]:.3f} vs scipy {ref['energy_diss']:.3f}"


def test_batch_energy_conservation_lossless():
    """С R=0 энергия конденсатора = кинетика + магнитное поле (в конце I=0 -> вся в КЭ+остаток)."""
    xp, _ = get_backend(prefer_gpu=False)
    c = dict(turns_per_layer=20, layers=4, tube_od=0.01, slug_d=0.008, slug_len=0.02, mu_r=200,
             b_sat=1.9, cap_uf=100, esr=0.0, charge_v=300, wire_rho=0.0, wire_gauge=1.0)
    model, cp, q0, x0, v0 = _setup(**c)
    params = params_from_models(xp, [model], [cp])
    out = integrate_batch_discharge(xp, xp.asarray([q0]), xp.asarray([x0]), xp.asarray([v0]), params, dt=1e-6, max_steps=40000)
    assert bool(out["feasible"][0])
    energy_in = 0.5 * cp.capacitance_f * (q0 / cp.capacitance_f) ** 2
    ke_before = 0.5 * cp.mass_kg * v0**2
    ke_after = 0.5 * cp.mass_kg * float(out["exit_v"][0]) ** 2
    remaining = float(out["exit_q"][0]) ** 2 / (2 * cp.capacitance_f)
    balance = remaining + (ke_after - ke_before)
    assert abs(balance - energy_in) <= 0.01 * energy_in, f"баланс {balance:.3f} vs вход {energy_in:.3f}"