"""Base configuration schema."""

import ast
import os
from enum import Enum, auto
from typing import Any

from hydra.core.config_store import ConfigStore
from omegaconf import MISSING, DictConfig
from pydantic.dataclasses import dataclass


@dataclass(config={"arbitrary_types_allowed": True})
class CommStack(DictConfig):
    """Communication stack configuration.

    Attributes
    ----------
    shm: bool = MISSING
        Whether to use SharedMemory communication
    s3: bool = MISSING
        Whether to use S3 communication
    ray: bool = MISSING
        Whether to use Ray communication
    n_batches: int | None = MISSING
        Number of batches to split the communication of parameters into

    """

    s3: bool = MISSING
    shm: bool = MISSING
    ray: bool = MISSING
    n_batches: int | None = MISSING


@dataclass(config={"arbitrary_types_allowed": True})
class Centralized(DictConfig):
    """Centralized configuration.

    Attributes
    ----------
    store_init_model: bool = MISSING
        Whether to store the initial model
    store_final_model: bool = MISSING
        Whether to store the final model
    stream_id: str | None = MISSING
        Stream id to pass to the data configuration
    eval_only: bool = MISSING
        Whether to only execute the evaluation
    split_eval: bool = MISSING
        Whether to report performance separately on each stream
    reset_timestamp: bool = MISSING
        Whether to reset the timestamp

    """

    store_init_model: bool = MISSING
    store_final_model: bool = MISSING
    stream_id: str | None = MISSING
    eval_only: bool = MISSING
    split_eval: bool = MISSING
    reset_timestamp: bool = MISSING


@dataclass(config={"arbitrary_types_allowed": True})
class Checkpointing(DictConfig):
    """Checkpoint configuration.

    Attributes
    ----------
    local_checkpoint: bool = MISSING
        Whether to checkpoint the model to local storage.
    remote_checkpoint: bool = MISSING
        Whether to checkpoint the model to a remote S3 bucket.
    local_cleanup_checkpoints: bool = MISSING
        Whether to clean up all the local checkpoints at the end
    local_cleanup_checkpoints_per_round: bool = MISSING
        Whether to clean up the local checkpoints at the end of each round
    remote_cleanup_checkpoints: bool = MISSING
        Whether to clean up all the remote checkpoints at the end
    remote_cleanup_checkpoints_per_round: bool = MISSING
        Whether to clean up the remote checkpoints at the end of each round
    per_round_cleanup_staleness: int = MISSING
        How many checkpoints to delete when removing checkpoints from local/remote.
        For example, if cleanup_staleness is 3, we will keep the 3 most recent
        checkpoints. The stateleness factor needs to be an integer greater than zero.
    copy_client_checkpoints: bool = MISSING
        Whether to copy the client checkpoints

    """

    local_checkpoint: bool = MISSING
    remote_checkpoint: bool = MISSING
    local_cleanup_checkpoints: bool = MISSING
    local_cleanup_checkpoints_per_round: bool = MISSING
    remote_cleanup_checkpoints: bool = MISSING
    remote_cleanup_checkpoints_per_round: bool = MISSING
    per_round_cleanup_staleness: int = MISSING
    copy_client_checkpoints: bool = MISSING


@dataclass(config={"arbitrary_types_allowed": True})
class CentralisedRestore(DictConfig):
    """Centralised run restoration configuration.

    Attributes
    ----------
    preserve_optimizer_states: bool = MISSING
        Whether to preserve the optimizer states from the
        centralised checkpoint
    checkpoint_dir: str = MISSING
        The directory from which to restore the centralized model
    run_batches: int = MISSING
        Number of batches to restore from the centralized model
    run_uuid: str = MISSING
        Run UUID to restore the centralized model

    """

    preserve_optimizer_states: bool = MISSING
    checkpoint_dir: str = MISSING
    run_batches: int | None = MISSING
    run_uuid: str = MISSING


@dataclass(config={"arbitrary_types_allowed": True})
class mt_dao(DictConfig):
    """mt_dao configuration.

    Attributes
    ----------
    n_nodes: int = MISSING
        Number of nodes in the cluster
    refresh_period: int | None = MISSING
        Refresh period for the client workers
    local_init: bool
        Whether to initialise a new run from a previous local state.
    remote_init: bool
        Whether to initialise a new run from a previous remote state.
    restore_run_uuid: str = MISSING
        Run UUID to restore the model
    resume_round: int | None = MISSING
        Round to resume from, None implies start anew
        negative indices are counted from the last round.
    is_production: bool = MISSING
        Whether to run in production mode
    comm_stack: CommStack = Missing
        The communication stack configuration
    checkpointing: Checkpointing = MISSING
        The checkpointing configuration
    restore_cent : CentralisedRestore = MISSING
        The configuration for restoring from centralised
        runs.

    """

    n_nodes: int = MISSING
    refresh_period: int | None = MISSING
    local_init: bool = MISSING
    remote_init: bool = MISSING
    restore_run_uuid: str | None = MISSING
    resume_round: int | None = MISSING
    comm_stack: CommStack = MISSING
    is_production: bool = MISSING
    checkpointing: Checkpointing = MISSING
    restore_cent: CentralisedRestore = MISSING


class StrategyName(str, Enum):  # noqa: UP042
    """Strategy type."""

    @staticmethod
    def _generate_next_value_(
        name: str,
        start: int,  # noqa: ARG004
        count: int,  # noqa: ARG004
        last_values: list[int],  # noqa: ARG004
    ) -> str:
        """Generate the next value.

        Replacement for StrEnum to support python 3.10

        Parameters
        ----------
        name: str
            Name of the strategy
        start: int
            Start value
        count: int
            Count value
        last_values: list[int]
            List of last values

        Returns
        -------
        str
            Lowercase name

        """
        return name.lower()

    NESTEROV = auto()
    FEDMOM = auto()
    FEDAVG = auto()
    FEDYOGI = auto()
    FEDADAM = auto()


class ParamSchedulerName(str, Enum):  # noqa: UP042
    """Param scheduler type."""

    @staticmethod
    def _generate_next_value_(
        name: str,
        start: int,  # noqa: ARG004
        count: int,  # noqa: ARG004
        last_values: list[int],  # noqa: ARG004
    ) -> str:
        """Generate the next value.

        Replacement for StrEnum to support python 3.10

        Parameters
        ----------
        name: str
            Name of the strategy
        start: int
            Start value
        count: int
            Count value
        last_values: list[int]
            List of last values

        Returns
        -------
        str
            Lowercase name

        """
        return name.lower()

    ALL = auto()
    FREQ = auto()
    LAYER_FILTER = auto()


class StrategyKWArgs(dict[str, Any], DictConfig):  # type: ignore[reportIncompatibleMethodOverride,misc]
    """StrategyKWArgs configuration."""


class ParamSchedulerKWArgs(dict[str, Any], DictConfig):  # type: ignore[reportIncompatibleMethodOverride,misc]
    """ParamSchedulerKWArgs configuration."""


class DeltaModes(str, Enum):  # noqa: UP042
    """Name of modes for L2 norm calculation of client and server deltas."""

    @staticmethod
    def _generate_next_value_(
        name: str,
        start: int,  # noqa: ARG004
        count: int,  # noqa: ARG004
        last_values: list[int],  # noqa: ARG004
    ) -> str:
        """Generate the next value.

        Replacement for StrEnum to support python 3.10

        Parameters
        ----------
        name: str
            Name of the strategy
        start: int
            Start value
        count: int
            Count value
        last_values: list[int]
            List of last values

        Returns
        -------
        str
            Lowercase name

        """
        return name.lower()

    TRANSMISSION = auto()
    AGGREGATION = auto()
    INTERSECTION = auto()
    ALL_PARAMS = auto()


class ExpertsAssignmentPolicyName(str, Enum):  # noqa: UP042
    """Defines the experts assignment policy name."""

    @staticmethod
    def _generate_next_value_(
        name: str,
        start: int,  # noqa: ARG004
        count: int,  # noqa: ARG004
        last_values: list[int],  # noqa: ARG004
    ) -> str:
        """Generate the next value.

        Replacement for StrEnum to support python 3.10

        Parameters
        ----------
        name: str
            Name of the experts assignment policy
        start: int
            Start value
        count: int
            Count value
        last_values: list[int]
            List of last values

        Returns
        -------
        str
            Lowercase name

        """
        return name.lower()

    RANDOM = auto()
    FIXED = auto()


@dataclass(config={"arbitrary_types_allowed": True})
class ExpertsAssignmentPolicy(DictConfig):
    """Defines the experts assignment policy.

    Parameters
    ----------
    experts_assignment_policy_name: ExpertsAssignmentPolicyName
        Name of the experts assignment policy
    shared_expert_ids: list[int]
        Number of shared experts
    experts_overlapping_factor: int
        Overlapping factor for experts

    """

    experts_assignment_policy_name: ExpertsAssignmentPolicyName = MISSING
    shared_expert_ids: list[int] = MISSING
    experts_overlapping_factor: int = MISSING


@dataclass(config={"arbitrary_types_allowed": True})
class FL(DictConfig):
    """Federated learning configuration.

    Attributes
    ----------
    n_total_clients: int = MISSING
        Total number of clients
    n_clients_per_round: int = MISSING
        Number of clients per round
    n_rounds: int = MISSING
        Number of rounds
    reset_checkpoint: bool = MISSING
        Whether to reset the checkpoint
    reset_optimizer: bool = MISSING
        Whether to reset the optimizer
    reset_dataset_state: bool = MISSING
        Whether to reset the dataset state
    resize_vocab: int | None = MISSING
        Resize the vocabulary
    n_local_steps: int = MISSING
        Number of local steps
    use_unigram_metrics: bool = MISSING
        Whether to use unigram metrics
    allow_unigram_metrics_failures: bool = MISSING
        Whether to allow missing frequency dictionaries
    random_layers: list[str] = MISSING
        List of random layers
    random_init_freq: int = MISSING
        Random initialization frequency, 0 means never
    truly_random_init: bool = MISSING
        Whether to truly random initialization
    personalized_layers: list[str] = MISSING
        List of personalized layers
    frozen_layers : list[str] | None
        Specify layers to freeze.
    unfrozen_layers : list[str] | None
        Specify layers to keep unfrozen and freeze everything else.
    eval_period: int | None = MISSING
        Federated evaluation period, None means no evaluation is executed
    split_eval: bool = MISSING
        Whether to report performance separately on each stream
    strategy_name: StrategyName = MISSING
        Strategy name
    strategy_kwargs: StrategyKWArgs = MISSING
        Strategy kwargs
    aggregation_num_workers_process_ndarrays: int | None = MISSING
        Number of workers to parallelise aggregation across the parameters
        of a single update
    aggregation_num_workers_across_clients: int | None = MISSING
        Number of workers to parallelize client updates during aggregation
    dropout_ratio: float = MISSING
        Ratio of clients dropping out just after sampling.
    dropout_function_name: str = MISSING
        Method name to use for client dropout.
    parameter_scheduler_name: ParamSchedulerName = MISSING
        Parameter scheduler name,
        used to dispatch the parameter scheduler which controls the
        transmission and aggregation
        masks
    parameter_scheduler_kwargs: ParamSchedulerKWArgs = MISSING
        Parameter scheduler kwargs, passed
        to the dispatch function
        to configure behavior
    delta_modes: list[DeltaModes] = MISSING
        The mode(s) that we wish to use to calculate the L2 norms between server
        and client representations.
    experts_overlapping_factor: int = MISSING
        The number of clients that will share any single expert.
        This is used to control the overlap of experts across clients.
        A value of 1 means no overlap, so that each client has its own set of experts.
        A value of 2 means that each expert is shared by two clients, and so on.

    """

    n_total_clients: int = MISSING
    n_clients_per_round: int = MISSING
    n_rounds: int = MISSING
    reset_checkpoint: bool = MISSING
    reset_optimizer: bool = MISSING
    reset_dataset_state: bool = MISSING
    reset_timestamp: bool = MISSING
    resize_vocab: int | None = MISSING
    n_local_steps: int = MISSING
    use_unigram_metrics: bool = MISSING
    allow_unigram_metrics_failures: bool = MISSING
    random_layers: list[str] = MISSING
    random_init_freq: int = MISSING
    truly_random_init: bool = MISSING
    personalized_layers: list[str] = MISSING
    frozen_layers: list[str] | None = MISSING
    unfrozen_layers: list[str] | None = MISSING
    delta_modes: list[DeltaModes] = MISSING

    eval_period: int | None = MISSING
    split_eval: bool = MISSING

    strategy_name: StrategyName = MISSING
    strategy_kwargs: StrategyKWArgs = MISSING

    parameter_scheduler_name: ParamSchedulerName = MISSING
    parameter_scheduler_kwargs: ParamSchedulerKWArgs = MISSING

    dropout_ratio: float = MISSING
    dropout_function_name: str = MISSING

    experts_assignment_policy: ExpertsAssignmentPolicy = MISSING

    def __post_init__(self) -> None:
        """Execute post init checks."""
        # Check the number of clients per round against experts' overlapping factor
        assert (
            self.n_clients_per_round
            >= self.experts_assignment_policy.experts_overlapping_factor
        ), (
            "Incompatible configuration: n_clients_per_round"
            f" ({self.n_clients_per_round}) must be greater than or equal to"
            " experts_overlapping_factor"
            f" ({self.experts_assignment_policy.experts_overlapping_factor})."
        )
        if (
            self.experts_assignment_policy.experts_assignment_policy_name
            == ExpertsAssignmentPolicyName.FIXED
        ):
            # Check if the environmental variable exists
            mt_dao_client_to_experts_ids = ast.literal_eval(
                os.getenv("mt_dao_CLIENT_TO_EXPERTS_IDS", "{}"),
            )
            assert mt_dao_client_to_experts_ids != {}, (
                "When using the experts assignment policy"
                f" {ExpertsAssignmentPolicyName.FIXED}, the environmental variable"
                " mt_dao_CLIENT_TO_EXPERTS_IDS must be set."
            )
            # Check if the keys are consistent with the clients
            env_var_keys_set = set(mt_dao_client_to_experts_ids.keys())
            client_ids = set(range(self.n_clients_per_round))
            assert env_var_keys_set == client_ids, (
                f"The keys of mt_dao_CLIENT_TO_EXPERTS_IDS ({env_var_keys_set}) must"
                f" match the available client IDs ({client_ids})."
            )


@dataclass(config={"arbitrary_types_allowed": True})
class ClientConfig(DictConfig):
    """Client configuration."""

    connect_timeout: int = MISSING
    read_timeout: int = MISSING


@dataclass(config={"arbitrary_types_allowed": True})
class BackendKwargs(DictConfig):
    """Backend configuration.

    Attributes
    ----------
    client_config: DictConfig
        Configuration for the client

    """

    client_config: ClientConfig = MISSING


@dataclass(config={"arbitrary_types_allowed": True})
class S3CommConfig(DictConfig):
    """S3 communication configuration.

    Attributes
    ----------
    bucket_name: str = MISSING
        Name of the S3 bucket
    num_attempts: int = MISSING
        Number of attempts
    backend_kwargs: BackendKwargs
        Backend configuration

    """

    bucket_name: str = MISSING
    num_attempts: int = MISSING
    backend_kwargs: BackendKwargs = MISSING


@dataclass(config={"arbitrary_types_allowed": True})
class WandbSetup(DictConfig):
    """Wand setup configuration.

    Attributes
    ----------
    project: str = MISSING
        Name of the project
    group: str = MISSING
        Name of the group
    tags: list[str] = MISSING
        List of tags
    entity: str = MISSING
        Name of the entity
    mode: str = MISSING
        Mode of the run: "online", "offline"
    name: str = MISSING
        Name of the run, {Config.run_uuid}
    resume: str = MISSING
        Weather to allow resumption
    id: str = MISSING
        ID of the run, {Config.run_uuid}
    allow_val_change: bool = MISSING
        Allows changing the value of the config when resuming

    """

    project: str = MISSING
    group: str = MISSING
    tags: list[str] = MISSING
    entity: str | None = MISSING
    mode: str = MISSING
    name: str = MISSING
    resume: str = MISSING
    id: str = MISSING
    allow_val_change: bool = MISSING


@dataclass(config={"arbitrary_types_allowed": True})
class Wandb(DictConfig):
    """Wandb configuration.

    Attributes
    ----------
    setup: WandbSetup
        Wandb setup configuration

    """

    setup: WandbSetup = MISSING


class Dataset(dict[str, Any], DictConfig):  # type: ignore[reportIncompatibleMethodOverride,misc]
    """Dataset configuration."""


class LLMConfig(dict[str, Any], DictConfig):  # type: ignore[reportIncompatibleMethodOverride,misc]
    """LLM configuration."""


@dataclass(config={"arbitrary_types_allowed": True})
class BaseConfig(DictConfig):
    """Base configuration.

    Attributes
    ----------
    run_uuid: str = MISSING
        Run UUID
    seed: int = MISSING
        Seed
    is_test: bool = MISSING
        Whether to run in test mode
    pretrained_model_path: str = MISSING
        Path to the pretrained model
    wte_parameters_path: str = MISSING
        Path to the model from which to take the WTE parameters
    centralized: Centralized = MISSING
        Centralized configuration
    mt_dao: mt_dao
        mt_dao configuration
    fl: FL
        Federated learning configuration
    s3_comm_config: S3CommConfig
        S3 communication configuration
    use_wandb: bool = MISSING
        Whether to use Wandb
    wandb: Wandb
        Wandb configuration
    llm_config: LLMConfig
        LLM configuration
    dataset: Dataset
        Dataset configuration

    """

    run_uuid: str = MISSING
    seed: int = MISSING
    is_test: bool = MISSING
    pretrained_model_path: str | None = MISSING
    wte_parameters_path: str | None = MISSING
    centralized: Centralized = MISSING
    mt_dao: mt_dao = MISSING
    fl: FL = MISSING
    s3_comm_config: S3CommConfig = MISSING
    use_wandb: bool = MISSING
    wandb: Wandb = MISSING

    # NOTE: MosaicML specific, do not include in the base schema
    llm_config: LLMConfig = MISSING
    dataset: Dataset = MISSING

    def __post_init__(self) -> None:
        """Execute post init checks."""
        # Check the number of total experts against the number of shared experts
        if (
            hasattr(self.llm_config.model, "ffn_config")
            and hasattr(self.llm_config.model.ffn_config, "ff_n_experts")
            and hasattr(self.llm_config.model.ffn_config, "ff_n_active_experts")
            and len(self.fl.experts_assignment_policy.shared_expert_ids) > 0
        ):
            n_total_experts = self.llm_config.model.ffn_config.ff_n_experts
            assert n_total_experts >= len(
                self.fl.experts_assignment_policy.shared_expert_ids,
            ), (
                "Incompatible configuration: n_total_experts"
                f" ({n_total_experts}) must be greater than or equal to"
                " the number of shared_expert_ids"
                f" ({len(self.fl.experts_assignment_policy.shared_expert_ids)})."
            )
            assert min(self.fl.experts_assignment_policy.shared_expert_ids) >= 0, (
                "Incompatible configuration: all shared expert IDs"
                f" ({self.fl.experts_assignment_policy.shared_expert_ids}) must be"
                " non-negative."
            )
            assert (
                max(self.fl.experts_assignment_policy.shared_expert_ids)
                < n_total_experts
            ), (
                "Incompatible configuration: all shared expert IDs"
                f" ({self.fl.experts_assignment_policy.shared_expert_ids}) must be less"
                f" than n_total_experts ({n_total_experts})."
            )


def register_config(name: "str") -> None:
    """Register the base configuration schema."""
    cs = ConfigStore.instance()
    cs.store(name=name, node=BaseConfig)
