# (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 copy
import logging
from typing import TYPE_CHECKING

import numpy as np
import torch
from pytorch_lightning.utilities import rank_zero_only

from anemoi.training.diagnostics.callbacks.plot import GraphTrainableFeaturesPlot as _GraphTrainableFeaturesPlot
from anemoi.training.diagnostics.callbacks.plot import PlotHistogram as _PlotHistogram
from anemoi.training.diagnostics.callbacks.plot import PlotLoss as _PlotLoss
from anemoi.training.diagnostics.callbacks.plot import PlotSample as _PlotSample
from anemoi.training.diagnostics.callbacks.plot import PlotSpectrum as _PlotSpectrum

if TYPE_CHECKING:
    from typing import Any
    from typing import Union

    import pytorch_lightning as pl
    from omegaconf import DictConfig

    from anemoi.training.schemas.base_schema import BaseSchema

LOGGER = logging.getLogger(__name__)


class EnsemblePlotMixin:
    """Mixin class for ensemble-specific plotting."""

    def _handle_ensemble_batch_and_output(
        self,
        pl_module: pl.LightningModule,
        output: tuple[torch.Tensor, list[dict[str, torch.Tensor]]],
        batch: dict[str, torch.Tensor],
    ) -> tuple[dict[torch.Tensor], tuple[torch.Tensor, list[dict[str, torch.Tensor]]]]:
        """Handle ensemble batch and output structure.

        Returns
        -------
        tuple
            Processed batch and predictions
        """
        # For ensemble models, batch is a tuple - allgather the full batch first
        batch = {
            dataset: pl_module.allgather_batch(batch[dataset], pl_module.grid_indices[dataset], pl_module.grid_dim)
            for dataset in batch
        }
        # Extract ensemble predictions
        loss, y_preds = output
        y_preds = [
            {
                dataset: pl_module.allgather_batch(
                    pred[dataset],
                    pl_module.grid_indices[dataset],
                    pl_module.grid_dim,
                )
                for dataset in pred
            }
            for pred in y_preds
        ]

        # Return batch (normalized data) and structured output like regular forecaster
        return batch, [loss, y_preds]

    def _get_output_times(self, config: BaseSchema, pl_module: pl.LightningModule) -> tuple:
        """Return times outputted by the model."""
        if config["training"]["model_task"] == "anemoi.training.train.tasks.GraphEnsInterpolator":
            output_times = (len(config.training.explicit_times.target), "time_interp")
        else:
            output_times = (getattr(pl_module, "rollout", 0), "forecast")
        return output_times

    def process(
        self,
        pl_module: pl.LightningModule,
        dataset_name: str,
        outputs: tuple[torch.Tensor, list[dict[str, torch.Tensor]]],
        batch: dict[str, torch.Tensor],
        output_times: tuple,
        members: Union[int, list[int]] = 0,
    ) -> tuple[np.ndarray, np.ndarray]:
        """Process ensemble outputs for metrics plotting.

        Note: Return only the first ensemble member!!!

        Parameters
        ----------
        pl_module : pl.LightningModule
            Lightning module object
        outputs : list
            List of outputs from the model
        batch : torch.Tensor
            Batch tensor (bs, input_steps + forecast_steps, latlon, nvar)
        output_times : tuple
        members : int, list[int], optional
            Ensemble members to plot. If None, all members are returned. Default to 0.

        Returns
        -------
        tuple
            Processed batch and predictions
        """
        # When running in Async mode, it might happen that in the last epoch these tensors
        # have been moved to the cpu (and then the denormalising would fail as the 'input_tensor' would be on CUDA
        # but internal ones would be on the cpu), The lines below allow to address this problem
        if self.latlons is None:
            self.latlons = {}

        if dataset_name not in self.latlons:
            self.latlons[dataset_name] = pl_module.model.model._graph_data[dataset_name][
                pl_module.model.model._graph_name_data
            ].x.detach()
            self.latlons[dataset_name] = np.rad2deg(self.latlons[dataset_name].cpu().numpy())

        total_targets = output_times[0]
        if output_times[1] == "forecast":
            total_targets *= pl_module.n_step_output

        input_tensor = (
            batch[dataset_name][
                :,
                pl_module.n_step_input - 1 : pl_module.n_step_input + total_targets + 1,
                ...,
                pl_module.data_indices[dataset_name].data.output.full,
            ]
            .detach()
            .cpu()
        )
        data = self.post_processors[dataset_name](input_tensor)[self.sample_idx]
        output_tensor = torch.cat(
            tuple(
                self.post_processors[dataset_name](x[dataset_name][:, ...].detach().cpu(), in_place=False)[
                    self.sample_idx : self.sample_idx + 1,
                    :,
                    members,
                    ...,
                ]
                for x in outputs[1]
            ),
        )
        if output_times[1] == "time_interp" and output_tensor.ndim == 5 and output_tensor.shape[0] == 1:
            output_tensor = output_tensor.squeeze(0)
        output_tensor = pl_module.output_mask[dataset_name].apply(output_tensor, dim=-2, fill_value=np.nan).numpy()
        data[1:, ...] = pl_module.output_mask[dataset_name].apply(data[1:, ...], dim=-2, fill_value=np.nan)
        data = data.numpy()

        return data, output_tensor


class EnsemblePerBatchPlotMixin(EnsemblePlotMixin):
    """Mixin for per-batch ensemble plotting callbacks."""

    def on_validation_batch_end(
        self,
        trainer: pl.Trainer,
        pl_module: pl.LightningModule,
        output: tuple[torch.Tensor, list[dict[str, torch.Tensor]]],
        batch: dict[str, torch.Tensor],
        batch_idx: int,
        **kwargs,
    ) -> None:
        if (
            self.config.diagnostics.plot.asynchronous
            and self.config.dataloader.read_group_size > 1
            and pl_module.local_rank == 0
        ):
            LOGGER.warning("Asynchronous plotting can result in NCCL timeouts with reader_group_size > 1.")

        if batch_idx % self.every_n_batches == 0:
            processed_batch, processed_output = self._handle_ensemble_batch_and_output(pl_module, output, batch)
            # When running in Async mode, it might happen that in the last epoch these tensors
            # have been moved to the cpu (and then the denormalising would fail as the 'input_tensor' would be on CUDA
            # but internal ones would be on the cpu), The lines below allow to address this problem
            self.post_processors = copy.deepcopy(pl_module.model.post_processors)
            for dataset_name in self.post_processors:
                for post_processor in self.post_processors[dataset_name].processors.values():
                    if hasattr(post_processor, "nan_locations"):
                        post_processor.nan_locations = pl_module.allgather_batch(
                            post_processor.nan_locations,
                            pl_module.grid_indices[dataset_name],
                            pl_module.grid_dim,
                        )
                self.post_processors[dataset_name] = self.post_processors[dataset_name].cpu()

            output_times = self._get_output_times(self.config, pl_module)

            self.plot(
                trainer,
                pl_module,
                self.dataset_names,
                processed_output,
                processed_batch,
                batch_idx,
                epoch=trainer.current_epoch,
                output_times=output_times,
                **kwargs,
            )


class BaseEnsemblePlotCallback(EnsemblePerBatchPlotMixin):
    """Base class for ensemble plotting callbacks that ensures proper inheritance order."""

    def __init_subclass__(cls, **kwargs):
        """Ensure ensemble mixin comes first in MRO."""
        super().__init_subclass__(**kwargs)
        mro = cls.__mro__

        # Find positions of our key classes
        ensemble_mixin_pos = None
        base_plot_pos = None

        for i, base in enumerate(mro):
            if base.__name__ == "EnsemblePerBatchPlotMixin":
                ensemble_mixin_pos = i
            elif hasattr(base, "__name__") and "BasePerBatchPlotCallback" in base.__name__:
                base_plot_pos = i
                break

        # Warn if ordering might cause issues
        if ensemble_mixin_pos is not None and base_plot_pos is not None and ensemble_mixin_pos > base_plot_pos:
            import warnings

            warnings.warn(
                f"In {cls.__name__}, EnsemblePerBatchPlotMixin should come before "
                f"BasePerBatchPlotCallback in inheritance hierarchy to ensure proper method resolution.",
                UserWarning,
            )


class PlotEnsSample(EnsemblePerBatchPlotMixin, _PlotSample):
    """Plots a post-processed ensemble sample: input, target and prediction."""

    def __init__(
        self,
        config: DictConfig,
        sample_idx: int,
        parameters: list[str],
        accumulation_levels_plot: list[float],
        output_steps: int,
        precip_and_related_fields: list[str] | None = None,
        colormaps: dict[str] | None = None,
        per_sample: int = 6,
        every_n_batches: int | None = None,
        dataset_names: list[str] | None = None,
        members: list | None = None,
        focus_area: list[dict] | None = None,
        **kwargs: Any,
    ) -> None:
        # Initialize PlotSample first
        _PlotSample.__init__(
            self,
            config,
            sample_idx,
            parameters,
            accumulation_levels_plot,
            output_steps,
            precip_and_related_fields,
            colormaps,
            per_sample,
            every_n_batches,
            dataset_names,
            focus_area,
            **kwargs,
        )
        self.plot_members = members

    @rank_zero_only
    def _plot(
        self,
        trainer: pl.Trainer,
        pl_module: pl.LightningModule,
        dataset_names: list[str],
        outputs: tuple[torch.Tensor, list[dict[str, torch.Tensor]]],
        batch: dict[str, torch.Tensor],
        batch_idx: int,
        epoch: int,
        output_times: tuple,
    ) -> None:
        from anemoi.training.diagnostics.plots import plot_predicted_ensemble

        logger = trainer.logger

        for dataset_name in dataset_names:

            # Build dictionary of indices and parameters to be plotted
            diagnostics = (
                []
                if self.config.data.datasets[dataset_name].diagnostic is None
                else self.config.data.datasets[dataset_name].diagnostic
            )
            plot_parameters_dict = {
                pl_module.data_indices[dataset_name].model.output.name_to_index[name]: (name, name not in diagnostics)
                for name in self.parameters
            }

            data, output_tensor = self.process(
                pl_module,
                dataset_name,
                outputs,
                batch,
                output_times=output_times,
                members=self.plot_members,
            )

            # Apply spatial mask
            _, data, output_tensor = self.focus_mask.apply(
                pl_module.model.model._graph_data,
                self.latlons[dataset_name],
                data,
                output_tensor,
            )

            local_rank = pl_module.local_rank
            if output_times[1] == "forecast" and pl_module.n_step_output > 1:
                max_out_steps = min(pl_module.n_step_output, self.output_steps)
                for rollout_step in range(output_times[0]):
                    for out_step in range(max_out_steps):
                        truth_idx = rollout_step * pl_module.n_step_output + out_step + 1
                        fig = plot_predicted_ensemble(
                            parameters=plot_parameters_dict,
                            n_plots_per_sample=4,
                            latlons=self.latlons[dataset_name],
                            clevels=self.accumulation_levels_plot,
                            y_true=data[truth_idx, ...].squeeze(),
                            y_pred=output_tensor[rollout_step, out_step, ...].squeeze(),
                            datashader=self.datashader_plotting,
                            precip_and_related_fields=self.precip_and_related_fields,
                            colormaps=self.colormaps,
                        )

                        self._output_figure(
                            logger,
                            fig,
                            epoch=epoch,
                            tag=(
                                "pred_val_sample_"
                                f"{dataset_name}_rstep{rollout_step:02d}_out{out_step:02d}_"
                                f"batch{batch_idx:04d}_rank{local_rank:01d}{self.focus_mask.tag}"
                            ),
                            exp_log_tag=(
                                "pred_val_sample_"
                                f"{dataset_name}_rstep{rollout_step:02d}_out{out_step:02d}_"
                                f"rank{local_rank:01d}{self.focus_mask.tag}"
                            ),
                        )
            else:
                for rollout_step in range(output_times[0]):
                    fig = plot_predicted_ensemble(
                        parameters=plot_parameters_dict,
                        n_plots_per_sample=4,
                        latlons=self.latlons[dataset_name],
                        clevels=self.accumulation_levels_plot,
                        y_true=data[rollout_step + 1, ...].squeeze(),
                        y_pred=output_tensor[rollout_step, ...].squeeze(),
                        datashader=self.datashader_plotting,
                        precip_and_related_fields=self.precip_and_related_fields,
                        colormaps=self.colormaps,
                    )

                    self._output_figure(
                        logger,
                        fig,
                        epoch=epoch,
                        tag=(
                            f"pred_val_sample_{dataset_name}_rstep{rollout_step:02d}_batch{batch_idx:04d}_"
                            f"rank{local_rank:01d}{self.focus_mask.tag}"
                        ),
                        exp_log_tag=f"pred_val_sample_{dataset_name}_rstep{rollout_step:02d}_rank{local_rank:01d}{self.focus_mask.tag}",
                    )


# Overload callbacks from single forecaster by using them with the first ensemble member
# ================================
class PlotLoss(_PlotLoss):
    """Plots the unsqueezed loss over rollouts for ensemble models."""

    def on_validation_batch_end(
        self,
        trainer: pl.Trainer,
        pl_module: pl.LightningModule,
        outputs: tuple[torch.Tensor, list[dict[str, torch.Tensor]]],
        batch: dict[str, torch.Tensor],
        batch_idx: int,
    ) -> None:
        batch_without_ensemble_dim = {dataset: data[:, :, 0, :, :] for dataset, data in batch.items()}
        super().on_validation_batch_end(
            trainer,
            pl_module,
            outputs,
            batch_without_ensemble_dim,
            batch_idx,
        )


class PlotSpectrum(BaseEnsemblePlotCallback, _PlotSpectrum):
    """Plots Spectrum of first ensemble member using regular PlotSpectrum logic."""

    def __init__(
        self,
        config: DictConfig,
        sample_idx: int,
        parameters: list[str],
        output_steps: int,
        min_delta: float | None = None,
        every_n_batches: int | None = None,
        dataset_names: list[str] | None = None,
        focus_area: list[dict] | None = None,
    ) -> None:
        """Initialise the PlotSpectrum callback."""
        _PlotSpectrum.__init__(
            self,
            config,
            sample_idx,
            parameters,
            output_steps,
            min_delta,
            every_n_batches,
            dataset_names,
            focus_area,
        )


class PlotSample(BaseEnsemblePlotCallback, _PlotSample):
    """Plots a post-processed sample using regular PlotSample logic on first ensemble member."""

    def __init__(
        self,
        config: DictConfig,
        sample_idx: int,
        parameters: list[str],
        accumulation_levels_plot: list[float],
        output_steps: int,
        precip_and_related_fields: list[str] | None = None,
        colormaps: dict[str] | None = None,
        per_sample: int = 6,
        every_n_batches: int | None = None,
        dataset_names: list[str] | None = None,
        focus_area: list[dict] | None = None,
        **kwargs: Any,
    ) -> None:
        """Initialise the PlotSample callback."""
        _PlotSample.__init__(
            self,
            config,
            sample_idx,
            parameters,
            accumulation_levels_plot,
            output_steps,
            precip_and_related_fields,
            colormaps,
            per_sample,
            every_n_batches,
            dataset_names,
            focus_area,
            **kwargs,
        )


class PlotHistogram(BaseEnsemblePlotCallback, _PlotHistogram):
    """Plots histograms comparing target and prediction for ensemble models using first member."""

    def __init__(
        self,
        config: DictConfig,
        sample_idx: int,
        parameters: list[str],
        output_steps: int,
        precip_and_related_fields: list[str] | None = None,
        log_scale: bool = False,
        every_n_batches: int | None = None,
        dataset_names: list[str] | None = None,
        focus_area: list[dict] | None = None,
    ) -> None:
        """Initialise the PlotHistogram callback."""
        _PlotHistogram.__init__(
            self,
            config,
            sample_idx,
            parameters,
            output_steps,
            precip_and_related_fields,
            log_scale,
            every_n_batches,
            dataset_names,
            focus_area,
        )


class GraphTrainableFeaturesPlot(_GraphTrainableFeaturesPlot):
    """Visualize the node & edge trainable features for ensemble models."""

    def __init__(
        self,
        config: DictConfig,
        dataset_names: list[str] | None = None,
        every_n_epochs: int | None = None,
    ) -> None:
        """Initialise the GraphTrainableFeaturesPlot callback."""
        _GraphTrainableFeaturesPlot.__init__(self, config, dataset_names, every_n_epochs)
