"""Handle aggregation in-place and potentially async."""

from collections import defaultdict
from concurrent.futures import Future, ThreadPoolExecutor
from queue import Queue

import numpy as np

from mt_dao.utils import ModelStateNames


def _aggregate_helper(
    prev_tuple: tuple[np.ndarray, int],
    param: np.ndarray,
    num_examples: int,
) -> tuple[np.ndarray, int]:
    """Compute weighted average for one parameter.

    Parameters
    ----------
    prev_tuple : Tuple[np.ndarray, int]
        Previously accumulated parameter and count.
    param : np.ndarray
        Current parameter.
    num_examples : int
        Number of current examples.

    Returns
    -------
    Tuple[np.ndarray, int]
        Updated parameter and total sample count.

    """
    prev_params, prev_total_examples = prev_tuple
    new_total_samples = prev_total_examples + num_examples
    acc_scaling_factor = float(prev_total_examples) / new_total_samples
    scaling_factor = float(num_examples) / new_total_samples
    updated_params = prev_params * acc_scaling_factor + param * scaling_factor
    return updated_params, new_total_samples


def _remove_futures_done(
    accumulator: dict[tuple[str, ModelStateNames], tuple[np.ndarray, int]],
    futures: dict[tuple[str, ModelStateNames], Future[tuple[np.ndarray, int]]],
) -> None:
    """Remove futures that are done from the futures dictionary.

    Parameters
    ----------
    accumulator : dict[tuple[str, ModelStateNames], tuple[np.ndarray, int]]
        A dictionary to accumulate the results from the futures.
    futures : dict[tuple[str, ModelStateNames], Future[tuple[np.ndarray, int]]]
        Dictionary of futures to check and remove done futures from.

    """
    futures_to_remove: list[tuple[str, ModelStateNames]] = []
    for key, future in futures.items():
        if future.done():
            accumulator[key] = future.result()
            futures_to_remove.append(key)
    for key in futures_to_remove:
        futures.pop(key)


def wait_for_all_futures(
    accumulator: dict[tuple[str, ModelStateNames], tuple[np.ndarray, int]],
    futures: dict[tuple[str, ModelStateNames], Future[tuple[np.ndarray, int]]],
) -> None:
    """Wait for all futures to complete and collect their results.

    This function collects the results from the futures and updates the
    accumulator with the results from each future.

    Parameters
    ----------
    accumulator : dict[tuple[str, ModelStateNames], tuple[np.ndarray, int]]
        A dictionary to accumulate the results from the futures.
    futures : dict[tuple[str, ModelStateNames], Future[tuple[np.ndarray, int]]]
        A dictionary of futures to wait for and collect results from.

    """
    # Collect the remaining results from futures
    while futures:
        _remove_futures_done(accumulator, futures)


def submit_aggregation_tasks_to_executor(
    executor: ThreadPoolExecutor,
    accumulator: dict[tuple[str, ModelStateNames], tuple[np.ndarray, int]],
    aggregation_queue: Queue[tuple[np.ndarray, int, str, ModelStateNames]],
    futures: dict[tuple[str, ModelStateNames], Future[tuple[np.ndarray, int]]],
) -> dict[tuple[str, ModelStateNames], tuple[np.ndarray, int]]:
    """Submit aggregation tasks to the executor.

    The function processes the aggregation queue and submits tasks to the executor.
    It checks if the key is already in the accumulator and if a future is already
    running for that key. If so, it puts the item back in the queue.
    If the key is not in the accumulator, it adds it. If a future is running for
    that key, it submits a task to the executor to aggregate the parameters.

    Parameters
    ----------
    executor : ThreadPoolExecutor
        The executor to submit tasks to.
    accumulator : dict[tuple[str, ModelStateNames], tuple[np.ndarray, int]]
        The accumulator containing parameters and their corresponding number of
        examples.
    aggregation_queue : Queue[tuple[np.ndarray, int, str, ModelStateNames]]
        The queue containing parameters to be aggregated.
    futures : dict[tuple[str, ModelStateNames], Future[tuple[np.ndarray, int]]]
        The futures mapping keys to their corresponding Future objects.

    Returns
    -------
    dict[tuple[str, ModelStateNames], tuple[np.ndarray, int]]
        Updated accumulator with aggregated parameters and their total number of
        examples.

    """
    # Remove futures that are done to ensure they won't impact the first iteration over
    # the aggregation queue
    _remove_futures_done(accumulator, futures)
    # Process the queue with the executor
    while not aggregation_queue.empty():
        # Get the next parameter from the queue
        param, num_examples, name, typ = aggregation_queue.get()
        # Extract the key for the accumulator
        key = (name, typ)
        # Check if there's any future we are waiting for using this key
        is_busy = (name, typ) in futures
        # If the key is busy, put it back in the queue
        if is_busy:
            aggregation_queue.put((param, num_examples, name, typ))
            continue
        # If the key is already in the accumulator, submit a task to the executor
        if key in accumulator:
            futures[key] = executor.submit(
                _aggregate_helper,
                accumulator[key],
                param,
                num_examples,
            )
        else:
            accumulator[key] = (param, num_examples)
        # Remove futures that are done to ensure we can eventually empty the queue and
        # exit the while loop
        _remove_futures_done(accumulator, futures)
    return accumulator


def weighted_average(
    metrics: list[tuple[int, dict]],
) -> dict:
    """Compute a weighted average over pre-defined metrics.

    Parameters
    ----------
    metrics : List[Tuple[int, Dict]]
        The metrics to aggregate.

    Returns
    -------
    Dict
        The weighted average over pre-defined metrics.

    """
    # Filter out empty metrics
    metrics = [
        (num_examples, metric) for num_examples, metric in metrics if metric != {}
    ]
    total_num_examples = sum(num_examples for num_examples, _ in metrics)
    weighted_metrics: dict = defaultdict(float)

    for num_examples, metric in metrics:
        # NOTE: These should now all be metrics and not contain any weird stuff
        for key, value in metric.items():
            if not isinstance(value, str):
                weighted_metrics[key] += num_examples * value

    return {
        key: value / total_num_examples for key, value in weighted_metrics.items()
    }
