from __future__ import annotations

import warnings
from abc import ABC, abstractmethod
from collections import abc as cabc
from collections.abc import Iterator, Mapping, Sequence
from copy import copy
from typing import (
    TYPE_CHECKING,
    ClassVar,
    Literal,
    TypeVar,
    Union,
)

import numpy as np
import pandas as pd
from scipy.sparse import spmatrix

from anndata._warnings import ExperimentalFeatureWarning, ImplicitModificationWarning
from anndata.compat import AwkArray

from ..utils import deprecated, dim_len, ensure_df_homogeneous, warn_once
from .access import ElementRef
from .index import _subset
from .views import as_view, view_update

if TYPE_CHECKING:
    from .anndata import AnnData
    from .raw import Raw


OneDIdx = Union[Sequence[int], Sequence[bool], slice]
TwoDIdx = tuple[OneDIdx, OneDIdx]

I = TypeVar("I", OneDIdx, TwoDIdx, covariant=True)
# TODO: pd.DataFrame only allowed in AxisArrays?
V = Union[pd.DataFrame, spmatrix, np.ndarray]


class AlignedMapping(cabc.MutableMapping, ABC):
    """\
    An abstract base class for Mappings containing array-like values aligned
    to either one or both AnnData axes.
    """

    _allow_df: ClassVar[bool]
    """If this mapping supports heterogeneous DataFrames"""

    _view_class: ClassVar[type[AlignedViewMixin]]
    """The view class for this aligned mapping."""

    _actual_class: ClassVar[type[AlignedActualMixin]]
    """The actual class (which has it’s own data) for this aligned mapping."""

    def __repr__(self):
        return f"{type(self).__name__} with keys: {', '.join(self.keys())}"

    def _ipython_key_completions_(self) -> list[str]:
        return list(self.keys())

    def _validate_value(self, val: V, key: str) -> V:
        """Raises an error if value is invalid"""
        if isinstance(val, AwkArray):
            warn_once(
                "Support for Awkward Arrays is currently experimental. "
                "Behavior may change in the future. Please report any issues you may encounter!",
                ExperimentalFeatureWarning,
                # stacklevel=3,
            )
        for i, axis in enumerate(self.axes):
            if self.parent.shape[axis] == dim_len(val, i):
                continue
            right_shape = tuple(self.parent.shape[a] for a in self.axes)
            actual_shape = tuple(dim_len(val, a) for a, _ in enumerate(self.axes))
            if actual_shape[i] is None and isinstance(val, AwkArray):
                dim = ("obs", "var")[i]
                msg = (
                    f"The AwkwardArray is of variable length in dimension {dim}.",
                    f"Try ak.to_regular(array, {i}) before including the array in AnnData",
                )
            else:
                dims = tuple(("obs", "var")[ax] for ax in self.axes)
                msg = (
                    f"Value passed for key {key!r} is of incorrect shape. "
                    f"Values of {self.attrname} must match dimensions {dims} of parent. "
                    f"Value had shape {actual_shape} while it should have had {right_shape}."
                )
            raise ValueError(msg)

        if not self._allow_df and isinstance(val, pd.DataFrame):
            name = self.attrname.title().rstrip("s")
            val = ensure_df_homogeneous(val, f"{name} {key!r}")
        return val

    @property
    @abstractmethod
    def attrname(self) -> str:
        """What attr for the AnnData is this?"""
        pass

    @property
    @abstractmethod
    def axes(self) -> tuple[Literal[0, 1], ...]:
        """Which axes of the parent is this aligned to?"""
        pass

    @property
    @abstractmethod
    def is_view(self) -> bool:
        pass

    @property
    def parent(self) -> AnnData | Raw:
        return self._parent

    def copy(self):
        d = self._actual_class(self.parent, self._axis)
        for k, v in self.items():
            if isinstance(v, AwkArray):
                # Shallow copy since awkward array buffers are immutable
                d[k] = copy(v)
            else:
                d[k] = v.copy()
        return d

    def _view(self, parent: AnnData, subset_idx: I):
        """Returns a subset copy-on-write view of the object."""
        return self._view_class(self, parent, subset_idx)

    @deprecated("dict(obj)", FutureWarning)
    def as_dict(self) -> dict:
        return dict(self)


class AlignedViewMixin:
    parent: AnnData
    """Reference to parent AnnData view"""

    attrname: str
    """What attribute in the parent is this?"""

    parent_mapping: Mapping[str, V]
    """The object this is a view of."""

    is_view = True

    def __getitem__(self, key: str) -> V:
        return as_view(
            _subset(self.parent_mapping[key], self.subset_idx),
            ElementRef(self.parent, self.attrname, (key,)),
        )

    def __setitem__(self, key: str, value: V):
        value = self._validate_value(value, key)  # Validate before mutating
        warnings.warn(
            f"Setting element `.{self.attrname}['{key}']` of view, "
            "initializing view as actual.",
            ImplicitModificationWarning,
            stacklevel=2,
        )
        with view_update(self.parent, self.attrname, ()) as new_mapping:
            new_mapping[key] = value

    def __delitem__(self, key: str):
        if key not in self:
            raise KeyError(
                "'{key!r}' not found in view of {self.attrname}"
            )  # Make sure it exists before bothering with a copy
        warnings.warn(
            f"Removing element `.{self.attrname}['{key}']` of view, "
            "initializing view as actual.",
            ImplicitModificationWarning,
            stacklevel=2,
        )
        with view_update(self.parent, self.attrname, ()) as new_mapping:
            del new_mapping[key]

    def __contains__(self, key: str) -> bool:
        return key in self.parent_mapping

    def __iter__(self) -> Iterator[str]:
        return iter(self.parent_mapping)

    def __len__(self) -> int:
        return len(self.parent_mapping)


class AlignedActualMixin:
    _data: dict[str, V]
    """Underlying mapping to the data"""

    is_view = False

    def __getitem__(self, key: str) -> V:
        return self._data[key]

    def __setitem__(self, key: str, value: V):
        value = self._validate_value(value, key)
        self._data[key] = value

    def __contains__(self, key: str) -> bool:
        return key in self._data

    def __delitem__(self, key: str):
        del self._data[key]

    def __iter__(self) -> Iterator[str]:
        return iter(self._data)

    def __len__(self) -> int:
        return len(self._data)


class AxisArraysBase(AlignedMapping):
    """\
    Mapping of key→array-like,
    where array-like is aligned to an axis of parent AnnData.
    """

    _allow_df = True
    _dimnames = ("obs", "var")

    @property
    def attrname(self) -> str:
        return f"{self.dim}m"

    @property
    def axes(self) -> tuple[Literal[0, 1]]:
        """Axes of the parent this is aligned to"""
        return (self._axis,)

    @property
    def dim(self) -> str:
        """Name of the dimension this aligned to."""
        return self._dimnames[self._axis]

    def flipped(self) -> AxisArraysBase:
        """Transpose."""
        new = self.copy()
        new.dimension = abs(self._axis - 1)
        return new

    def to_df(self) -> pd.DataFrame:
        """Convert to pandas dataframe."""
        df = pd.DataFrame(index=self.dim_names)
        for key in self.keys():
            value = self[key]
            for icolumn, column in enumerate(value.T):
                df[f"{key}{icolumn + 1}"] = column
        return df

    def _validate_value(self, val: V, key: str) -> V:
        if (
            hasattr(val, "index")
            and isinstance(val.index, cabc.Collection)
            and not val.index.equals(self.dim_names)
        ):
            # Could probably also re-order index if it’s contained
            try:
                pd.testing.assert_index_equal(val.index, self.dim_names)
            except AssertionError as e:
                msg = f"value.index does not match parent’s {self.dim} names:\n{e}"
                raise ValueError(msg) from None
            else:
                msg = "Index.equals and pd.testing.assert_index_equal disagree"
                raise AssertionError(msg)
        return super()._validate_value(val, key)

    @property
    def dim_names(self) -> pd.Index:
        return (self.parent.obs_names, self.parent.var_names)[self._axis]


class AxisArrays(AlignedActualMixin, AxisArraysBase):
    def __init__(
        self,
        parent: AnnData | Raw,
        axis: int,
        vals: Mapping | AxisArraysBase | None = None,
    ):
        self._parent = parent
        if axis not in (0, 1):
            raise ValueError()
        self._axis = axis
        self._data = dict()
        if vals is not None:
            self.update(vals)


class AxisArraysView(AlignedViewMixin, AxisArraysBase):
    def __init__(
        self,
        parent_mapping: AxisArraysBase,
        parent_view: AnnData,
        subset_idx: OneDIdx,
    ):
        self.parent_mapping = parent_mapping
        self._parent = parent_view
        self.subset_idx = subset_idx
        self._axis = parent_mapping._axis


AxisArraysBase._view_class = AxisArraysView
AxisArraysBase._actual_class = AxisArrays


class LayersBase(AlignedMapping):
    """\
    Mapping of key: array-like, where array-like is aligned to both axes of the
    parent anndata.
    """

    _allow_df = False
    attrname = "layers"
    axes = (0, 1)

    # TODO: I thought I had a more elegant solution to overriding this...
    def copy(self) -> Layers:
        d = self._actual_class(self.parent)
        for k, v in self.items():
            d[k] = v.copy()
        return d


class Layers(AlignedActualMixin, LayersBase):
    def __init__(self, parent: AnnData, vals: Mapping | None = None):
        self._parent = parent
        self._data = dict()
        if vals is not None:
            self.update(vals)


class LayersView(AlignedViewMixin, LayersBase):
    def __init__(
        self,
        parent_mapping: LayersBase,
        parent_view: AnnData,
        subset_idx: TwoDIdx,
    ):
        self.parent_mapping = parent_mapping
        self._parent = parent_view
        self.subset_idx = subset_idx


LayersBase._view_class = LayersView
LayersBase._actual_class = Layers


class PairwiseArraysBase(AlignedMapping):
    """\
    Mapping of key: array-like, where both axes of array-like are aligned to
    one axis of the parent anndata.
    """

    _allow_df = False
    _dimnames = ("obs", "var")

    @property
    def attrname(self) -> str:
        return f"{self.dim}p"

    @property
    def axes(self) -> tuple[Literal[0], Literal[0]] | tuple[Literal[1], Literal[1]]:
        """Axes of the parent this is aligned to"""
        return self._axis, self._axis

    @property
    def dim(self) -> str:
        """Name of the dimension this aligned to."""
        return self._dimnames[self._axis]


class PairwiseArrays(AlignedActualMixin, PairwiseArraysBase):
    def __init__(
        self,
        parent: AnnData,
        axis: int,
        vals: Mapping | None = None,
    ):
        self._parent = parent
        if axis not in (0, 1):
            raise ValueError()
        self._axis = axis
        self._data = dict()
        if vals is not None:
            self.update(vals)


class PairwiseArraysView(AlignedViewMixin, PairwiseArraysBase):
    def __init__(
        self,
        parent_mapping: PairwiseArraysBase,
        parent_view: AnnData,
        subset_idx: OneDIdx,
    ):
        self.parent_mapping = parent_mapping
        self._parent = parent_view
        self.subset_idx = (subset_idx, subset_idx)
        self._axis = parent_mapping._axis


PairwiseArraysBase._view_class = PairwiseArraysView
PairwiseArraysBase._actual_class = PairwiseArrays
