import dataclasses
import os
from collections.abc import Callable, Mapping, Sequence
from typing import Any

import torch

from fme.ace.aggregator import (
    InferenceEvaluatorAggregatorConfig,
    OneStepAggregatorConfig,
)
from fme.ace.aggregator.inference.main import InferenceEvaluatorAggregator
from fme.ace.data_loading.config import DataLoaderConfig
from fme.ace.data_loading.getters import get_gridded_data, get_inference_data
from fme.ace.data_loading.gridded_data import (
    ErrorInferenceData,
    GriddedData,
    InferenceGriddedData,
)
from fme.ace.data_loading.inference import InferenceDataLoaderConfig
from fme.ace.requirements import (
    DataRequirements,
    NullDataRequirements,
    PrognosticStateDataRequirements,
)
from fme.ace.stepper import Stepper
from fme.ace.stepper.single_module import StepperConfig
from fme.core.cli import ResumeResultsConfig
from fme.core.dataset.data_typing import VariableMetadata
from fme.core.dataset_info import DatasetInfo
from fme.core.distributed import Distributed
from fme.core.ema import EMAConfig, EMATracker
from fme.core.generics.trainer import EndOfBatchCallback, EndOfEpochCallback
from fme.core.logging_utils import LoggingConfig
from fme.core.optimization import Optimization, OptimizationConfig
from fme.core.rand import set_seed
from fme.core.typing_ import Slice, TensorDict, TensorMapping
from fme.core.weight_ops import CopyWeightsConfig


@dataclasses.dataclass
class WeatherEvaluationConfig:
    """
    Parameters:
        loader: configuration for the data loader used during weather evaluation
        n_forward_steps: number of forward steps to take
        forward_steps_in_memory: number of forward steps to take before
            re-reading data from disk
        epochs: epochs on which to run weather evaluation. By default runs
            weather evaluation every epoch.
        aggregator: configuration of weather evaluation aggregator.
    """

    loader: InferenceDataLoaderConfig
    n_forward_steps: int = 2
    forward_steps_in_memory: int = 2
    epochs: Slice = dataclasses.field(default_factory=lambda: Slice())
    aggregator: InferenceEvaluatorAggregatorConfig = dataclasses.field(
        default_factory=lambda: InferenceEvaluatorAggregatorConfig(
            log_global_mean_time_series=False, log_global_mean_norm_time_series=False
        )
    )

    def __post_init__(self):
        dist = Distributed.get_instance()
        if self.loader.start_indices.n_initial_conditions % dist.world_size != 0:
            raise ValueError(
                "Number of inference initial conditions must be divisible by the "
                "number of parallel workers, got "
                f"{self.loader.start_indices.n_initial_conditions} and "
                f"{dist.world_size}."
            )
        if (
            self.aggregator.log_global_mean_time_series
            or self.aggregator.log_global_mean_norm_time_series
        ):
            # Both of log_global_mean_time_series and
            # log_global_mean_norm_time_series must be False for inline inference.
            self.aggregator.log_global_mean_time_series = False
            self.aggregator.log_global_mean_norm_time_series = False

    def get_inference_data(
        self,
        window_requirements: DataRequirements,
        initial_condition: PrognosticStateDataRequirements,
    ) -> InferenceGriddedData:
        return get_inference_data(
            config=self.loader,
            total_forward_steps=self.n_forward_steps,
            window_requirements=window_requirements,
            initial_condition=initial_condition,
        )


@dataclasses.dataclass
class InlineInferenceConfig:
    """
    Parameters:
        loader: configuration for the data loader used during inference
        n_forward_steps: number of forward steps to take
        forward_steps_in_memory: number of forward steps to take before
            re-reading data from disk
        epochs: epochs on which to run inference. By default runs inference every epoch.
        aggregator: configuration of inline inference aggregator.
    """

    loader: InferenceDataLoaderConfig
    n_forward_steps: int = 2
    forward_steps_in_memory: int = 2
    epochs: Slice = dataclasses.field(default_factory=lambda: Slice())
    aggregator: InferenceEvaluatorAggregatorConfig = dataclasses.field(
        default_factory=lambda: InferenceEvaluatorAggregatorConfig(
            log_global_mean_time_series=False, log_global_mean_norm_time_series=False
        )
    )

    def __post_init__(self):
        dist = Distributed.get_instance()
        if self.loader.start_indices.n_initial_conditions % dist.world_size != 0:
            raise ValueError(
                "Number of inference initial conditions must be divisible by the "
                "number of parallel workers, got "
                f"{self.loader.start_indices.n_initial_conditions} and "
                f"{dist.world_size}."
            )
        if (
            self.aggregator.log_global_mean_time_series
            or self.aggregator.log_global_mean_norm_time_series
        ):
            # Both of log_global_mean_time_series and
            # log_global_mean_norm_time_series must be False for inline inference.
            self.aggregator.log_global_mean_time_series = False
            self.aggregator.log_global_mean_norm_time_series = False

    def get_inference_data(
        self,
        window_requirements: DataRequirements,
        initial_condition: PrognosticStateDataRequirements,
    ) -> InferenceGriddedData:
        return get_inference_data(
            config=self.loader,
            total_forward_steps=self.n_forward_steps,
            window_requirements=window_requirements,
            initial_condition=initial_condition,
        )


@dataclasses.dataclass
class TrainConfig:
    """
    Configuration for training a model.

    Arguments:
        train_loader: Configuration for the training data loader.
        validation_loader: Configuration for the validation data loader.
        stepper: Configuration for the stepper.
        optimization: Configuration for the optimization.
        logging: Configuration for logging.
        max_epochs: Total number of epochs to train for.
        save_checkpoint: Whether to save checkpoints.
        experiment_dir: Directory where checkpoints and logs are saved.
        inference: Configuration for inline inference.
            If None, no inline inference is run,
            and no "best_inline_inference" checkpoint will be saved.
        weather_evaluation: Configuration for weather evaluation.
            If None, no weather evaluation is run. Weather evaluation is not
            used to select checkpoints, but is used to provide metrics.
        n_forward_steps: Number of forward steps during training. Cannot be given
            at the same time as train_n_forward_steps in StepperConfig.
        seed: Random seed for reproducibility. If set, is used for all types of
            randomization, including data shuffling and model initialization.
            If unset, weight initialization is not reproducible but data shuffling is.
        copy_weights_after_batch: Configuration for copying weights from the
            base model to the training model after each batch.
        ema: Configuration for exponential moving average of model weights.
        validate_using_ema: Whether to validate and perform inference using
            the EMA model.
        checkpoint_save_epochs: How often to save epoch-based checkpoints,
            if save_checkpoint is True. If None, checkpoints are only saved
            for the most recent epoch
            (and the best epochs if validate_using_ema == False).
        ema_checkpoint_save_epochs: How often to save epoch-based EMA checkpoints,
            if save_checkpoint is True. If None, EMA checkpoints are only saved
            for the most recent epoch
            (and the best epochs if validate_using_ema == True).
        log_train_every_n_batches: How often to log batch_loss during training.
        checkpoint_every_n_batches: How often to save latest checkpoint during training.
            If 0 is given, checkpoints will not be saved based on batch progress,
            only other factors like pre-emption or being at the end of an epoch.
        segment_epochs: Exit after training for at most this many epochs
            in current job, without exceeding `max_epochs`. Use this if training
            must be run in segments, e.g. due to wall clock limit.
        save_per_epoch_diagnostics: Whether to save per-epoch diagnostics from
            training, validation and inline inference aggregators.
        validation_aggregator: Configuration for the validation aggregator.
        evaluate_before_training: Whether to run validation and inline inference before
            any training is done.
        save_best_inference_epoch_checkpoints: Whether to save a separate checkpoint
            for each epoch where best_inference_error achieves a new minimum.
            Checkpoints are saved as best_inference_ckpt_XXXX.tar.
        resume_results:  Configuration for resuming a previously stopped or finished
            training job. When provided and experiment_dir has no training_checkpoints
            subdirectory, then it is assumed that this is a new run to resume a
            previously completed run and resume_results.existing_dir is recursively
            copied to experiment_dir.
    """

    train_loader: DataLoaderConfig
    validation_loader: DataLoaderConfig
    stepper: StepperConfig
    optimization: OptimizationConfig
    logging: LoggingConfig
    max_epochs: int
    save_checkpoint: bool
    experiment_dir: str
    inference: InlineInferenceConfig | None
    n_forward_steps: int | None = None
    seed: int | None = None
    copy_weights_after_batch: list[CopyWeightsConfig] = dataclasses.field(
        default_factory=list
    )
    ema: EMAConfig = dataclasses.field(default_factory=lambda: EMAConfig())
    weather_evaluation: WeatherEvaluationConfig | None = None
    validate_using_ema: bool = False
    checkpoint_save_epochs: Slice | None = None
    ema_checkpoint_save_epochs: Slice | None = None
    log_train_every_n_batches: int = 100
    checkpoint_every_n_batches: int = 1000
    segment_epochs: int | None = None
    save_per_epoch_diagnostics: bool = False
    validation_aggregator: OneStepAggregatorConfig = dataclasses.field(
        default_factory=lambda: OneStepAggregatorConfig()
    )
    evaluate_before_training: bool = False
    save_best_inference_epoch_checkpoints: bool = False
    resume_results: ResumeResultsConfig | None = None

    def __post_init__(self):
        if (
            isinstance(self.stepper, StepperConfig)
            and self.stepper.train_n_forward_steps is not None
            and self.n_forward_steps is not None
        ):
            raise ValueError(
                "stepper.train_n_forward_steps may not be given at the same time as "
                "n_forward_steps at the top level"
            )

    def set_random_seed(self):
        if self.seed is not None:
            set_seed(self.seed)

    @property
    def inference_n_forward_steps(self) -> int:
        if self.inference is None:
            return 0
        return self.inference.n_forward_steps

    @property
    def inference_aggregator(self) -> InferenceEvaluatorAggregatorConfig | None:
        if self.inference is None:
            return None
        return self.inference.aggregator

    @property
    def checkpoint_dir(self) -> str:
        """
        The directory where checkpoints are saved.
        """
        return os.path.join(self.experiment_dir, "training_checkpoints")

    @property
    def output_dir(self) -> str:
        """
        The directory where output files are saved.
        """
        return os.path.join(self.experiment_dir, "output")

    def get_inference_epochs(self) -> list[int]:
        if self.inference is None:
            return []
        start_epoch = 0 if self.evaluate_before_training else 1
        all_epochs = list(range(start_epoch, self.max_epochs + 1))
        return all_epochs[self.inference.epochs.slice]


class TrainBuilders:
    def __init__(self, config: TrainConfig):
        self.config = config

    def _get_train_window_data_requirements(self) -> DataRequirements:
        return self.config.stepper.get_train_window_data_requirements(
            default_n_forward_steps=self.config.n_forward_steps
        )

    def _get_evaluation_window_data_requirements(self) -> DataRequirements:
        if self.config.inference is None:
            return NullDataRequirements
        return self.config.stepper.get_evaluation_window_data_requirements(
            self.config.inference.forward_steps_in_memory
        )

    def _get_initial_condition_data_requirements(
        self,
    ) -> PrognosticStateDataRequirements:
        return self.config.stepper.get_prognostic_state_data_requirements()

    def get_train_data(self) -> GriddedData:
        data_requirements = self._get_train_window_data_requirements()
        return get_gridded_data(
            self.config.train_loader,
            requirements=data_requirements,
            train=True,
        )

    def get_validation_data(self) -> GriddedData:
        data_requirements = self._get_train_window_data_requirements()
        return get_gridded_data(
            self.config.validation_loader,
            requirements=data_requirements,
            train=False,
        )

    def get_evaluation_inference_data(
        self,
    ) -> InferenceGriddedData:
        if self.config.inference is None:
            return ErrorInferenceData()  # type: ignore
        else:
            return self.config.inference.get_inference_data(
                window_requirements=self._get_evaluation_window_data_requirements(),
                initial_condition=self._get_initial_condition_data_requirements(),
            )

    def get_optimization(self, modules: torch.nn.ModuleList) -> Optimization:
        return self.config.optimization.build(modules, self.config.max_epochs)

    def get_stepper(
        self,
        dataset_info: DatasetInfo,
    ) -> Stepper:
        return self.config.stepper.get_stepper(
            dataset_info=dataset_info,
        )

    def get_ema(self, modules) -> EMATracker:
        return self.config.ema.build(modules)

    def get_end_of_batch_ops(
        self,
        modules: list[torch.nn.Module],
        base_weights: list[Mapping[str, Any]] | None,
    ) -> EndOfBatchCallback:
        if base_weights is not None:

            def copy_after_batch():
                for module, copy_config in zip(
                    modules, self.config.copy_weights_after_batch
                ):
                    copy_config.apply(weights=base_weights, modules=[module])
                return

            return copy_after_batch
        return lambda: None

    def get_end_of_epoch_callback(
        self,
        inference_one_epoch: Callable[
            [InferenceGriddedData, InferenceEvaluatorAggregator, str, int],
            Mapping[str, Any],
        ],
        normalize: Callable[[TensorMapping], TensorDict],
        output_dir: str,
        variable_metadata: Mapping[str, VariableMetadata],
        channel_mean_names: Sequence[str] | None,
        save_diagnostics: bool,
        n_ic_timesteps: int,
    ) -> EndOfEpochCallback:
        if self.config.weather_evaluation is not None:
            data = self.config.weather_evaluation.get_inference_data(
                window_requirements=self._get_evaluation_window_data_requirements(),
                initial_condition=self._get_initial_condition_data_requirements(),
            )
            dataset_info = data.dataset_info.update_variable_metadata(variable_metadata)
            aggregator = self.config.weather_evaluation.aggregator.build(
                dataset_info=dataset_info,
                n_timesteps=self.config.weather_evaluation.n_forward_steps
                + n_ic_timesteps,
                initial_time=data.initial_time,
                normalize=normalize,
                output_dir=output_dir,
                record_step_20=self.config.weather_evaluation.n_forward_steps >= 20,
                channel_mean_names=channel_mean_names,
                save_diagnostics=save_diagnostics,
            )

            def end_of_epoch_ops(epoch: int) -> Mapping[str, Any]:
                if self.config.weather_evaluation is not None:
                    if self.config.weather_evaluation.epochs.contains(epoch):
                        return inference_one_epoch(
                            data,
                            aggregator,
                            "weather_eval",
                            epoch,
                        )
                return {}

            return end_of_epoch_ops

        return lambda epoch: {}
