"""Dispatch strategy based on configuration."""

from collections.abc import Callable

import numpy as np

from mt_dao.conf.base_schema import BaseConfig, StrategyName
from mt_dao.constants import FIRST_MOMENTUM, MODEL_PARAMETERS, SECOND_MOMENTUM
from mt_dao.masks_utils import ModelStateNames
from mt_dao.server.server_state import ServerState
from mt_dao.strategy_functional.fedadam import federated_adam
from mt_dao.strategy_functional.fedavg_eff import federated_averaging_efficient
from mt_dao.strategy_functional.fedmom import federated_averaging_with_momentum
from mt_dao.strategy_functional.fednesterov import (
    federated_averaging_with_nesterov_momentum,
)
from mt_dao.strategy_functional.fedyogi import federated_yogi

ServerOptimizer = Callable[
    [
        ServerState,
        dict[tuple[str, ModelStateNames], np.ndarray],
    ],
    None,
]


class StrategyNotImplementedError(Exception):
    """Custom exception for unsupported strategies."""

    def __init__(self, strategy_name: str) -> None:
        """Initialize the exception with the strategy name.

        Parameters
        ----------
        strategy_name : str
            The name of the strategy that is not implemented or supported.

        """
        super().__init__(f"Strategy '{strategy_name}' is not implemented or supported.")
        self.strategy_name = strategy_name


def dispatch_strategy_state_keys(cfg: BaseConfig) -> tuple[str, ...]:
    """Dispatch the state keys based on the strategy configuration.

    Parameters
    ----------
    cfg : BaseConfig
        The configuration object containing the settings for the server.

    Returns
    -------
    tuple[str, ...]
        The state keys to use for the specified strategy.

    Raises
    ------
    StrategyNotImplementedError
        If the strategy is not implemented or supported.

    """
    # NOTE: We need to lowercase the match as the Enum instantiate with auto() in Python
    # lowercases by default
    match cfg.fl.strategy_name.lower():
        case StrategyName.NESTEROV:
            return (
                MODEL_PARAMETERS,
                FIRST_MOMENTUM,
            )
        case StrategyName.FEDMOM:
            return (
                MODEL_PARAMETERS,
                FIRST_MOMENTUM,
            )
        case StrategyName.FEDAVG:
            return (MODEL_PARAMETERS,)
        case StrategyName.FEDYOGI:
            return (
                MODEL_PARAMETERS,
                FIRST_MOMENTUM,
                SECOND_MOMENTUM,
            )
        case StrategyName.FEDADAM:
            return (
                MODEL_PARAMETERS,
                FIRST_MOMENTUM,
                SECOND_MOMENTUM,
            )
        case _:
            msg = "Unknown strategy"
            raise StrategyNotImplementedError(msg)


def dispatch_strategy_functional(
    cfg: BaseConfig,
) -> ServerOptimizer:
    """Dispatch the appropriate server optimizer based on the configuration.

    Parameters
    ----------
    cfg : BaseConfig
        The configuration object containing the settings for the server.

    Returns
    -------
    ServerOptimizer
        The server optimizer function to use for federated learning.

    Raises
    ------
    StrategyNotImplementedError
        If the strategy is not implemented.

    """
    # NOTE: We need to lowercase the match as the Enum instantiate with auto() in Python
    # lowercases by default
    match cfg.fl.strategy_name.lower():
        case StrategyName.NESTEROV:
            assert (
                cfg.fl.strategy_kwargs.server_learning_rate is not None
            ), "Server learning rate is required for Nesterov strategy."
            assert (
                cfg.fl.strategy_kwargs.server_momentum is not None
            ), "Server momentum is required for Nesterov strategy."
            return lambda server_state, fedavg_result: (
                federated_averaging_with_nesterov_momentum(
                    server_state=server_state,
                    fedavg_result=fedavg_result,
                    server_learning_rate=cfg.fl.strategy_kwargs.server_learning_rate,
                    server_momentum=cfg.fl.strategy_kwargs.server_momentum,
                )
            )
        case StrategyName.FEDMOM:
            assert (
                cfg.fl.strategy_kwargs.server_learning_rate is not None
            ), "Server learning rate is required for FedMom strategy."
            assert (
                cfg.fl.strategy_kwargs.server_momentum is not None
            ), "Server momentum is required for FedMom strategy."
            return lambda server_state, fedavg_result: (
                federated_averaging_with_momentum(
                    server_state=server_state,
                    fedavg_result=fedavg_result,
                    server_learning_rate=cfg.fl.strategy_kwargs.server_learning_rate,
                    server_momentum=cfg.fl.strategy_kwargs.server_momentum,
                )
            )
        case StrategyName.FEDAVG:
            return lambda server_state, fedavg_result: (
                federated_averaging_efficient(
                    server_state=server_state,
                    fedavg_result=fedavg_result,
                    server_learning_rate=1.0,
                )
            )
        case StrategyName.FEDYOGI:
            assert (
                cfg.fl.strategy_kwargs.eta is not None
            ), "Server learning rate is required for FedYogi strategy."
            assert (
                cfg.fl.strategy_kwargs.tau is not None
            ), "Server momentum is required for FedYogi strategy."
            assert (
                cfg.fl.strategy_kwargs.beta_1 is not None
            ), "Beta 1 is required for FedYogi strategy."
            assert (
                cfg.fl.strategy_kwargs.beta_2 is not None
            ), "Beta 2 is required for FedYogi strategy."
            return lambda server_state, fedavg_result: (
                federated_yogi(
                    server_state=server_state,
                    fedavg_result=fedavg_result,
                    beta_1=cfg.fl.strategy_kwargs.beta_1,
                    beta_2=cfg.fl.strategy_kwargs.beta_2,
                    eta=cfg.fl.strategy_kwargs.eta,
                    tau=cfg.fl.strategy_kwargs.tau,
                )
            )
        case StrategyName.FEDADAM:
            assert (
                cfg.fl.strategy_kwargs.eta is not None
            ), "Server learning rate is required for FedAdam strategy."
            assert (
                cfg.fl.strategy_kwargs.tau is not None
            ), "Server momentum is required for FedAdam strategy."
            assert (
                cfg.fl.strategy_kwargs.beta_1 is not None
            ), "Beta 1 is required for FedAdam strategy."
            assert (
                cfg.fl.strategy_kwargs.beta_2 is not None
            ), "Beta 2 is required for FedAdam strategy."
            return lambda server_state, fedavg_result: (
                federated_adam(
                    server_state=server_state,
                    fedavg_result=fedavg_result,
                    beta_1=cfg.fl.strategy_kwargs.beta_1,
                    beta_2=cfg.fl.strategy_kwargs.beta_2,
                    eta=cfg.fl.strategy_kwargs.eta,
                    tau=cfg.fl.strategy_kwargs.tau,
                )
            )
        case _:
            msg = "Unknown strategy"
            raise StrategyNotImplementedError(msg)
