# (C) Copyright 2024 Anemoi contributors.
#
# This software is licensed under the terms of the Apache Licence Version 2.0
# which can be obtained at http://www.apache.org/licenses/LICENSE-2.0.
#
# In applying this licence, ECMWF does not waive the privileges and immunities
# granted to it by virtue of its status as an intergovernmental organisation
# nor does it submit to any jurisdiction.


from __future__ import annotations

import logging
from abc import ABC
from abc import abstractmethod
from typing import TYPE_CHECKING
from typing import Any

import pytorch_lightning as pl
import torch
from hydra.utils import instantiate
from omegaconf import OmegaConf
from timm.scheduler import CosineLRScheduler

from anemoi.models.data_indices.collection import IndexCollection
from anemoi.models.distributed.graph import gather_tensor
from anemoi.models.distributed.shapes import apply_shard_shapes
from anemoi.models.interface import AnemoiModelInterface
from anemoi.models.utils.config import get_multiple_datasets_config
from anemoi.training.losses import get_loss_function
from anemoi.training.losses.base import BaseLoss
from anemoi.training.losses.loss import get_metric_ranges
from anemoi.training.losses.scaler_tensor import grad_scaler
from anemoi.training.losses.scalers import create_scalers
from anemoi.training.losses.scalers.base_scaler import AvailableCallbacks
from anemoi.training.losses.scalers.base_scaler import BaseScaler
from anemoi.training.losses.utils import print_variable_scaling
from anemoi.training.utils.enums import TensorDim
from anemoi.training.utils.variables_metadata import ExtractVariableGroupAndLevel

if TYPE_CHECKING:
    from collections.abc import Mapping

    from torch.distributed.distributed_c10d import ProcessGroup
    from torch_geometric.data import HeteroData

    from anemoi.models.data_indices.collection import IndexCollection
    from anemoi.training.schemas.base_schema import BaseSchema

LOGGER = logging.getLogger(__name__)


class BaseGraphModule(pl.LightningModule, ABC):
    """Abstract base class for Anemoi GNN forecasters using PyTorch Lightning.

    This class encapsulates the shared functionality for distributed training,
    scaling, and evaluation of graph-based neural network models across multiple GPUs and nodes.
    It provides hooks for defining losses, metrics, optimizers, and distributed sharding strategies.

    Key Features
    ------------
    - Supports model and data parallelism through model and reader process groups.
    - Handles graph data via `torch_geometric.data.HeteroData` format.
    - Supports sharded input batches and reconstruction via `allgather`.
    - Integrates modular loss and metric functions with support for variable scaling.
    - Enables deferred creation of variable scalers post-model instantiation.
    - Fully compatible with PyTorch Lightning training and validation loops.

    Subclass Responsibilities
    -------------------------
    Child classes must implement the `_step` method, which defines the forward and loss computation
    for training and validation steps.

    Parameters
    ----------
    config : BaseSchema
        Configuration object defining all parameters.
    graph_data : HeteroData
        Graph-structured input data containing node and edge features.
    statistics : dict
        Dictionary of training statistics (mean, std, etc.) used for normalization.
    statistics_tendencies : dict
        Statistics related to tendencies (if used).
    data_indices : IndexCollection
        Maps feature names to index ranges used for training and loss functions.
    metadata : dict
        Dictionary with metadata such as dataset provenance and variable descriptions.
    supporting_arrays : dict
        Numpy arrays (e.g., topography, masks) needed during inference and stored in checkpoints.

    Attributes
    ----------
    model : AnemoiModelInterface
        Wrapper for the underlying GNN model and its pre/post-processing logic.
    loss : BaseLoss
        Training loss function, optionally supporting variable scaling and sharding.
    metrics : dict[str, BaseLoss | Callable]
        Dictionary of validation metrics (often loss-style) computed during evaluation.
    scalers : dict
        Variable-wise scaling functions (e.g., standardization).
    val_metric_ranges : dict
        Mapping of variable groups for which to calculate validation metrics.
    output_mask : nn.Module
        Masking module that filters outputs during inference.
    multi_step : bool
        Flag to enable autoregressive rollouts (used in multi-step forecasting).
    keep_batch_sharded : bool
        Whether to keep input batches split across GPUs instead of gathering them.

    Distributed Training
    --------------------
    The module can be configured to work in multi-node, multi-GPU environments with support for:
    - Custom communication groups for model and reader parallelism
    - Sharded input and output tensors
    - Support for `ZeroRedundancyOptimizer` and learning rate warmup

    Notes
    -----
    - This class should not be used directly. Subclass it and override `_step`.

    See Also
    --------
    - `AnemoiModelInterface`
    - `BaseLoss`
    - `IndexCollection`
    - `CosineLRScheduler`
    - `create_scalers`, `grad_scaler`

    """

    def __init__(
        self,
        *,
        config: BaseSchema,
        graph_data: HeteroData,
        statistics: dict,
        statistics_tendencies: dict,
        data_indices: IndexCollection,
        metadata: dict,
        supporting_arrays: dict,
    ) -> None:
        """Initialize graph neural network forecaster.

        Parameters
        ----------
        config : DictConfig
            Job configuration
        graph_data : HeteroData
            Graph object
        statistics : dict
            Statistics of the training data
        data_indices : IndexCollection
            Indices of the training data,
        metadata : dict
            Provenance information
        supporting_arrays : dict
            Supporting NumPy arrays to store in the checkpoint

        """
        super().__init__()

        # Handle dictionary of graph_data
        graph_data = {name: data.to(self.device) for name, data in graph_data.items()}
        self.dataset_names = list(graph_data.keys())

        # Create output_mask dictionary for each dataset
        self.output_mask = {}
        for name in self.dataset_names:
            self.output_mask[name] = instantiate(config.model.output_mask, graph_data=graph_data[name])

        # Handle supporting_arrays merge with all output masks
        combined_supporting_arrays = supporting_arrays.copy()
        for dataset_name, mask in self.output_mask.items():
            combined_supporting_arrays[dataset_name].update(mask.supporting_arrays)

        if not hasattr(self.__class__, "task_type"):
            msg = """Subclasses of BaseGraphModule must define a `task_type` class attribute,
                indicating the type of task (e.g., 'forecaster', 'time-interpolator')."""
            raise AttributeError(msg)

        metadata["metadata_inference"]["task"] = self.task_type

        self.model = AnemoiModelInterface(
            statistics=statistics,
            statistics_tendencies=statistics_tendencies,
            data_indices=data_indices,
            metadata=metadata,
            supporting_arrays=combined_supporting_arrays,
            graph_data=graph_data,
            config=config,
        )
        self.config = config

        self.data_indices = data_indices

        self.save_hyperparameters()

        self.statistics_tendencies = statistics_tendencies

        self.logger_enabled = config.diagnostics.log.wandb.enabled or config.diagnostics.log.mlflow.enabled

        # Initialize components for multi-dataset
        self.latlons_data = {}  # plotting only, dict of tensors
        self.scalers = {}  # dict of dict of tensors
        self.updating_scalars = {}  # dict of dict of objects
        self.val_metric_ranges = {}  # dict of dict of lists
        self._scaling_values_log = {}  # dict of dict[str, float]
        self.loss = torch.nn.ModuleDict()
        self.metrics = torch.nn.ModuleDict()

        dataset_variable_groups = get_multiple_datasets_config(self.config.training.variable_groups)
        loss_configs = get_multiple_datasets_config(config.training.training_loss)
        scalers_configs = get_multiple_datasets_config(config.training.scalers)
        val_metrics_configs = get_multiple_datasets_config(config.training.validation_metrics)
        metrics_to_log = get_multiple_datasets_config(config.training.metrics)
        for dataset_name in self.dataset_names:
            self.latlons_data[dataset_name] = graph_data[dataset_name][config.graph.data].x

            # Create dataset-specific metadata extractor
            metadata_extractor = ExtractVariableGroupAndLevel(
                variable_groups=dataset_variable_groups[dataset_name],
                metadata_variables=metadata["dataset"][dataset_name].get("variables_metadata"),
            )

            dataset_scalers, dataset_updating_scalars = create_scalers(
                scalers_configs[dataset_name],
                data_indices=data_indices[dataset_name],
                graph_data=graph_data[dataset_name],
                statistics=statistics[dataset_name],
                statistics_tendencies=(
                    statistics_tendencies[dataset_name] if statistics_tendencies is not None else None
                ),
                metadata_extractor=metadata_extractor,
                output_mask=self.output_mask[dataset_name],
            )
            self.scalers[dataset_name] = dataset_scalers
            self.updating_scalars[dataset_name] = dataset_updating_scalars

            self.val_metric_ranges[dataset_name] = get_metric_ranges(
                metadata_extractor,
                output_data_indices=data_indices[dataset_name].model.output,
                metrics_to_log=metrics_to_log[dataset_name],
            )

            self.loss[dataset_name] = get_loss_function(
                loss_configs[dataset_name],
                dataset_scalers,
                data_indices[dataset_name],
            )

            self.metrics[dataset_name] = self._build_metrics_for_dataset(
                val_metrics_configs[dataset_name],
                scalers=dataset_scalers,
                data_indices=data_indices[dataset_name],
            )
            self._scaling_values_log[dataset_name] = print_variable_scaling(
                self.loss[dataset_name],
                data_indices[dataset_name],
            )

        if config.training.loss_gradient_scaling:
            # Multi-dataset: register hook for each loss
            for loss_fn in self.loss.values():
                loss_fn.register_full_backward_hook(grad_scaler, prepend=False)

        self.is_first_step = True
        self.multi_step = config.training.multistep_input
        self.lr = (
            config.system.hardware.num_nodes
            * config.system.hardware.num_gpus_per_node
            * config.training.lr.rate
            / config.system.hardware.num_gpus_per_model
        )
        self.lr_iterations = config.training.lr.iterations
        self.lr_warmup = config.training.lr.warmup
        self.lr_min = config.training.lr.min
        self.optimizer_settings = config.training.optimizer

        self.model_comm_group = None
        self.reader_groups = None

        reader_group_size = self.config.dataloader.read_group_size

        self.grid_indices = {}
        grid_indices_configs = get_multiple_datasets_config(self.config.dataloader.grid_indices)
        for dataset_name in self.dataset_names:
            self.grid_indices[dataset_name] = instantiate(
                grid_indices_configs[dataset_name],
                reader_group_size=reader_group_size,
            )
            self.grid_indices[dataset_name].setup(graph_data[dataset_name])
        self.grid_dim = -2

        # check sharding support
        self.keep_batch_sharded = self.config.model.keep_batch_sharded
        read_group_supports_sharding = reader_group_size == self.config.system.hardware.num_gpus_per_model
        assert read_group_supports_sharding or not self.keep_batch_sharded, (
            f"Reader group size {reader_group_size} does not match the number of GPUs per model "
            f"{self.config.system.hardware.num_gpus_per_model}, but `model.keep_batch_sharded=True` was set. ",
            "Please set `model.keep_batch_sharded=False` or set `dataloader.read_group_size` ="
            "`hardware.num_gpus_per_model`.",
        )

        # set flag if loss and metrics support sharding
        self._check_sharding_support()

        LOGGER.debug("Multistep: %d", self.multi_step)

        # lazy init model and reader group info, will be set by the DDPGroupStrategy:
        self.model_comm_group_id = 0
        self.model_comm_group_rank = 0
        self.model_comm_num_groups = 1
        self.model_comm_group_size = 1

        self.reader_group_id = 0
        self.reader_group_rank = 0
        self.reader_group_size = 1

        self.grid_shard_shapes = dict.fromkeys(self.dataset_names, None)
        self.grid_shard_slice = dict.fromkeys(self.dataset_names, None)

    def _get_loss_name(self) -> str:
        """Get the loss name for multi-dataset cases."""
        # For multi-dataset, use a generic name or combine dataset names
        return "multi_dataset"

    def _check_sharding_support(self) -> None:
        self.loss_supports_sharding = all(getattr(loss, "supports_sharding", False) for loss in self.loss.values())
        self.metrics_support_sharding = all(
            getattr(metric, "supports_sharding", False)
            for dataset_metrics in self.metrics.values()
            for metric in dataset_metrics.values()
        )

        if not self.loss_supports_sharding and self.keep_batch_sharded:
            unsupported_losses = [
                loss.name for loss in self.loss.values() if not getattr(loss, "supports_sharding", False)
            ]
            LOGGER.warning(
                "Some loss functions do not support sharding: %s. "
                "This may lead to increased memory usage and slower training.",
                ", ".join(unsupported_losses),
            )
        if not self.metrics_support_sharding and self.keep_batch_sharded:
            unsupported_metrics = [
                f"{dataset_name}.{metric_name}"
                for dataset_name, dataset_metrics in self.metrics.items()
                for metric_name, metric in dataset_metrics.items()
                if not getattr(metric, "supports_sharding", False)
            ]
            LOGGER.warning(
                "Some validation metrics do not support sharding: %s. "
                "This may lead to increased memory usage and slower training.",
                ", ".join(unsupported_metrics),
            )

    def _build_metrics_for_dataset(
        self,
        validation_metrics_configs: dict,
        scalers: dict,
        data_indices: IndexCollection,
    ) -> torch.nn.ModuleDict:
        return torch.nn.ModuleDict(
            {
                metric_name: get_loss_function(val_metric_config, scalers=scalers, data_indices=data_indices)
                for metric_name, val_metric_config in validation_metrics_configs.items()
            },
        )

    def forward(self, x: torch.Tensor, **kwargs) -> torch.Tensor:
        """Forward method.

        This method calls the model's forward method with the appropriate
        communication group and sharding information.
        """
        return self.model(
            x,
            model_comm_group=self.model_comm_group,
            grid_shard_shapes=self.grid_shard_shapes,
            **kwargs,
        )

    def on_load_checkpoint(self, checkpoint: torch.nn.Module) -> None:
        self._ckpt_model_name_to_index = {
            dataset_name: data_indices.name_to_index
            for dataset_name, data_indices in checkpoint["hyper_parameters"]["data_indices"].items()
        }

    def _update_scaler_for_dataset(
        self,
        name: str,
        scaler_builder: BaseScaler,
        callback: AvailableCallbacks,
        loss_obj: torch.nn.Module,
        metrics_dict: dict,
        dataset_name: str | None = None,
    ) -> None:
        """Update a single scaler for loss and metrics objects."""
        kwargs = {"model": self.model}
        if dataset_name is not None:
            kwargs["dataset_name"] = dataset_name

        scaler = scaler_builder.update_scaling_values(callback, **kwargs)
        if scaler is None:  # If scalar is None, no update to be applied
            return

        if name in loss_obj.scaler:  # If scalar in loss, update it
            loss_obj.update_scaler(scaler=scaler[1], name=name)  # Only update the values

        for metric in metrics_dict.values():  # If scalar in metrics, update it
            if name in metric.scaler:
                metric.update_scaler(scaler=scaler[1], name=name)  # Only update the values

    def update_scalers(self, callback: AvailableCallbacks) -> None:
        """Update scalers, calling the defined function on them, updating if not None."""
        # Multi-dataset case: {'dataset_a': {'nan_mask_weights': scaler, ...}, 'dataset_b': {...}}
        for dataset_name, dataset_scalers in self.updating_scalars.items():
            for name, scaler_builder in dataset_scalers.items():
                self._update_scaler_for_dataset(
                    name,
                    scaler_builder,
                    callback,
                    self.loss[dataset_name],
                    self.metrics[dataset_name],
                    dataset_name=dataset_name,
                )

    def set_model_comm_group(
        self,
        model_comm_group: ProcessGroup,
        model_comm_group_id: int,
        model_comm_group_rank: int,
        model_comm_num_groups: int,
        model_comm_group_size: int,
    ) -> None:
        self.model_comm_group = model_comm_group
        self.model_comm_group_id = model_comm_group_id
        self.model_comm_group_rank = model_comm_group_rank
        self.model_comm_num_groups = model_comm_num_groups
        self.model_comm_group_size = model_comm_group_size

    def set_reader_groups(
        self,
        reader_groups: list[ProcessGroup],
        reader_group_id: int,
        reader_group_rank: int,
        reader_group_size: int,
    ) -> None:
        self.reader_groups = reader_groups
        self.reader_group_id = reader_group_id
        self.reader_group_rank = reader_group_rank
        self.reader_group_size = reader_group_size

    def _prepare_tensors_for_loss(
        self,
        y_pred: torch.Tensor,
        y: torch.Tensor,
        validation_mode: bool = False,
        dataset_name: str | None = None,
    ) -> tuple[torch.Tensor, torch.Tensor, slice | None]:
        """Prepare tensors for loss computation, handling sharding if necessary.

        Parameters
        ----------
        y_pred : torch.Tensor
            Predicted values
        y : torch.Tensor
            Target values
        validation_mode : bool
            Whether in validation mode

        Returns
        -------
        tuple[torch.Tensor, torch.Tensor, slice | None]
            Prepared y_pred, y, and grid_shard_slice
        """
        # Handle multi-dataset case for grid shard slice and shapes
        assert dataset_name is not None, "dataset_name must be provided for multi-dataset case"
        grid_shard_slice = self.grid_shard_slice[dataset_name]
        grid_shard_shapes = self.grid_shard_shapes[dataset_name]

        is_sharded = grid_shard_slice is not None

        sharding_supported = (self.loss_supports_sharding or validation_mode) and (
            self.metrics_support_sharding or not validation_mode
        )

        if is_sharded and not sharding_supported:  # gather tensors if loss or metrics do not support sharding
            shard_shapes = apply_shard_shapes(y_pred, self.grid_dim, grid_shard_shapes)
            y_pred_full = gather_tensor(torch.clone(y_pred), self.grid_dim, shard_shapes, self.model_comm_group)
            y_full = gather_tensor(torch.clone(y), self.grid_dim, shard_shapes, self.model_comm_group)
            final_grid_shard_slice = None
        else:
            y_pred_full, y_full = y_pred, y
            final_grid_shard_slice = grid_shard_slice

        return y_pred_full, y_full, final_grid_shard_slice

    def _compute_loss(
        self,
        y_pred: torch.Tensor,
        y: torch.Tensor,
        grid_shard_slice: slice | None = None,
        dataset_name: str | None = None,
        **_kwargs,
    ) -> torch.Tensor:
        """Compute the loss function.

        Parameters
        ----------
        y_pred : torch.Tensor
            Predicted values
        y : torch.Tensor
            Target values
        grid_shard_slice : slice | None
            Grid shard slice for distributed training
        dataset_name : str | None
            Dataset name for multi-dataset scenarios
        **_kwargs
            Additional arguments

        Returns
        -------
        torch.Tensor
            Computed loss
        """
        assert dataset_name is not None, "dataset_name must be provided when using multiple datasets"

        return self.loss[dataset_name](
            y_pred,
            y,
            grid_shard_slice=grid_shard_slice,
            group=self.model_comm_group,
        )

    def _compute_metrics(
        self,
        y_pred: torch.Tensor,
        y: torch.Tensor,
        grid_shard_slice: slice | None = None,
        dataset_name: str | None = None,
        **_kwargs,
    ) -> dict[str, torch.Tensor]:
        """Compute validation metrics.

        Parameters
        ----------
        y_pred : torch.Tensor
            Predicted values
        y : torch.Tensor
            Target values
        grid_shard_slice : slice | None
            Grid shard slice for distributed training

        Returns
        -------
        dict[str, torch.Tensor]
            Computed metrics
        """
        return self.calculate_val_metrics(y_pred, y, grid_shard_slice=grid_shard_slice, dataset_name=dataset_name)

    def compute_dataset_loss_metrics(
        self,
        y_pred: torch.Tensor,
        y: torch.Tensor,
        validation_mode: bool = False,
        dataset_name: str | None = None,
        **kwargs,
    ) -> tuple[torch.Tensor | None, dict[str, torch.Tensor], torch.Tensor]:
        """Compute loss and metrics for the given predictions and targets.

        Parameters
        ----------
        y_pred : torch.Tensor
            Predicted values
        y : torch.Tensor
            Target values
        step : int, optional
            Current step
        validation_mode : bool, optional
            Whether to compute validation metrics
        **kwargs
            Additional arguments to pass to loss computation

        Returns
        -------
        tuple[torch.Tensor | None, dict[str, torch.Tensor], torch.Tensor]
            Loss, metrics dictionary (if validation_mode), and full predictions
        """
        # Prepare tensors for loss/metrics computation
        y_pred_full, y_full, grid_shard_slice = self._prepare_tensors_for_loss(
            y_pred,
            y,
            validation_mode=validation_mode,
            dataset_name=dataset_name,
        )

        loss = self._compute_loss(
            y_pred=y_pred_full,
            y=y_full,
            grid_shard_slice=grid_shard_slice,
            dataset_name=dataset_name,
            **kwargs,
        )

        # Compute metrics if in validation mode
        metrics_next = {}
        if validation_mode:
            metrics_next = self._compute_metrics(
                y_pred_full,
                y_full,
                grid_shard_slice=grid_shard_slice,
                dataset_name=dataset_name,
                **kwargs,
            )

        return loss, metrics_next, y_pred

    def compute_loss_metrics(
        self,
        y_pred: dict[str, torch.Tensor],
        y: dict[str, torch.Tensor],
        validation_mode: bool = False,
        **kwargs,
    ) -> tuple[torch.Tensor | None, dict[str, torch.Tensor], dict[str, torch.Tensor]]:
        """Compute loss and metrics for the given predictions and targets.

        Parameters
        ----------
        y_pred : dict[str, torch.Tensor]
            Predicted values
        y : dict[str, torch.Tensor]
            Target values
        step : int, optional
            Current step
        validation_mode : bool, optional
            Whether to compute validation metrics
        **kwargs
            Additional arguments to pass to loss computation

        Returns
        -------
        tuple[torch.Tensor | None, dict[str, torch.Tensor], dict[str, torch.Tensor]]
            Loss, metrics dictionary (if validation_mode), and full predictions
        """
        # Prepare tensors for loss/metrics computation
        total_loss, metrics_next, y_preds = None, {}, {}
        for dataset_name in self.dataset_names:
            dataset_loss, dataset_metrics, y_preds[dataset_name] = self.compute_dataset_loss_metrics(
                y_pred[dataset_name],
                y[dataset_name],
                validation_mode=validation_mode,
                dataset_name=dataset_name,
                **kwargs,
            )

            if dataset_loss is not None:
                dataset_loss_sum = dataset_loss.sum()  # collapse potential multi-scale loss
                total_loss = dataset_loss_sum if total_loss is None else total_loss + dataset_loss_sum

                if validation_mode:
                    loss_obj = self.loss[dataset_name]
                    loss_name = getattr(loss_obj, "name", loss_obj.__class__.__name__.lower())
                    metrics_next[f"{dataset_name}_{loss_name}_loss"] = dataset_loss

            # Prefix dataset name to metric keys
            for metric_name, metric_value in dataset_metrics.items():
                metrics_next[f"{dataset_name}_{metric_name}"] = metric_value

        return total_loss, metrics_next, y_preds

    def on_after_batch_transfer(self, batch: torch.Tensor, _: int) -> torch.Tensor:
        """Assemble batch after transfer to GPU by gathering the batch shards if needed.

        Also normalize the batch in-place if needed.

        Parameters
        ----------
        batch : torch.Tensor
            Batch to transfer

        Returns
        -------
        torch.Tensor
            Batch after transfer
        """
        # Gathering/sharding of batch
        batch = self._setup_batch_sharding(batch)

        # Batch normalization
        batch = self._normalize_batch(batch)

        # Prepare scalers, e.g. init delayed scalers and update scalers
        self._prepare_loss_scalers()

        return batch

    def _setup_batch_sharding(self, batch: dict) -> torch.Tensor:
        """Setup batch sharding before every step.

        If the batch is sharded, it will be setup with the grid shard shapes and slice.
        Otherwise, the batch will be allgathered.

        Parameters
        ----------
        batch : torch.Tensor
            Batch to setup

        Returns
        -------
        torch.Tensor
            Batch after setup
        """
        self.grid_shard_shapes = {}
        self.grid_shard_slice = {}

        for dataset_name in self.grid_indices:
            if self.keep_batch_sharded and self.model_comm_group_size > 1:
                self.grid_shard_shapes[dataset_name] = self.grid_indices[dataset_name].shard_shapes
                self.grid_shard_slice[dataset_name] = self.grid_indices[dataset_name].get_shard_slice(
                    self.reader_group_rank,
                )
            else:
                self.grid_shard_shapes[dataset_name] = None
                self.grid_shard_slice[dataset_name] = None
                batch[dataset_name] = self.allgather_batch(
                    batch[dataset_name],
                    self.grid_indices[dataset_name],
                    self.grid_dim,
                )
        return batch

    def transfer_batch_to_device(
        self,
        batch: dict,
        device: torch.device,
        _dataloader_idx: int = 0,
    ) -> dict:
        """Transfer batch to device, handling dictionary batches."""
        transferred_batch = {}
        for dataset_name, dataset_batch in batch.items():
            transferred_batch[dataset_name] = (
                dataset_batch.to(device, non_blocking=True)
                if isinstance(dataset_batch, torch.Tensor)
                else dataset_batch
            )
        return transferred_batch

    def _normalize_batch(self, batch: torch.Tensor) -> torch.Tensor:
        """Normalize batch for training and validation before every step.

        Parameters
        ----------
        batch : torch.Tensor
            Batch to prepare

        Returns
        -------
        torch.Tensor
            Normalized batch
        """
        for dataset_name in batch:
            batch[dataset_name] = self.model.pre_processors[dataset_name](batch[dataset_name])  # normalized in-place
        return batch

    def _prepare_loss_scalers(self) -> None:
        """Prepare scalers for training and validation before every step."""
        # Delayed scalers need to be initialized after the pre-processors once
        if self.is_first_step:
            self.update_scalers(callback=AvailableCallbacks.ON_TRAINING_START)
            self.is_first_step = False
        self.update_scalers(callback=AvailableCallbacks.ON_BATCH_START)
        return

    @abstractmethod
    def _step(
        self,
        batch: torch.Tensor,
        validation_mode: bool = False,
    ) -> tuple[torch.Tensor, Mapping[str, torch.Tensor]]:
        pass

    def allgather_batch(self, batch: torch.Tensor, grid_indices: dict, grid_dim: int) -> torch.Tensor:
        """Allgather the batch-shards across the reader group.

        Parameters
        ----------
        batch : torch.Tensor
            Batch-shard of current reader rank
        grid_indices :
            Grid indices object with shard_shapes and grid_size
        grid_dim : int
            Grid dimension

        Returns
        -------
        torch.Tensor
            Allgathered (full) batch
        """
        grid_shard_shapes = grid_indices.shard_shapes
        grid_size = grid_indices.grid_size

        if grid_size == batch.shape[grid_dim] or self.reader_group_size == 1:
            return batch  # already have the full grid

        shard_shapes = apply_shard_shapes(batch, grid_dim, grid_shard_shapes)
        tensor_list = [torch.empty(shard_shape, device=batch.device, dtype=batch.dtype) for shard_shape in shard_shapes]

        torch.distributed.all_gather(
            tensor_list,
            batch,
            group=self.reader_groups[self.reader_group_id],
        )

        return torch.cat(tensor_list, dim=grid_dim)

    def calculate_val_metrics(
        self,
        y_pred: torch.Tensor,
        y: torch.Tensor,
        grid_shard_slice: slice | None = None,
        dataset_name: str | None = None,
        step: int | None = None,
        **_kwargs,
    ) -> dict[str, torch.Tensor]:
        """Calculate metrics on the validation output.

        Parameters
        ----------
        y_pred: torch.Tensor
            Predicted ensemble
        y: torch.Tensor
            Ground truth (target).
        step: int, optional
            Step number

        Returns
        -------
        val_metrics : dict[str, torch.Tensor]
            validation metrics and predictions
        """
        metrics = {}

        # Handle multi-dataset case for post-processors
        assert dataset_name is not None, "dataset_name must be provided for multi-dataset case"
        post_processor = self.model.post_processors[dataset_name]
        metrics_dict = self.metrics[dataset_name]
        val_metric_ranges = self.val_metric_ranges[dataset_name]

        y_postprocessed = post_processor(y, in_place=False)
        y_pred_postprocessed = post_processor(y_pred, in_place=False)

        suffix = "" if step is None else f"/{step + 1}"
        for metric_name, metric in metrics_dict.items():
            if not isinstance(metric, BaseLoss):
                # If not a loss, we cannot feature scale, so call normally
                metrics[f"{metric_name}_metric/{dataset_name}{suffix}"] = metric(
                    y_pred_postprocessed,
                    y_postprocessed,
                    grid_shard_slice=grid_shard_slice,
                    model_comm_group=self.model_comm_group,
                    model_comm_group_size=self.model_comm_group_size,
                    grid_dim=self.grid_dim,
                    grid_shard_shapes=self.grid_shard_shapes,
                )
                continue

            for mkey, indices in val_metric_ranges.items():
                metric_step_name = f"{metric_name}_metric/{dataset_name}/{mkey}{suffix}"
                if len(metric.scaler.subset_by_dim(TensorDim.VARIABLE.value)):
                    exception_msg = (
                        "Validation metrics cannot be scaled over the variable dimension"
                        " in the post processed space."
                    )
                    raise ValueError(exception_msg)

                metrics[metric_step_name] = metric(
                    y_pred_postprocessed,
                    y_postprocessed,
                    scaler_indices=[..., indices],
                    grid_shard_slice=grid_shard_slice,
                    group=self.model_comm_group,
                    model_comm_group_size=self.model_comm_group_size,
                    grid_dim=self.grid_dim,
                    grid_shard_shapes=self.grid_shard_shapes,
                )

        return metrics

    def training_step(self, batch: dict[str, torch.Tensor], batch_idx: int) -> torch.Tensor:
        del batch_idx
        # Get batch size (handle dict of tensors)
        batch_size = next(iter(batch.values())).shape[0]

        train_loss, *_ = self._step(batch)
        train_loss = train_loss.sum()

        self.log(
            "train_" + self._get_loss_name() + "_loss",
            train_loss,
            on_epoch=True,
            on_step=True,
            prog_bar=True,
            logger=self.logger_enabled,
            batch_size=batch_size,
            sync_dist=True,
        )

        return train_loss

    def validation_step(self, batch: dict[str, torch.Tensor], batch_idx: int) -> None:
        """Calculate the loss over a validation batch using the training loss function.

        Parameters
        ----------
        batch : dict[str, torch.Tensor]
            Validation batch
        batch_idx : int
            Batch inces
        """
        del batch_idx

        # Get batch size (handle dict of tensors)
        batch_size = next(iter(batch.values())).shape[0]

        with torch.no_grad():
            val_loss_scales, metrics, *args = self._step(batch, validation_mode=True)
        val_loss = val_loss_scales.sum()

        self.log(
            "val_" + self._get_loss_name() + "_loss",
            val_loss,
            on_epoch=True,
            on_step=True,
            prog_bar=True,
            logger=self.logger_enabled,
            batch_size=batch_size,
            sync_dist=True,
        )

        if val_loss_scales.numel() > 1:
            for scale in range(val_loss_scales.numel()):
                self.log(
                    "val_" + self.loss.name + "_loss" + "_scale_" + str(scale),
                    val_loss_scales[scale],
                    on_epoch=True,
                    on_step=True,
                    prog_bar=False,
                    logger=self.logger_enabled,
                    batch_size=batch_size,
                    sync_dist=True,
                )

        for mname, mvalue in metrics.items():
            for scale in range(mvalue.numel()):

                log_val = mvalue[scale] if mvalue.numel() > 1 else mvalue

                self.log(
                    "val_" + mname + "_scale_" + str(scale),
                    log_val,
                    on_epoch=True,
                    on_step=False,
                    prog_bar=False,
                    logger=self.logger_enabled,
                    batch_size=batch_size,
                    sync_dist=True,
                )

        return val_loss, *args

    def lr_scheduler_step(self, scheduler: CosineLRScheduler, metric: None = None) -> None:
        """Step the learning rate scheduler by Pytorch Lightning.

        Parameters
        ----------
        scheduler : CosineLRScheduler
            Learning rate scheduler object.
        metric : Any
            Metric object for e.g. ReduceLRonPlateau. Default is None.

        """
        del metric
        scheduler.step(epoch=self.trainer.global_step)

    def on_train_epoch_end(self) -> None:
        pass

    def configure_optimizers(self) -> tuple[list[torch.optim.Optimizer], list[dict[str, Any]]]:
        """Create optimizer and LR scheduler based on Hydra config."""
        optimizer = self._create_optimizer_from_config(self.config.training.optimizer)
        scheduler = self._create_scheduler(optimizer)
        return [optimizer], [scheduler]

    def _create_optimizer_from_config(self, opt_cfg: Any) -> torch.optim.Optimizer:
        """Instantiate optimizer directly via Hydra config (_target_ style)."""
        params = filter(lambda p: p.requires_grad, self.parameters())

        # Convert schema to dict if needed
        if hasattr(opt_cfg, "model_dump"):
            opt_cfg = opt_cfg.model_dump(by_alias=True)

        return instantiate(opt_cfg, params=params, lr=self.lr)

    def _create_scheduler(self, optimizer: torch.optim.Optimizer) -> dict[str, Any]:
        """Helper to create the cosine LR scheduler."""
        scheduler = CosineLRScheduler(
            optimizer,
            lr_min=self.lr_min,
            t_initial=self.lr_iterations,
            warmup_t=self.lr_warmup,
        )
        return {"scheduler": scheduler, "interval": "step"}

    def setup(self, stage: str) -> None:
        """Lightning hook that is called after model is initialized but before training starts."""
        # The conditions should be separate, but are combined due to pre-commit hook
        if stage == "fit" and self.trainer.is_global_zero and self.logger is not None:
            # Log hyperparameters on rank 0
            hyper_params = OmegaConf.to_container(self.config, resolve=True)
            hyper_params.update({"variable_loss_scaling": self._scaling_values_log})
            # Log hyperparameters
            self.logger.log_hyperparams(hyper_params)
