# (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 datetime
import logging
from functools import cached_property
from typing import Any

from earthkit.data import Field
from earthkit.data.core.fieldlist import math
from numpy.typing import NDArray

from .coordinates import extract_single_value
from .coordinates import is_scalar
from .metadata import XArrayMetadata

LOG = logging.getLogger(__name__)


class EmptyFieldList:
    """A class to represent an empty list of fields."""

    def __len__(self) -> int:
        """Return the length of the field list."""
        return 0

    def __getitem__(self, i: int) -> Any:
        """Raise an IndexError when trying to access an item.

        Returns
        -------
        Any
            This method does not return anything as it raises an IndexError.

        Raises
        ------
        IndexError
            Always raised to indicate that the list is empty.

        Args
        ----
        i : int
            Index of the item to access.
        """
        raise IndexError(i)

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


class XArrayField(Field):
    """A class to represent a field in an XArray dataset."""

    def __init__(self, owner: Any, selection: Any) -> None:
        """Create a new XArrayField object.

        Parameters
        ----------
        owner : Variable
            The variable that owns this field.
        selection : XArrayDataArray
            A 2D sub-selection of the variable's underlying array.
            This is actually a nD object, but the first dimensions are always 1.
            The other two dimensions are latitude and longitude.
        """
        self.owner = owner
        self.selection = selection

        # Copy the metadata from the owner
        self._md = owner._metadata.copy()

        aliases = {}
        for coord_name, coord_value in self.selection.coords.items():
            if is_scalar(coord_value):
                # Extract the single value from the scalar dimension
                # and store it in the metadata
                coordinate = owner.by_name[coord_name]
                normalised = coordinate.normalise(extract_single_value(coord_value))
                self._md[coord_name] = normalised
                for alias in coordinate.mars_names:
                    aliases[alias] = normalised

        # Add metadata aliases (e.g. levelist == level) only if they are not already present
        for alias, value in aliases.items():
            if alias not in self._md:
                self._md[alias] = value

        # By now, the only dimensions should be latitude and longitude
        self._shape = tuple(list(self.selection.shape)[-2:])
        if math.prod(self._shape) != math.prod(self.selection.shape):
            raise ValueError(f"Invalid shape for selection {self._shape=}, {self.selection.shape=} {self.selection=}")

    @property
    def shape(self) -> tuple[int, int]:
        """Return the shape of the field."""
        return self._shape

    def to_numpy(self, flatten: bool = False, dtype: type | None = None, index: int | None = None) -> NDArray[Any]:
        """Convert the selection to a numpy array.

        Returns
        -------
        NDArray[Any]
            The selection converted to a numpy array.

        Args
        ----
        flatten : bool, optional
            Whether to flatten the array, by default False.
        dtype : Optional[type], optional
            Data type of the array, by default None.
        index : Optional[int], optional
            Index to select a specific element, by default None.
        """
        if index is not None:
            values = self.selection[index].values
        else:
            values = self.selection.values

        assert dtype is None

        if flatten:
            return values.flatten()

        return values

    @cached_property
    def _metadata(self) -> XArrayMetadata:
        """Return the metadata of the field."""
        return XArrayMetadata(self)

    def grid_points(self) -> Any:
        """Return the grid points of the field."""
        return self.owner.grid_points()

    def to_latlon(self, flatten: bool = True) -> dict[str, Any]:
        """Convert the selection to latitude and longitude coordinates.

        Returns
        -------
        Dict[str, Any]
            The latitude and longitude coordinates.

        Args
        ----
        flatten : bool, optional
            Whether to flatten the coordinates, by default True.
        """
        assert flatten
        return dict(lat=self.latitudes, lon=self.longitudes)

    @property
    def resolution(self) -> Any | None:
        """Return the resolution of the field."""
        return None

    @property
    def grid_mapping(self) -> Any:
        """Return the grid mapping of the field."""
        return self.owner.grid_mapping

    @property
    def latitudes(self) -> Any:
        """Return the latitudes of the field."""
        return self.owner.latitudes

    @property
    def longitudes(self) -> Any:
        """Return the longitudes of the field."""
        return self.owner.longitudes

    @property
    def forecast_reference_time(self) -> datetime.datetime:
        """Return the forecast reference time of the field."""
        date, time = self.metadata("date", "time")
        assert len(time) == 4, time
        assert len(date) == 8, date
        yyyymmdd = int(date)
        time = int(time) // 100
        return datetime.datetime(yyyymmdd // 10000, yyyymmdd // 100 % 100, yyyymmdd % 100, time)

    def __repr__(self) -> str:
        """Return a string representation of the field."""
        return f"XArrayField({self._metadata})"

    def _values(self, dtype: type | None = None) -> Any:
        """Return the values of the selection.

        Returns
        -------
        Any
            The values of the selection.

        Args
        ----
        dtype : Optional[type], optional
            Data type of the values, by default None.
        """
        # we don't use .values as this will download the data
        return self.selection
