# (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.

import logging
import warnings
from typing import Any

from anemoi.inference.config import Configuration
from anemoi.inference.context import Context
from anemoi.inference.types import ProcessorConfig
from anemoi.inference.types import State

from ..output import ForwardOutput

LOG = logging.getLogger(__name__)


class MaskedOutput(ForwardOutput):
    """Apply mask output class.

    Parameters
    ----------
    context : dict
        The context dictionary.
    mask : Any
        The mask.
    output : dict
        The output configuration dictionary.
    variables : list, optional
        The list of variables to extract, by default None.
    post_processors : Optional[List[ProcessorConfig]], default None
        Post-processors to apply to the input
    output_frequency : int, optional
        The frequency of output, by default None.
    write_initial_state : bool, optional
        Whether to write the initial state, by default None.
    """

    def __init__(
        self,
        context: Context,
        *,
        mask: Any,
        output: Configuration,
        variables: list[str] | None = None,
        post_processors: list[ProcessorConfig] | None = None,
        output_frequency: int | None = None,
        write_initial_state: bool | None = None,
    ) -> None:

        warnings.warn(
            f"The {self.__class__.__name__} is deprecated and will be removed in a future release. "
            "Use extraction post-processors instead.",
            DeprecationWarning,
        )
        super().__init__(
            context,
            output,
            variables=variables,
            post_processors=post_processors,
            output_frequency=output_frequency,
            write_initial_state=write_initial_state,
        )
        self.mask = mask

    def modify_state(self, state: State) -> State:
        """Apply the mask to the state.

        Parameters
        ----------
        state : State
            The state dictionary.

        Returns
        -------
        State
            The masked state dictionary.
        """
        state = state.copy()
        state["fields"] = state["fields"].copy()
        state["latitudes"] = state["latitudes"][self.mask]
        state["longitudes"] = state["longitudes"][self.mask]

        for field in state["fields"]:
            data = state["fields"][field]
            if data.ndim == 1:
                data = data[self.mask]
            else:
                data = data[..., self.mask]
            state["fields"][field] = data

        return state

    def __repr__(self) -> str:
        """Return a string representation of the object."""
        return f"{self.__class__.__name__}({self.mask}, {self.output})"
