from __future__ import annotations

import math
import numpy as np
import torch
import warnings

from inspect import stack
from mpi4py import MPI
from pathlib import Path
from typing import List, Union, Tuple, TypeVar, Optional, Iterable
from heat.core.dndarray import DNDarray

def __getitem__(self, key: Union[int, Tuple[int, ...], List[int, ...]]) -> DNDarray:
    """
    Global getter function for DNDarrays.
    Returns a new DNDarray composed of the elements of the original tensor selected by the indices
    given. This does *NOT* redistribute or rebalance the resulting tensor. If the selection of values is
    unbalanced then the resultant tensor is also unbalanced!
    To redistributed the ``DNDarray`` use :func:`balance()` (issue #187)

    Parameters
    ----------
    key : int, slice, Tuple[int,...], List[int,...]
        Indices to get from the tensor.

    Examples
    --------
    >>> a = ht.arange(10, split=0)
    (1/2) >>> tensor([0, 1, 2, 3, 4], dtype=torch.int32)
    (2/2) >>> tensor([5, 6, 7, 8, 9], dtype=torch.int32)
    >>> a[1:6]
    (1/2) >>> tensor([1, 2, 3, 4], dtype=torch.int32)
    (2/2) >>> tensor([5], dtype=torch.int32)
    >>> a = ht.zeros((4,5), split=0)
    (1/2) >>> tensor([[0., 0., 0., 0., 0.],
                        [0., 0., 0., 0., 0.]])
    (2/2) >>> tensor([[0., 0., 0., 0., 0.],
                        [0., 0., 0., 0., 0.]])
    >>> a[1:4, 1]
    (1/2) >>> tensor([0.])
    (2/2) >>> tensor([0., 0.])
    """
    # key can be: int, tuple, list, slice, DNDarray, torch tensor, numpy array, or sequence thereof

    if key is None:
        return self.expand_dims(0)
    if (
        key is ... or isinstance(key, slice) and key == slice(None)
    ):  # latter doesnt work with torch for 0-dim tensors
        return self

    original_split = self.split

    # Single-element indexing
    scalar = np.isscalar(key) or getattr(key, "ndim", 1) == 0
    if scalar:
        # single-element indexing on axis 0
        if self.ndim == 0:
            raise IndexError(
                "Too many indices for DNDarray: DNDarray is 0-dimensional, but 1 were indexed"
            )
        output_shape = self.gshape[1:]
        if original_split is None or original_split == 0:
            output_split = None
        else:
            output_split = original_split - 1
        split_key_is_ordered = 1
        out_is_balanced = True
        backwards_transpose_axes = tuple(range(self.ndim))
        key, root = self.__process_scalar_key(key, indexed_axis=0, return_local_indices=True)
        if root is None:
            # early out for single-element indexing not affecting split axis
            indexed_arr = self.larray[key]
            indexed_arr = DNDarray(
                indexed_arr,
                gshape=output_shape,
                dtype=self.dtype,
                split=output_split,
                device=self.device,
                comm=self.comm,
                balanced=out_is_balanced,
            )
            return indexed_arr
    else:
        # process multi-element key
        (
            self,
            key,
            output_shape,
            output_split,
            split_key_is_ordered,
            key_is_mask_like,
            out_is_balanced,
            root,
            backwards_transpose_axes,
        ) = self.__process_key(key, return_local_indices=True)

    if not self.is_distributed():
        # key is torch-proof, index underlying torch tensor
        indexed_arr = self.larray[key]
        # transpose array back if needed
        self = self.transpose(backwards_transpose_axes)
        return DNDarray(
            indexed_arr,
            gshape=output_shape,
            dtype=self.dtype,
            split=output_split,
            device=self.device,
            comm=self.comm,
            balanced=out_is_balanced,
        )

    if split_key_is_ordered == 1:
        if root is not None:
            # single-element indexing along split axis
            # prepare for Bcast: allocate buffer on all processes
            if self.comm.rank == root:
                indexed_arr = self.larray[key]
            else:
                indexed_arr = torch.zeros(
                    output_shape, dtype=self.larray.dtype, device=self.larray.device
                )
            # broadcast result to all processes
            self.comm.Bcast(indexed_arr, root=root)
            indexed_arr = DNDarray(
                indexed_arr,
                gshape=output_shape,
                dtype=self.dtype,
                split=output_split,
                device=self.device,
                comm=self.comm,
                balanced=out_is_balanced,
            )
            # transpose array back if needed
            self = self.transpose(backwards_transpose_axes)
            return indexed_arr

        # root is None, i.e. indexing does not affect split axis, apply as is
        indexed_arr = self.larray[key]
        # transpose array back if needed
        self = self.transpose(backwards_transpose_axes)
        return DNDarray(
            indexed_arr,
            gshape=output_shape,
            dtype=self.dtype,
            split=output_split,
            device=self.device,
            balanced=out_is_balanced,
            comm=self.comm,
        )

    # key along split axis is not ordered, indices are GLOBAL
    # prepare for communication of indices and data
    counts, displs = self.counts_displs()
    rank, size = self.comm.rank, self.comm.size

    key_is_single_tensor = isinstance(key, torch.Tensor)
    if key_is_single_tensor:
        split_key = key
    else:
        split_key = key[self.split]
    # split_key might be multi-dimensional, flatten it for communication
    if split_key.ndim > 1:
        original_split_key_shape = split_key.shape
        communication_split = output_split - (split_key.ndim - 1)
        split_key = split_key.flatten()
    else:
        communication_split = output_split

    # determine the number of elements to be received from each process
    recv_counts = torch.zeros((size, 1), dtype=torch.int64, device=self.larray.device)
    if key_is_mask_like:
        recv_indices = torch.zeros(
            (len(split_key), len(key)), dtype=split_key.dtype, device=self.larray.device
        )
    else:
        recv_indices = torch.zeros(
            (split_key.shape), dtype=split_key.dtype, device=self.larray.device
        )
    for p in range(size):
        cond1 = split_key >= displs[p]
        cond2 = split_key < displs[p] + counts[p]
        indices_from_p = torch.nonzero(cond1 & cond2, as_tuple=False)
        incoming_indices = split_key[indices_from_p].flatten()
        recv_counts[p, 0] = incoming_indices.numel()
        # store incoming indices in appropiate slice of recv_indices
        start = recv_counts[:p].sum().item()
        stop = start + recv_counts[p].item()
        if incoming_indices.numel() > 0:
            if key_is_mask_like:
                # apply selection to all dimensions
                for i in range(len(key)):
                    recv_indices[start:stop, i] = key[i][indices_from_p].flatten()
                recv_indices[start:stop, self.split] -= displs[p]
            else:
                recv_indices[start:stop] = incoming_indices - displs[p]
    # build communication matrix by sharing recv_counts with all processes
    # comm_matrix rows contain the send_counts for each process, columns contain the recv_counts
    comm_matrix = torch.zeros((size, size), dtype=torch.int64, device=self.larray.device)
    self.comm.Allgather(recv_counts, comm_matrix)
    send_counts = comm_matrix[:, rank]

    # active rank pairs:
    active_rank_pairs = torch.nonzero(comm_matrix, as_tuple=False)

    # Communication build-up:
    active_recv_indices_from = active_rank_pairs[torch.where(active_rank_pairs[:, 1] == rank)][
        :, 0
    ]
    active_send_indices_to = active_rank_pairs[torch.where(active_rank_pairs[:, 0] == rank)][
        :, 1
    ]
    rank_is_active = active_recv_indices_from.numel() > 0 or active_send_indices_to.numel() > 0

    # allocate recv_buf for incoming data
    recv_buf_shape = list(output_shape)
    if communication_split != output_split:
        # split key was flattened, flatten corresponding dims in recv_buf accordingly
        recv_buf_shape = (
            recv_buf_shape[:communication_split]
            + [recv_counts.sum().item()]
            + recv_buf_shape[output_split + 1 :]
        )
    else:
        recv_buf_shape[communication_split] = recv_counts.sum().item()
    recv_buf = torch.zeros(
        tuple(recv_buf_shape), dtype=self.larray.dtype, device=self.larray.device
    )
    if rank_is_active:
        # non-blocking send indices to `active_send_indices_to`
        send_requests = []
        for i in active_send_indices_to:
            start = recv_counts[:i].sum().item()
            stop = start + recv_counts[i].item()
            outgoing_indices = recv_indices[start:stop]
            send_requests.append(self.comm.Isend(outgoing_indices, dest=i))
            del outgoing_indices
        del recv_indices
        for i in active_recv_indices_from:
            # receive indices from `active_recv_indices_from`
            if key_is_mask_like:
                incoming_indices = torch.zeros(
                    (send_counts[i].item(), len(key)),
                    dtype=torch.int64,
                    device=self.larray.device,
                )
            else:
                incoming_indices = torch.zeros(
                    send_counts[i].item(), dtype=torch.int64, device=self.larray.device
                )
            self.comm.Recv(incoming_indices, source=i)
            # prepare send_buf for outgoing data
            if key_is_single_tensor:
                send_buf = self.larray[incoming_indices]
            else:
                if key_is_mask_like:
                    send_key = tuple(
                        incoming_indices[:, i].reshape(-1)
                        for i in range(incoming_indices.shape[1])
                    )
                    send_buf = self.larray[send_key]
                else:
                    send_key = list(key)
                    send_key[self.split] = incoming_indices
                    send_buf = self.larray[tuple(send_key)]
            # non-blocking send requested data to i
            send_requests.append(self.comm.Isend(send_buf, dest=i))
            del send_buf
        # allocate temporary recv_buf to receive data from all active processes
        tmp_recv_buf_shape = recv_buf_shape.copy()
        tmp_recv_buf_shape[communication_split] = recv_counts.max().item()
        tmp_recv_buf = torch.zeros(
            tuple(tmp_recv_buf_shape), dtype=self.larray.dtype, device=self.larray.device
        )
        for i in active_send_indices_to:
            # receive data from i
            tmp_recv_slice = [slice(None)] * tmp_recv_buf.ndim
            tmp_recv_slice[communication_split] = slice(0, recv_counts[i].item())
            self.comm.Recv(tmp_recv_buf[tmp_recv_slice], source=i)
            # write received data to appropriate portion of recv_buf
            cond1 = split_key >= displs[i]
            cond2 = split_key < displs[i] + counts[i]
            recv_buf_indices = torch.nonzero(cond1 & cond2, as_tuple=False).flatten()
            recv_buf_key = [slice(None)] * recv_buf.ndim
            recv_buf_key[communication_split] = recv_buf_indices
            recv_buf[recv_buf_key] = tmp_recv_buf[tmp_recv_slice]
        del tmp_recv_buf
        # wait for all non-blocking communication to finish
        for req in send_requests:
            req.Wait()
    if communication_split != output_split:
        # split_key has been flattened, bring back recv_buf to intended shape
        original_local_shape = (
            output_shape[:communication_split]
            + original_split_key_shape
            + output_shape[output_split + 1 :]
        )
        recv_buf = recv_buf.reshape(original_local_shape)

    # construct indexed array from recv_buf
    indexed_arr = DNDarray(
        recv_buf,
        gshape=output_shape,
        dtype=self.dtype,
        split=output_split,
        device=self.device,
        comm=self.comm,
        balanced=out_is_balanced,
    )
    # transpose array back if needed
    self = self.transpose(backwards_transpose_axes)
    return indexed_arr