# (C) Copyright 2025- 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
from pathlib import Path

import numpy as np
from anemoi.transform.fields import new_field_from_numpy
from anemoi.transform.fields import new_fieldlist_from_list

from anemoi.inference.context import Context
from anemoi.inference.decorators import main_argument
from anemoi.inference.types import BoolArray
from anemoi.inference.types import State

from ..processor import Processor
from . import pre_processor_registry

LOG = logging.getLogger(__name__)


class ExtractBase(Processor):
    """Base class for pre-processors that extract data from the state."""

    # this needs to be set in subclasses
    indexer: BoolArray | slice

    def __init__(self, context: Context) -> None:
        super().__init__(context)
        self._indexer: BoolArray | slice | None = None

    @property
    def indexer(self) -> BoolArray | slice:
        if self._indexer is None:
            raise RuntimeError(f"{type(self).__name__}.indexer is not set. Set it before process().")
        return self._indexer

    @indexer.setter
    def indexer(self, value: BoolArray | slice) -> None:
        if isinstance(value, slice):
            self._indexer = value
            return
        arr = np.asarray(value)
        self._indexer = arr

    def process(self, state: State) -> State:
        """Process the state to extract a subset of points based on the indexer."""
        state = state.copy()
        idx = self.indexer  # validate indexer is set and correct
        result = []
        for field in state["fields"]:
            data = field.to_numpy().flatten()[idx]
            result.append(new_field_from_numpy(data, template=field))
        state["fields"] = new_fieldlist_from_list(result)
        state["latitudes"] = state["latitudes"][idx]
        state["longitudes"] = state["longitudes"][idx]
        return state


@pre_processor_registry.register("extract_mask")
@main_argument("mask")
class ExtractMask(ExtractBase):
    """Extract a subset of points from the state based on a boolean mask.

    Parameters
    ----------
    context : Any
        The context in which the processor is running.
    mask : str
        Either a path to a `.npy` file containing the boolean mask or
        the name of a supporting array found in the checkpoint.
    """

    def __init__(self, context: Context, mask: str) -> None:
        super().__init__(context)

        self._maskname = mask

        if Path(mask).is_file():
            mask = np.load(mask)
        else:
            mask = context.checkpoint.load_supporting_array(mask)

        if not isinstance(mask, np.ndarray) or mask.dtype != bool:
            raise ValueError(
                f"Expected the mask to be a boolean numpy array. Got {type(mask)} with dtype {mask.dtype}."
            )

        self.indexer = mask
        self.npoints = np.sum(mask)

    def __repr__(self) -> str:
        """Return a string representation of the ExtractMask object.

        Returns
        -------
        str
            String representation of the object.
        """
        return f"ExtractMask({self._maskname}, points={self.npoints}/{self.indexer.size})"


@pre_processor_registry.register("extract_slice")
class ExtractSlice(ExtractBase):
    """Extract a subset of points from the state based on a slice.

    Parameters
    ----------
    context : Context
        The context in which the processor is running.
    slice_args : int
        Arguments to create a slice object. This can be a single integer or
        a tuple of integers representing the start, stop, and step of the slice.
    """

    def __init__(self, context: Context, *slice_args: int) -> None:
        super().__init__(context)
        self.indexer = slice(*slice_args)

    def __repr__(self) -> str:
        """Return a string representation of the ExtractSlice object.

        Returns
        -------
        str
            String representation of the object.
        """
        return f"ExtractSlice({self.indexer})"
