import json

from gausse.components.database import ComponentDatabase
from gausse.gpu.batch_sweep import run_gpu_sweep
from gausse.optim.search_space import SearchBounds, decode, genome_from_dict
from gausse.sim.coilgun import run_coilgun
from gausse.storage.database import count_runs, fetch_runs, open_connection

DB = ComponentDatabase.load()


def test_gpu_sweep_records_all_runs(tmp_path):
    db_path = tmp_path / "gpu.sqlite3"
    summary = run_gpu_sweep(db_path, n_runs=300, seed=3, prefer_gpu=False, batch_size=300)
    assert summary["n_runs"] == 300
    conn = open_connection(db_path)
    assert count_runs(conn) == 300
    rows = fetch_runs(conn)
    # многоступенчатые конфиги допустимы (до max_stages)
    stage_counts = set()
    for r in rows:
        g = genome_from_dict(json.loads(r.genome_json))
        stage_counts.add(len(g.stages))
        assert r.search_mode.startswith("gpu-sweep")
    assert max(stage_counts) >= 2, "должны встречаться многоступенчатые конфиги"


def test_gpu_sweep_matches_cpu_within_tolerance(tmp_path):
    """Честная сверка: GPU-батч (в т.ч. многоступ.) даёт те же exit_v, что CPU run_coilgun."""
    db_path = tmp_path / "gpu.sqlite3"
    run_gpu_sweep(db_path, n_runs=2000, seed=5, prefer_gpu=False, batch_size=2000)
    conn = open_connection(db_path)
    bounds = SearchBounds()
    feasible = fetch_runs(conn, feasible=True, order_by_efficiency_desc=True, limit=12)
    assert len(feasible) >= 3
    checked = 0
    for r in feasible:
        g = genome_from_dict(json.loads(r.genome_json))
        config, _, _ = decode(g, DB, bounds)
        cpu = run_coilgun(config)
        if not cpu.feasible:
            continue
        # фикс.шаг батча + аналитический триггер vs адаптивный solve_ivp: ~3%
        assert abs(r.exit_velocity_mps - cpu.exit_v_mps) <= 0.03 * abs(cpu.exit_v_mps) + 0.2, \
            f"GPU {r.exit_velocity_mps:.2f} vs CPU {cpu.exit_v_mps:.2f} ({len(g.stages)} ступ.)"
        checked += 1
    assert checked >= 3