"""Utilities specific to archives."""

from __future__ import annotations

from types import ModuleType

import numpy as np
from numpy.typing import DTypeLike

from ribs._utils import xp_namespace
from ribs.typing import BatchData, Float


def parse_dtype(dtype: DTypeLike, xp: ModuleType) -> DTypeLike:
    """Makes any necessary modifications to the input dtype.

    Any dtypes that are `None` are set to the default "real floating" dtype for the
    provided array backend.
    """
    if dtype is None:
        xp = xp_namespace(xp)
        # See here for info on array API inspection:
        # https://data-apis.org/array-api/latest/API_specification/inspection.html
        return xp.__array_namespace_info__().default_dtypes()["real floating"]  # ty: ignore[unresolved-attribute]
    return dtype


def parse_all_dtypes(
    dtype: DTypeLike,
    solution_dtype: DTypeLike,
    objective_dtype: DTypeLike,
    measures_dtype: DTypeLike,
    xp: ModuleType,
) -> DTypeLike:
    """Parses dtypes for the archives.

    Any dtypes that are `None` are set to the default "real floating" dtype for the
    provided array backend.

    Returns dtypes for the solution, objective, and measures.
    """
    if isinstance(dtype, dict):
        raise ValueError(
            "Passing a dict as `dtype` is deprecated in pyribs 0.9.0. Please use "
            "`solution_dtype`, `objective_dtype`, and/or `measures_dtype` instead."
        )

    use_dtype = dtype is not None
    use_individual_dtypes = (
        solution_dtype is not None
        or objective_dtype is not None
        or measures_dtype is not None
    )

    if use_dtype and use_individual_dtypes:
        raise ValueError(
            "dtype cannot be used at the same time as solution_dtype, "
            "objective_dtype, or measures_dtype. Either pass dtype on its own,"
            "or pass solution_dtype, objective_dtype, and/or measures_dtype."
        )
    elif use_individual_dtypes:
        # Parse each dtype individually.
        solution_dtype = parse_dtype(solution_dtype, xp)
        objective_dtype = parse_dtype(objective_dtype, xp)
        measures_dtype = parse_dtype(measures_dtype, xp)
    else:
        # Everything set to `dtype`, which may also be None to get the default case.
        dtype = parse_dtype(dtype, xp)
        solution_dtype = dtype
        objective_dtype = dtype
        measures_dtype = dtype

    return solution_dtype, objective_dtype, measures_dtype


def validate_cma_mae_settings(
    learning_rate: Float | None, threshold_min: Float, dtype: np.dtype
) -> tuple[np.floating, np.floating]:
    """Checks variables related to CMA-MAE, i.e., learning_rate and threshold_min."""
    if threshold_min != -np.inf and learning_rate is None:
        raise ValueError(
            "threshold_min was set without setting learning_rate. "
            "Please note that threshold_min is only used in CMA-MAE; "
            "it is not intended to be used for only filtering archive "
            "solutions. To run CMA-MAE, please also set learning_rate."
        )
    if learning_rate is None:
        learning_rate = 1.0  # Default value.
    if threshold_min == -np.inf and learning_rate != 1.0:
        raise ValueError("threshold_min can only be -np.inf if learning_rate is 1.0")
    learning_rate = np.asarray(learning_rate, dtype=dtype)
    threshold_min = np.asarray(threshold_min, dtype=dtype)
    return learning_rate, threshold_min


def fill_sentinel_values(occupied: np.ndarray, data: BatchData) -> None:
    """Fills unoccupied entries in data with sentinel values.

    Operates in-place on `data`.
    """
    unoccupied = ~occupied

    for name, arr in data.items():
        if arr.dtype == object:
            fill_val = None
        elif name == "index":
            fill_val = -1
        elif np.issubdtype(arr.dtype, np.integer):
            fill_val = 0
        else:  # Floating-point and other fields.
            fill_val = np.nan
        arr[unoccupied] = fill_val
