"""Parameter scheduler dispatcher.

These dispatchers are meant to extend the standard `get_fit_config`
and `get_eval_config` that the Flower Framework's Server uses to
compile the task instructions for the client. As such, the functions
dispatched should retain the same typing of the `get_fit_config` or
`get_eval_config`.
"""

from collections.abc import Callable
from typing import cast

from omegaconf import DictConfig

from mt_dao.conf.base_schema import BaseConfig, ParamSchedulerName
from mt_dao.constants import ALL
from mt_dao.utils import ModelStateNames


def get_freq_param_scheduler(
    kwargs: DictConfig | dict[str, int],
) -> Callable[[str | int, int], list[ModelStateNames]]:
    """Get a frequency parameter scheduler.

    Parameters
    ----------
    kwargs : DictConfig | dict[str, int]
        The keyword arguments, should be
            a dictionary of parameter names and frequencies.

    Returns
    -------
    Callable[[str | int, int], list[ModelStateNames]]
        The frequency parameter scheduler.

    """

    def freq_scheduler(
        cid: str | int,  # noqa: ARG001
        server_round: int,
    ) -> list[ModelStateNames]:
        return [
            ModelStateNames[cast("str", name).upper()]
            for name, freq in kwargs.items()
            if server_round % freq == 0
        ]

    return freq_scheduler


def get_layer_filter_scheduler(
    kwargs: DictConfig | dict[str, dict[str, int]],
) -> Callable[[str | int, int], dict[ModelStateNames, list[str]]]:
    """Get a parameter scheduler filtered by layer type.

    We assume that we receive a dictionary containing parameter names that we
    wish to preserve as the keys, with the frequency at which they should be
    transmitted as the value of the corresponding layer name. The layer name matching
    is done "loosely" in the sense that we look for whether the string specified is
    contained in the layer name. For example, specifiying a key of '10' preserves only
    those layers containing `10` in their name (i.e the 10th layer). A special case if
    considered if we specify the layer name 'all' where we wish to send all layers.

    Parameters
    ----------
    kwargs : DictConfig | dict[str, int]
        The keyword arguments, should be
            a dictionary of parameter names and frequencies.

    Returns
    -------
    Callable[[str | int, int], list[ModelStateNames]]
        The frequency parameter scheduler.

    """

    def layer_filter_scheduler(
        cid: str | int,  # noqa: ARG001
        server_round: int,
    ) -> dict[ModelStateNames, list[str]]:
        returned_model_state_names: dict[ModelStateNames, list[str]] = {}
        for name, layer_dictionary in kwargs.items():
            state_class = ModelStateNames[cast("str", name).upper()]
            # NOTE: Handling when we want to transmit all layers
            if ALL in layer_dictionary:  # noqa: SIM102
                if server_round % layer_dictionary[ALL] == 0:
                    returned_model_state_names[state_class] = []
                    continue
            local_selected_names = []
            for layer_name, freq in layer_dictionary.items():
                if server_round % freq == 0:
                    local_selected_names.append(layer_name)
            returned_model_state_names[state_class] = local_selected_names
        return returned_model_state_names

    return layer_filter_scheduler


def dispatch_model_state_scheduler(
    cfg: BaseConfig,
) -> Callable[
    [str | int, int],
    list[ModelStateNames] | dict[ModelStateNames, list[str]],
]:
    """Dispatch the model state scheduler.

    Parameters
    ----------
    cfg : BaseConfig
        The configuration.

    Returns
    -------
    Callable[[str | int, int], list[ModelStateNames]]
        The model state scheduler.

    Raises
    ------
    ValueError
        If the scheduler contains unrecognized parameter values.

    """
    match cfg.fl.parameter_scheduler_name.lower():
        case ParamSchedulerName.ALL:
            return lambda _cid, _server_round: [ModelStateNames.ALL]
        case ParamSchedulerName.FREQ:
            return get_freq_param_scheduler(
                cfg.fl.parameter_scheduler_kwargs,
            )
        case ParamSchedulerName.LAYER_FILTER:
            return get_layer_filter_scheduler(cfg.fl.parameter_scheduler_kwargs)
        case _:
            msg = "Invalid parameter scheduler name."
            raise ValueError(msg)
