from __future__ import annotations

from dataclasses import dataclass
from typing import Any, Callable, Mapping, Protocol, Tuple

import numpy as np


class MutatorFn(Protocol):
    """Callable that creates child params from parent params.

    The mutator receives:
    * parent params (dict)
    * RNG
    * arbitrary keyword arguments (context)

    It returns:
    * child params (dict)
    * metadata/info dict (may be empty)
    """

    def __call__(
        self,
        parent_params: Mapping[str, Any],
        rng: np.random.Generator,
        **kwargs: Any,
    ) -> Tuple[dict[str, Any], dict[str, Any]]:
        raise NotImplementedError


class Performant(Protocol):
    """Strategy object for `ZenFronts.perform_new` semantics."""

    def perform_new(
        self,
        parent_id: int,
        *,
        get_params: Callable[[int], dict[str, Any]],
        add_point: Callable[[Mapping[str, Any]], int],
        delete_point: Callable[[int], None],
        rng: np.random.Generator,
        mutator: MutatorFn,
        looser_id: int | None = None,
        remove_looser: bool = True,
        **kwargs: Any,
    ) -> tuple[int, dict[str, Any]]:
        raise NotImplementedError


@dataclass(frozen=True)
class DefaultPerformant:
    """Default `perform_new` semantics.

    * Calls mutator(parent_params, rng, **kwargs)
    * Adds new point
    * Optionally removes the provided looser
    """

    def perform_new(
        self,
        parent_id: int,
        *,
        get_params: Callable[[int], dict[str, Any]],
        add_point: Callable[[Mapping[str, Any]], int],
        delete_point: Callable[[int], None],
        rng: np.random.Generator,
        mutator: MutatorFn,
        looser_id: int | None = None,
        remove_looser: bool = True,
        **kwargs: Any,
    ) -> tuple[int, dict[str, Any]]:
        pid = int(parent_id)
        parent_params = get_params(pid)

        child_params, mut_info = mutator(parent_params, rng, **kwargs)
        new_id = int(add_point(child_params))

        info: dict[str, Any] = {
            "parent_id": pid,
            "new_id": new_id,
            "looser_id": int(looser_id) if looser_id is not None else None,
            "remove_looser": bool(remove_looser),
        }
        if isinstance(mut_info, dict):
            info.update(mut_info)

        if looser_id is not None and bool(remove_looser):
            delete_point(int(looser_id))

        return new_id, info
