#!/usr/bin/env python
# -*- coding: utf-8 -*-
from __future__ import annotations

from collections import Counter
from numbers import Integral
from typing import Any, List, Literal, Optional, Tuple, Union

import dask.array as da
import numpy as np
import xarray as xr

from . import types
from .exceptions import ConflictingArgumentsError, UnexpectedShapeError
from .image_container import ImageContainer

###############################################################################


def reduce_to_slice(L: Union[List, Tuple]) -> Union[int, List, slice, Tuple]:
    # if the list only has one element, then just use it
    if len(L) == 1:
        return slice(L[0], L[0] + 1)
    # if the list has at least 2 elements we can check for sliceable
    # it is convertable to a slice if the step size between each
    # consecutive pair of elements is equal and positive
    # 1. get all the deltas in a list:
    steps = [(L[i + 1] - L[i]) for i in range(len(L) - 1)]
    # 2. check if all the deltas are equal and positive
    if steps[0] > 0 and steps.count(steps[0]) == len(steps):
        return slice(min(L), max(L) + 1, steps[0])
    # if we can't convert to a slice, then just return the list unmodified
    return L


def transpose_to_dims(
    data: types.ArrayLike,
    given_dims: str,
    return_dims: str,
) -> types.ArrayLike:
    """
    This shuffles the data dimensions from given_dims to return_dims. Each dimension
    must be present in given_dims must be used in return_dims

    Parameters
    ----------
    data: types.ArrayLike
        Either a dask array or numpy.ndarray of arbitrary shape but with the dimensions
        specified in given_dims
    given_dims: str
        The dimension ordering of data, "CZYX", "VBTCXZY" etc
    return_dims: str
        The dimension ordering of the return data

    Returns
    -------
    data: types.ArrayLike
        The data with the specified dimension ordering.

    Raises
    ------
    ConflictingArgumentsError
        given_dims and return_dims are incompatible.
    """
    # Use a counter to track that the contents are composed of the same letters
    # and that no letter is repeated
    if (
        Counter(given_dims) != Counter(return_dims)
        or max(Counter(given_dims).values()) > 1
    ):
        raise ConflictingArgumentsError(
            f"given_dims={given_dims} and return_dims={return_dims} are incompatible."
        )

    # Resort the data into return_dims order
    match_map = {dim: given_dims.find(dim) for dim in given_dims}
    transposer = []
    for dim in return_dims:
        transposer.append(match_map[dim])
    data = data.transpose(transposer)

    return data


def reshape_data(
    data: types.ArrayLike, given_dims: str, return_dims: str, **kwargs: Any
) -> types.ArrayLike:
    """
    Reshape the data into return_dims, pad missing dimensions, and prune extra
    dimensions. Warns the user to use the base reader if the depth of the Dimension
    being removed is not 1.

    Parameters
    ----------
    data: types.ArrayLike
        Either a dask array or numpy.ndarray of arbitrary shape but with the dimensions
        specified in given_dims
    given_dims: str
        The dimension ordering of data, "CZYX", "VBTCXZY" etc
    return_dims: str
        The dimension ordering of the return data
    kwargs:
        * C=1 => desired specific channel, if C in the input data has depth 3 then C=1
          returns the 2nd slice (0 indexed)
        * Z=10 => desired specific channel, if Z in the input data has depth 20 then
          Z=10 returns the 11th slice
        * T=[0, 1] => desired specific timepoints, if T in the input data has depth 100
          then T=[0, 1] returns the 1st and 2nd slice (0 indexed)
        * T=(0, 1) => desired specific timepoints, if T in the input data has depth 100
          then T=(0, 1) returns the 1st and 2nd slice (0 indexed)
        * T=(0, -1) => desired specific timepoints, if T in the input data has depth 100
          then T=(0, -1) returns the first and last slice
        * T=range(10) => desired specific timepoints, if T in the input data has depth
          100 then T=range(10) returns the first ten slices
        * T=slice(0, -1, 5) => desired specific timepoints, T=slice(0, -1, 5) returns
          every fifth timepoint

    Returns
    -------
    data: types.ArrayLike
        The data with the specified dimension ordering.

    Raises
    ------
    ConflictingArgumentsError
        Missing dimension in return dims when using range, slice, or multi-index
        dimension selection for the requested dimension.

    IndexError
        Requested dimension index not present in data.

    Examples
    --------
    Specific index selection

    >>> data = np.random.rand((10, 100, 100))
    ... z1 = reshape_data(data, "ZYX", "YX", Z=1)

    List of index selection

    >>> data = np.random.rand((10, 100, 100))
    ... first_and_second = reshape_data(data, "ZYX", "YX", Z=[0, 1])

    Tuple of index selection

    >>> data = np.random.rand((10, 100, 100))
    ... first_and_last = reshape_data(data, "ZYX", "YX", Z=(0, -1))

    Range of index selection

    >>> data = np.random.rand((10, 100, 100))
    ... first_three = reshape_data(data, "ZYX", "YX", Z=range(3))

    Slice selection

    >>> data = np.random.rand((10, 100, 100))
    ... every_other = reshape_data(data, "ZYX", "YX", Z=slice(0, -1, 2))

    Empty dimension expansion

    >>> data = np.random.rand((10, 100, 100))
    ... with_time = reshape_data(data, "ZYX", "TZYX")

    Dimension order shuffle

    >>> data = np.random.rand((10, 100, 100))
    ... as_zx_base = reshape_data(data, "ZYX", "YZX")

    Selections, empty dimension expansions, and dimension order shuffle

    >>> data = np.random.rand((10, 100, 100))
    ... example = reshape_data(data, "CYX", "BSTCZYX", C=slice(0, -1, 3))
    """
    # Check for parameter conflicts
    for dim in given_dims:
        # return_dims='CZYX' and iterable dimensions 'T=range(10)'
        # Dimension is in kwargs
        # Dimension is an iterable
        # Dimension is not in return dimensions
        if (
            isinstance(kwargs.get(dim), (list, tuple, range, slice))
            and dim not in return_dims
        ):
            raise ConflictingArgumentsError(
                f"When selecting a multiple dimension indices, the specified "
                f"dimension must be provided in return_dims. "
                f"return_dims={return_dims}, dimension {dim} = {kwargs.get(dim)}"
            )

    # Process each dimension available
    new_dims = given_dims
    dim_specs = []
    for dim in given_dims:
        # Store index of the dim as it is in given data
        dim_index = given_dims.index(dim)

        # Handle dim in return_dims which means that it is
        # an iterable or None selection
        if dim in return_dims:
            # Specific iterable requested
            if dim in kwargs:
                # Actual dim specification
                # The specification provided for this dimension in the kwargs
                dim_spec = kwargs.get(dim)
                display_dim_spec = dim_spec

                if isinstance(dim_spec, int):
                    dim_spec = slice(dim_spec, dim_spec + 1)

                # Convert operator to standard list or slice
                # dask.Array and numpy.ndarray both natively support
                # List[int] and slices being passed to getitem so no need to cast them
                # to anything different
                if isinstance(dim_spec, (tuple, range)):
                    dim_spec = list(dim_spec)

                # Get the largest absolute value index in the list using min and max
                if isinstance(dim_spec, list):
                    check_selection_max = max([abs(min(dim_spec)), max(dim_spec)])
                    # try to convert to slice if possible
                    dim_spec = reduce_to_slice(dim_spec)

                # Get the largest absolute value index from start and stop of slice
                if isinstance(dim_spec, slice):
                    check_selection_max = max([abs(dim_spec.stop), abs(dim_spec.start)])
            else:
                # Nothing was requested from this dimension
                dim_spec = slice(None, None, None)
                display_dim_spec = dim_spec

                # No op means that it doesn't matter how much data is in this dimension
                check_selection_max = 0

        # Not in return_dims means that it is a fixed integer selection
        else:
            if dim in kwargs:
                # Integer requested
                dim_spec = kwargs.get(dim)
                display_dim_spec = dim_spec

                # Check that integer
                if not isinstance(dim_spec, Integral):
                    raise TypeError(
                        "Dimensions not in output must be integers. "
                        f"Got {type(dim_spec).__name__} for {dim}."
                    )
                check_selection_max = dim_spec
            else:
                dim_spec = 0
                display_dim_spec = dim_spec
                check_selection_max = 0

            # Remove dim from new dims as it is fixed size
            new_dims = new_dims.replace(dim, "")

        # Check that fixed integer request isn't outside of request
        if check_selection_max > data.shape[dim_index]:
            raise IndexError(
                f"Dimension specified with {dim}={display_dim_spec} "
                f"but Dimension shape is {data.shape[dim_index]}."
            )

        # All checks and operations passed, append dim operation to getitem ops
        dim_specs.append(dim_spec)

    # Run getitems
    data = data[tuple(dim_specs)]

    # Add empty dims where dimensions were requested but data doesn't exist
    # Add dimensions to new dims where empty dims are added
    for i, dim in enumerate(return_dims):
        # This dimension wasn't processed
        if dim not in given_dims:
            new_dims = f"{new_dims[:i]}{dim}{new_dims[i:]}"
            data = data.reshape(*data.shape[:i], 1, *data.shape[i:])

    # Any extra dimensions have been removed, only a problem if the depth is > 1
    return transpose_to_dims(
        data, given_dims=new_dims, return_dims=return_dims
    )  # don't pass kwargs or 2 copies


def generate_stack(
    image_container: ImageContainer,
    mode: Literal["data", "dask_data", "xarray_data", "xarray_dask_data"],
    drop_non_matching_scenes: bool = False,
    select_scenes: Optional[
        Union[list[Union[str, int]], tuple[Union[str, int], ...]]
    ] = None,
    scene_character: str = "I",
    scene_coord_values: str = "index",
) -> types.MetaArrayLike:
    """
    Stack each scene contained in the reader into a
    single array. This method handles the logic of determining which
    stack function to use (dask or numpy) and whether or not to return a
    labelled array (xr.DataArray). Users should prefer
    to use one of get_stack, get_dask_stack, get_xarray_stack, or
    get_xarray_dask_stack.

    Parameters
    ----------
    mode: Literal["data", "dask_data", "xarray_data", "xarray_dask_data"]
        String describing the style of data to return. Should be one of:
        "data", "dask_data", "xarray_data", "xarray_dask_data".
    drop_non_matching_scenes: bool
        During the scene iteration process, if the next scene to be added
        to the stack has different shape or dtype, should it be dropped or
        raise an error.  Default: False (raise an error)
    select_scenes: Optional[
            Union[List[Union[str, int]], Tuple[Union[str, int], ...]]]
        Which scenes to stack into a single array. Scenes can be provided
        as a list or tuple of scene indices or names. It is recommended to
        use the scene integer index instead of the scene name to avoid
        duplicate scene name lookup issues.
        Default: None (stack all scenes)
    scene_character: str
        Character to use as the name of the scene dimension on the output
        array. Default "I"
    scene_coord_values : str
        How to assign coordinates to the scene dimension of the final
        array. If scene_coord_values="names" use the scene name from
        the reader object. If scene_coord_values="index" don't attach any
        coordinates and fall back to integer values.
        Default: "index"

    Returns
    -------
    stack: types.MetaArrayLike
        The fully stacked array. This can be 6+ dimensions with Scene being
        the first dimension.

    """

    mode_check = ["data", "dask_data", "xarray_data", "xarray_dask_data"]
    if mode not in mode_check:
        raise ValueError(
            f"Invalid mode kwarg. Found {mode} but should be one of:"
            f"{', '.join(mode_check)}."
        )

    scene_stacks = []
    scene_names = []

    if select_scenes is None:
        select_scenes = list(range(len(image_container.scenes)))

    for i, s in enumerate(select_scenes):
        image_container.set_scene(s)
        data = getattr(image_container, mode)

        # Store the shape and dtype of the first scenes data
        # to check against later scenes. If returning a DataArray
        # store the coords and dims to use in the final output
        if i == 0:
            shape = data.shape
            dtype = data.dtype

            if "xarray" in mode:
                coords = dict(data.coords)
                dims = data.dims

                if scene_character in dims:
                    raise ValueError(
                        f"Provided scene dimension character '{scene_character}' "
                        f"was found in the existing dimensions of the data {dims}"
                    )

        # Check other scenes against the first scene
        else:
            if data.shape != shape:
                if not drop_non_matching_scenes:
                    raise UnexpectedShapeError(
                        f"All scenes must have same shape. Found shape"
                        f"{data.shape} in scene {s} but expected"
                        f"{shape} based on scene {select_scenes[0]}"
                    )
                else:
                    continue
            if data.dtype != dtype:
                if not drop_non_matching_scenes:
                    raise TypeError(
                        f"All scenes must have the same dtype. Found data"
                        f"with dtype {data.dtype} in scene {s} but expected"
                        f"dtype {dtype} based on scene {select_scenes[0]}"
                    )
                else:
                    continue

        scene_stacks.append(data)
        scene_names.append(image_container.current_scene)

    stack = da.stack if "dask" in mode else np.stack

    if "xarray" in mode:
        all_data = stack([x.data for x in scene_stacks])
        if scene_coord_values == "names":
            coords = {scene_character: scene_names, **coords}

        return xr.DataArray(
            all_data,
            dims=(scene_character, *dims),
            coords=coords,
        )

    else:
        return stack(scene_stacks)
