from arena.core.runner import run_arena
from arena.core.types import RunLimits
from arena.optimizers.registry import get_example


def test_run_arena_gradient_descent_linear_regression():
    example = get_example("Gradient Descent")
    result = run_arena(
        task_name="Linear Regression",
        optimizer_code=example.code,
        optimizer_config=example.default_config,
        task_config={"n": 64, "d": 4},
        device_name="cpu",
        steps=20,
        seed=0,
        limits=RunLimits(max_steps=100, max_seconds=10.0, log_every=1),
    )
    assert result.status == "success"
    assert not result.history.empty
    assert result.summary["final_loss"] <= result.summary["initial_loss"]


def test_iter_arena_streams_partial_results():
    from arena.core.runner import iter_arena

    example = get_example("Gradient Descent")
    results = list(
        iter_arena(
            task_name="Quadratic Bowl",
            optimizer_code=example.code,
            optimizer_config=example.default_config,
            task_config={"dim": 8},
            device_name="cpu",
            steps=5,
            seed=0,
            limits=RunLimits(max_steps=10, max_seconds=10.0, log_every=1),
        )
    )
    assert len(results) >= 2
    assert results[-1].status == "success"
    assert results[-1].summary["status"] == "success"