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

import numpy as np
import xarray as xr

from .field import XArrayField

LOG = logging.getLogger(__name__)


class Variable:
    """Represents a variable in an xarray dataset.

    Attributes
    ----------
    ds : xr.Dataset
        The xarray dataset.
    variable : xr.DataArray
        The data array representing the variable.
    coordinates : List[Any]
        List of coordinates associated with the variable.
    grid : Any
        The grid associated with the variable.
    time : Any
        The time dimension associated with the variable.
    metadata : Dict[str, Any]
        Metadata associated with the variable.
    """

    def __init__(
        self,
        *,
        ds: xr.Dataset,
        variable: xr.DataArray,
        coordinates: list[Any],
        grid: Any,
        time: Any,
        metadata: dict[str, Any],
    ):
        """Initialize the Variable object.

        Parameters
        ----------
        ds : xr.Dataset
            The xarray dataset.
        variable : xr.DataArray
            The data array representing the variable.
        coordinates : List[Any]
            List of coordinates associated with the variable.
        grid : Any
            The grid associated with the variable.
        time : Any
            The time dimension associated with the variable.
        metadata : Dict[str, Any]
            Metadata associated with the variable.
        """
        self.ds = ds
        self.variable = variable

        self.grid = grid
        self.coordinates = coordinates

        self._metadata = metadata.copy()
        self._metadata.update({"variable": variable.name, "param": variable.name})

        self.time = time

        self.shape = tuple(
            len(c.variable) for c in coordinates if c.is_dim and not c.scalar and not c.is_grid and not c.is_point
        )
        self.names = {
            c.variable.name: c for c in coordinates if c.is_dim and not c.scalar and not c.is_grid and not c.is_point
        }
        self.by_name = {c.variable.name: c for c in coordinates}

        # We need that alias for the time dimension
        self._aliases = dict(valid_datetime="time")

        self.length = math.prod(self.shape)

    @property
    def name(self) -> str:
        """Return the name of the variable."""
        return str(self.variable.name)

    def __len__(self) -> int:
        """Return the length of the variable.

        Returns
        -------
        int
            The length of the variable.
        """
        return self.length

    @property
    def grid_mapping(self) -> dict[str, Any] | None:
        """Return the grid mapping of the variable."""
        grid_mapping = self.variable.attrs.get("grid_mapping", None)
        if grid_mapping is None:
            return None
        return self.ds[grid_mapping].attrs

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

        Returns
        -------
        Any
            The grid points of the variable.
        """
        return self.grid.grid_points

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

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

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

        Returns
        -------
        str
            A string representation of the variable.
        """
        return "Variable[name={},coordinates={},metadata={}]".format(
            self.variable.name,
            self.coordinates,
            self._metadata,
        )

    def __getitem__(self, i: int) -> "XArrayField":
        """Get a 2D field from the variable.

        Parameters
        ----------
        i : int
            Index of the field.

        Returns
        -------
        XArrayField
            The 2D field at the specified index.

        Raises
        ------
        IndexError
            If the index is out of range.
        """
        if i >= self.length:
            raise IndexError(i)

        coords = np.unravel_index(i, self.shape)
        kwargs = {k: v for k, v in zip(self.names, coords)}
        return XArrayField(self, self.variable.isel(kwargs))

    def sel(self, missing: dict[str, Any], **kwargs: Any) -> Optional["Variable"]:
        """Select a subset of the variable based on the given coordinates.

        Parameters
        ----------
        missing : Dict[str, Any]
            Dictionary to store missing coordinates.
        **kwargs : Any
            Coordinates to select.

        Returns
        -------
        Optional[Variable]
            The selected subset of the variable.
        """
        if not kwargs:
            return self

        k, v = kwargs.popitem()

        user_provided_k = k

        if k == "valid_datetime":
            # Ask the Time object to select the valid datetime
            k = self.time.select_valid_datetime(self)
            if k is None:
                return None

        c = self.by_name.get(k)

        # assert c is not None, f"Could not find coordinate {k} in {self.variable.name} {self.coordinates} {list(self.by_name)}"

        if c is None:
            missing[k] = v
            return self.sel(missing, **kwargs)

        i = c.index(v)
        if i is None:
            if k != user_provided_k:
                LOG.warning(f"Could not find {user_provided_k}={v} in {c} (alias of {k})")
            else:
                LOG.warning(f"Could not find {k}={v} in {c}")
            return None

        if c.scalar and i == 0:
            variable = self
        else:
            coordinates = [x.reduced(i) if c is x else x for x in self.coordinates]

            metadata = self._metadata.copy()
            metadata.update({k: v})

            variable = Variable(
                ds=self.ds,
                variable=self.variable.isel({k: i}),
                coordinates=coordinates,
                grid=self.grid,
                time=self.time,
                metadata=metadata,
            )

        return variable.sel(missing, **kwargs)

    def match(self, **kwargs: Any) -> tuple[bool, dict[str, Any] | None]:
        """Match the variable based on the given metadata.

        Parameters
        ----------
        **kwargs : Any
            Metadata to match.

        Returns
        -------
        Tuple[bool, Optional[Dict[str, Any]]]
            A tuple containing a boolean indicating if the match was successful and the remaining metadata.
        """
        if "param" in kwargs:
            assert "variable" not in kwargs
            kwargs["variable"] = kwargs.pop("param")

        if "variable" in kwargs:
            name = kwargs.pop("variable")
            if not isinstance(name, (list, tuple)):
                name = [name]
            if self.variable.name not in name:
                return False, None
            return True, kwargs
        return True, kwargs


class FilteredVariable:
    """Represents a filtered variable based on metadata.

    Attributes
    ----------
    variable : Variable
        The variable to filter.
    kwargs : Any
        Metadata to filter the variable.
    """

    def __init__(self, variable: Variable, **kwargs: Any):
        """Initialize the FilteredVariable object.

        Parameters
        ----------
        variable : Variable
            The variable to filter.
        **kwargs : Any
            Metadata to filter the variable.
        """
        self.variable = variable
        self.kwargs = kwargs

    @cached_property
    def fields(self) -> list["XArrayField"]:
        """Filter the fields of a variable based on metadata."""
        return [
            field
            for field in self.variable
            if all(field.metadata(k, default=None) == v for k, v in self.kwargs.items())
        ]

    @property
    def length(self) -> int:
        """Return the length of the filtered variable."""
        return len(self.fields)

    def __len__(self) -> int:
        """Return the length of the filtered variable.

        Returns
        -------
        int
            The length of the filtered variable.
        """
        return self.length

    def __getitem__(self, i: int) -> "XArrayField":
        """Get a field from the filtered variable.

        Parameters
        ----------
        i : int
            Index of the field.

        Returns
        -------
        XArrayField
            The field at the specified index.

        Raises
        ------
        IndexError
            If the index is out of range.
        """
        if i >= self.length:
            raise IndexError(i)
        return self.fields[i]
