from typing import Dict, TypeVar

from ray.data._internal.execution.interfaces import ExecutionResources

# The math functions defined in this module use a generic type rather than
# `PhysicalOperator` so it's easier to test. We already pass in all of the necessary
# inputs, so the actual type doesn't matter.
T = TypeVar("T")


_SCHEDULABLE_RESOURCE_NAMES = ("cpu", "gpu", "memory")


def allocate_resources(
    throughput: float,
    *,
    rates: Dict[T, float],
    resource_requirements: Dict[T, ExecutionResources],
) -> Dict[T, ExecutionResources]:
    """Allocate resources for a pipeline to sustain the given throughput.

    Key insight: in a pipeline, all operators must sustain the same throughput T.
    Operator i with per-task rate r_i needs T/r_i tasks to sustain T. So maximizing
    throughput is equivalent to finding the largest feasible T, then deriving task
    counts from it.

    Args:
        throughput: The throughput for the pipeline in the same units as the rates.
        rates: The rate at which a task or actor produces outputs for each operator.
        resource_requirements: The logical resources required to schedule a task or
            actor for each operator.

    Returns:
        A dictionary mapping operators to the allocated resources.
    """
    assert throughput >= 0, "Throughput must be non-negative"
    assert all(rate > 0 for rate in rates.values()), "Rates must be positive"

    if not rates:
        return {}

    if throughput == 0:
        return {op: ExecutionResources.zero() for op in rates}

    # NOTE: This implementation computes fractional task counts. In practice, you
    # can't schedule a fractional task or actor, so the allocations might be infeasible.
    task_counts = {op: throughput / rate for op, rate in rates.items()}
    return {op: resource_requirements[op].scale(task_counts[op]) for op in rates}


def compute_optimal_throughput(
    *,
    rates: Dict[T, float],
    resource_requirements: Dict[T, ExecutionResources],
    resource_limits: ExecutionResources,
    concurrency_limits: Dict[T, int | None],
) -> float:
    """Compute the optimal throughput for a pipeline.

    The optimal throughput is bounded by two constraints (we take the tightest):
      1. Resource limits — total resource usage across all operators must fit the
         budget.
      2. Concurrency limits — each operator's task count cannot exceed its limit.

    Args:
        rates: The rate at which a task or actor produces outputs for each operator.
        resource_requirements: The logical resources required to schedule a task or
            actor for each operator.
        resource_limits: The resource limits for the cluster.
        concurrency_limits: The maximum number of tasks or actors that can be scheduled
            concurrently for each operator.

    Returns:
        The optimal throughput for the pipeline in the same units as the rates.
    """
    assert rates, "Rates must be non-empty"
    return min(
        _max_throughput_from_resources(rates, resource_requirements, resource_limits),
        _max_throughput_from_concurrency(rates, concurrency_limits),
    )


def _max_throughput_from_resources(
    rates: Dict[T, float],
    resource_requirements: Dict[T, ExecutionResources],
    resource_limits: ExecutionResources,
) -> float:
    """For each resource type, compute the max throughput the resource budget allows."""
    assert rates, "Rates must be non-empty"
    assert all(rate > 0 for rate in rates.values()), "Rates must be positive"
    assert (
        rates.keys() <= resource_requirements.keys()
    ), "You must provide a resource requirement for each operator with a rate."

    max_throughput = float("inf")

    for resource_name in _SCHEDULABLE_RESOURCE_NAMES:
        resource_limit = getattr(resource_limits, resource_name)
        resource_cost_per_unit_throughput = sum(
            getattr(resource_requirements[op], resource_name) / rates[op]
            for op in rates
        )
        if resource_cost_per_unit_throughput > 0:
            max_throughput = min(
                max_throughput, resource_limit / resource_cost_per_unit_throughput
            )

    assert max_throughput >= 0, "Max throughput must be non-negative"
    return max_throughput


def _max_throughput_from_concurrency(
    rates: Dict[T, float],
    concurrency_limits: Dict[T, int | None],
) -> float:
    """Each operator's throughput is capped at rate * concurrency_limit."""
    assert rates, "Rates must be non-empty"
    assert (
        rates.keys() <= concurrency_limits.keys()
    ), "You must provide a concurrency limit for each operator with a rate."

    # Convert `None` to float("inf") for operators with no concurrency limit
    normalized_concurrency_limits: Dict[T, float] = {
        op: limit if limit is not None else float("inf")
        for op, limit in concurrency_limits.items()
    }
    return min(rates[op] * normalized_concurrency_limits[op] for op in rates)