from __future__ import annotations

from typing import Any

import numpy as np

from zen_fronts import ZenFronts
from zen_fronts.selection import SelectionConfig


def test_happy_path_refresh_choose_parent_perform_new_removes_loser_and_updates_stats() -> None:
    # Minimal sampler/mutator for this sprint: single numeric parameter.
    def sampler(rng: np.random.Generator) -> dict[str, Any]:
        return {"p": float(rng.uniform(-1.0, 1.0))}

    seen_kwargs: dict[str, Any] = {}

    def mutator(
        parent: dict[str, Any], rng: np.random.Generator, **kwargs: Any
    ) -> tuple[dict[str, Any], dict[str, Any]]:
        # Capture kwargs to ensure semantics.
        seen_kwargs.clear()
        seen_kwargs.update(kwargs)

        # Small gaussian mutation.
        child = dict(parent)
        child["p"] = float(child["p"] + rng.normal(scale=0.1))
        return child, {"mutated": True}

    z = ZenFronts(
        crits={"cr1": "max", "cr2": "max", "cr3": "min"},
        selection=SelectionConfig(n_samples=200, percentile=0.2, seed=123),
        sampler=sampler,
        mutator=mutator,
    )

    ids = z.add_random_point(8)
    assert len(z.points) == 8

    # Make one point clearly bad.
    loser_id = ids[0]
    for pid in ids:
        if pid == loser_id:
            z.update_crits(pid, {"cr1": -100.0, "cr2": -100.0, "cr3": 100.0}, t=1.0)
        else:
            z.update_crits(pid, {"cr1": 10.0 + pid, "cr2": 5.0 + pid, "cr3": -1.0}, t=1.0)

    losers = z.refresh(now=1.0)
    assert loser_id in losers

    # After refresh(), individuals should have selection stats stored.
    info_loser = z.info(loser_id)
    assert info_loser["selection"] is not None
    assert "place" in info_loser["selection"]
    assert "front" in info_loser["selection"]
    assert "within" in info_loser["selection"]

    parent_id = z.choose_parent(loser_id, seed=999)
    assert parent_id in z.points
    assert parent_id != loser_id

    new_id, info = z.perform_new(parent_id, looser=loser_id, remove_looser=True, seed=999, step=5)
    assert new_id in z.points
    assert info["parent_id"] == parent_id
    assert info["looser_id"] == loser_id
    assert info["mutated"] is True

    # Ensure mutator received contextual kwargs.
    assert seen_kwargs["zf_parent_id"] == parent_id
    assert seen_kwargs["zf_loser_id"] == loser_id
    assert seen_kwargs["step"] == 5

    # The loser is removed from params storage.
    assert loser_id not in z.points
    try:
        z.params(loser_id)
        assert False, "expected KeyError"
    except KeyError:
        pass
