""" Lower level layer for slicer.
Mom's spaghetti.
"""
# TODO: Consider boolean array indexing.

from typing import Any, AnyStr, Union, List, Tuple
from abc import abstractmethod
import numbers


class AtomicSlicer:
    """ Wrapping object that will unify slicing across data structures.

    What we support:
        Basic indexing (return references):
        - (start:stop:step) slicing
        - support ellipses
        Advanced indexing (return references):
        - integer array indexing

    Numpy Reference:
        Basic indexing (return views):
        - (start:stop:step) slicing
        - support ellipses and newaxis (alias for None)
        Advanced indexing (return copy):
        - integer array indexing, i.e. X[[1,2], [3,4]]
        - boolean array indexing
        - mixed array indexing (has integer array, ellipses, newaxis in same slice)
    """

    def __init__(self, o: Any, max_dim: Union[None, int, AnyStr] = "auto"):
        """ Provides a consistent slicing API to the object provided.

        Args:
            o: Object to enable consistent slicing.
                Currently supports numpy dense arrays, recursive lists ending with list or numpy.
            max_dim: Max number of dimensions the wrapped object has.
                If set to "auto", max dimensions will be inferred. This comes at compute cost.
        """
        self.o = o
        self.max_dim = max_dim
        if self.max_dim == "auto":
            self.max_dim = UnifiedDataHandler.max_dim(o)

    def __repr__(self) -> AnyStr:
        """ Override default repr for human readability.

        Returns:
            String to display.
        """
        return f"{self.__class__.__name__}({self.o.__repr__()})"

    def __getitem__(self, item: Any) -> Any:
        """ Consistent slicing into wrapped object.

        Args:
            item: Slicing key of type integer or slice.

        Returns:
            Sliced object.

        Raises:
            ValueError: If slicing is not compatible with wrapped object.
        """
        # Turn item into tuple if not already.
        index_tup = unify_slice(item, self.max_dim)

        # Slice according to object type.
        return UnifiedDataHandler.slice(self.o, index_tup, self.max_dim)


def unify_slice(item: Any, max_dim: int, alias_lookup=None) -> Tuple:
    """ Resolves aliases and ellipses in a slice item.

    Args:
        item: Slicing key that is passed to __getitem__.
        max_dim: Max dimension of object to be sliced.
        alias_lookup: AliasLookup structure.

    Returns:
        A tuple representation of the item.
    """
    item = _normalize_slice_key(item)
    index_tup = _normalize_subkey_types(item)
    index_tup = _handle_newaxis_ellipses(index_tup, max_dim)
    if alias_lookup:
        index_tup = _handle_aliases(index_tup, alias_lookup)
    return index_tup


def _normalize_subkey_types(index_tup: Tuple) -> Tuple:
    """ Casts subkeys into basic types such as int.

    Args:
        key: Slicing key that is passed within __getitem__.

    Returns:
        Tuple with subkeys casted to basic types.
    """
    new_index_tup = []  # Gets casted to tuple at the end

    np_int_types = {
        "int8",
        "int16",
        "int32",
        "int64",
        "uint8",
        "uint16",
        "uint32",
        "uint64",
    }
    for subkey in index_tup:
        if _safe_isinstance(subkey, "numpy", np_int_types):
            new_subkey = int(subkey)
        elif _safe_isinstance(subkey, "numpy", "ndarray"):
            if len(subkey.shape) == 1:
                new_subkey = subkey.tolist()
            else:
                raise ValueError(f"Cannot use array of shape {subkey.shape} as subkey.")
        else:
            new_subkey = subkey

        new_index_tup.append(new_subkey)
    return tuple(new_index_tup)


def _normalize_slice_key(key: Any) -> Tuple:
    """ Normalizes slice key into always being a top-level tuple.

    Args:
        key: Slicing key that is passed within __getitem__.

    Returns:
        Expanded slice as a tuple.
    """
    if not isinstance(key, tuple):
        return (key,)
    else:
        return key


def _handle_newaxis_ellipses(index_tup: Tuple, max_dim: int) -> Tuple:
    """ Expands newaxis and ellipses within a slice for simplification.
    This code is mostly adapted from: https://github.com/clbarnes/h5py_like/blob/master/h5py_like/shape_utils.py#L111

    Args:
        index_tup: Slicing key as a tuple.
        max_dim: Maximum number of dimensions in the respective sliceable object.

    Returns:
        Expanded slice as a tuple.
    """
    non_indexes = (None, Ellipsis)
    concrete_indices = sum(idx not in non_indexes for idx in index_tup)
    index_list = []
    # newaxis_at = []
    has_ellipsis = False
    int_count = 0
    for item in index_tup:
        if isinstance(item, numbers.Number):
            int_count += 1

        # NOTE: If we need locations of new axis, re-enable this.
        if item is None:  # pragma: no cover
            pass
            # newaxis_at.append(len(index_list) + len(newaxis_at) - int_count)
        elif item == Ellipsis:
            if has_ellipsis:  # pragma: no cover
                raise IndexError("an index can only have a single ellipsis ('...')")
            has_ellipsis = True
            initial_len = len(index_list)
            while len(index_list) + (concrete_indices - initial_len) < max_dim:
                index_list.append(slice(None))
        else:
            index_list.append(item)

    if len(index_list) > max_dim:  # pragma: no cover
        raise IndexError("too many indices for array")
    while len(index_list) < max_dim:
        index_list.append(slice(None))

    # return index_list, newaxis_at
    return tuple(index_list)


def _handle_aliases(index_tup: Tuple, alias_lookup) -> Tuple:
    new_index_tup = []

    def resolve(item, dim):
        if isinstance(item, slice):
            return item
        # Replace element if in alias lookup, otherwise use original.
        item = alias_lookup.get(dim, item, item)
        return item

    # Go through each element within the index and resolve if needed.
    for dim, item in enumerate(index_tup):
        if isinstance(item, list):
            new_item = []
            for sub_item in item:
                new_item.append(resolve(sub_item, dim))
        else:
            new_item = resolve(item, dim)
        new_index_tup.append(new_item)

    return tuple(new_index_tup)


class Tracked(AtomicSlicer):
    """ Tracked defines an object that slicer wraps."""

    def __init__(self, o: Any, dim: Union[int, List, tuple, None, str] = "auto"):
        """ Defines an object that will be wrapped by slicer.

        Args:
            o: Object that will be tracked for slicer.
            dim: Target dimension(s) slicer will index on for this object.
        """
        super().__init__(o)

        # Protected attribute that can be overriden.
        self._name = None

        # Place dim into coordinate form.
        if dim == "auto":
            self.dim = list(range(self.max_dim))
        elif dim is None:
            self.dim = []
        elif isinstance(dim, int):
            self.dim = [dim]
        elif isinstance(dim, list):
            self.dim = dim
        elif isinstance(dim, tuple):
            self.dim = list(dim)
        else:  # pragma: no cover
            raise ValueError(f"Cannot handle dim of type: {type(dim)}")


class Obj(Tracked):
    """ An object that slicer wraps. """
    def __init__(self, o, dim="auto"):
        super().__init__(o, dim)


class Alias(Tracked):
    """ Defines a tracked object as well as additional __getitem__ keys. """
    def __init__(self, o, dim):
        if not (
            isinstance(dim, int) or (isinstance(dim, (list, tuple)) and len(dim) <= 1)
        ):  # pragma: no cover
            raise ValueError("Aliases must track a single dimension")
        super().__init__(o, dim)


class AliasLookup:
    def __init__(self, aliases):
        self._lookup = {}

        # Populate lookup and merge indexes.
        for _, alias in aliases.items():
            self.update(alias)

    def update(self, alias):
        if alias.dim is None or len(alias.dim) == 0:
            return

        dim = alias.dim[0]
        if dim not in self._lookup:
            self._lookup[dim] = {}

        dim_lookup = self._lookup[dim]
        # NOTE: Alias must be backed by either a list or dictionary.
        itr = enumerate(alias.o) if isinstance(alias.o, list) else alias.o.items()
        for i, x in itr:
            if x not in dim_lookup:
                dim_lookup[x] = set()
            dim_lookup[x].add(i)

    def delete(self, alias):
        '''Delete an alias that exists from lookup'''
        dim = alias.dim[0]
        dim_lookup = self._lookup[dim]
        # NOTE: Alias must be backed by either a list or dictionary.
        itr = enumerate(alias.o) if isinstance(alias.o, list) else alias.o.items()
        for i, x in itr:
            del dim_lookup[x]

    def get(self, dim, target, default=None):
        if dim not in self._lookup:
            return default

        indexes = self._lookup[dim].get(target, None)
        if indexes is None:
            return default

        if len(indexes) == 1:
            return next(iter(indexes))
        else:
            return list(indexes)


def resolve_dim(slicer_index: Tuple, slicer_dim: List) -> List:
    """ Extracts new dim after applying slicing index and maps it back to the original index list. """

    new_slicer_dim = []
    reduced_mask = []

    for _, curr_idx in enumerate(slicer_index):
        if isinstance(curr_idx, (tuple, list, slice)):
            reduced_mask.append(0)
        else:
            reduced_mask.append(1)

    for curr_dim in slicer_dim:
        if reduced_mask[curr_dim] == 0:
            new_slicer_dim.append(curr_dim - sum(reduced_mask[:curr_dim]))

    return new_slicer_dim


def reduced_o(tracked: Tracked) -> Union[List, Any]:
    os = [t.o for t in tracked]
    os = os[0] if len(os) == 1 else os
    return os


class BaseHandler:
    @classmethod
    @abstractmethod
    def head_slice(cls, o, index_tup, max_dim):
        raise NotImplementedError()  # pragma: no cover

    @classmethod
    @abstractmethod
    def tail_slice(cls, o, tail_index, max_dim, flatten=True):
        raise NotImplementedError()  # pragma: no cover

    @classmethod
    @abstractmethod
    def max_dim(cls, o):
        raise NotImplementedError()  # pragma: no cover

    @classmethod
    def default_alias(cls, o):
        return []


class SeriesHandler(BaseHandler):
    @classmethod
    def head_slice(cls, o, index_tup, max_dim):
        head_index = index_tup[0]
        is_element = True if isinstance(head_index, int) else False
        sliced_o = o.iloc[head_index]

        return is_element, sliced_o, 1

    @classmethod
    def tail_slice(cls, o, tail_index, max_dim, flatten=True):
        # NOTE: Series only has one dimension,
        #       call slicer again to end the recursion.
        return AtomicSlicer(o, max_dim=max_dim)[tail_index]

    @classmethod
    def max_dim(cls, o):
        return len(o.shape)

    @classmethod