# stdlib
from typing import Any
from typing import Dict
from typing import Sequence
from typing import Tuple
from typing import Union

# third party
import numpy as np
from numpy.typing import NDArray

# relative
from ...adp.data_subject_list import DataSubjectArray
from ..autodp.gamma_tensor import GammaTensor
from ..autodp.phi_tensor import PhiTensor
from ..lazy_repeat_array import lazyrepeatarray


def dp_maximum(
    x: Union[PhiTensor, GammaTensor], y: Union[np.ndarray, PhiTensor, GammaTensor]
) -> Union[PhiTensor, GammaTensor]:
    # TODO: Make this work for GammaTensors
    x_data = x.child
    y_data = y.child if hasattr(y, "child") else y

    output = np.maximum(x_data, y_data)

    min_v, max_v = output.min(), output.max()
    dsl = x.data_subjects  # TODO: fix later
    return PhiTensor(
        child=output,
        data_subjects=dsl,
        min_vals=min_v,
        max_vals=max_v,
    )


def dp_log(input: Union[PhiTensor, GammaTensor]) -> Union[PhiTensor, GammaTensor]:
    if isinstance(input, PhiTensor):
        output = np.log(input.child)
        # min_v, max_v = output.min(), output.max()  # These bounds are a function of private data
        min_v = lazyrepeatarray(
            data=np.log(input.min_vals.data.min()), shape=output.shape
        )
        max_v = lazyrepeatarray(
            data=np.log(input.max_vals.data.max()), shape=output.shape
        )
        dsl = input.data_subjects

        return PhiTensor(
            child=output,
            data_subjects=dsl,
            min_vals=min_v,
            max_vals=max_v,
        )
    elif isinstance(input, GammaTensor):
        output = np.log(input.child)
        # min_v, max_v = output.min(), output.max()  # These bounds are a function of private data
        min_v = lazyrepeatarray(
            data=np.log(input.min_vals.data.min()), shape=output.shape
        )
        max_v = lazyrepeatarray(
            data=np.log(input.max_vals.data.max()), shape=output.shape
        )
        dsl = input.data_subjects

        return GammaTensor(
            child=output,
            data_subjects=dsl,
            min_vals=min_v,
            max_vals=max_v,
        )
    else:
        raise NotImplementedError(f"Undefined behaviour for type: {type(input)}")


def dp_zeros(shape: Tuple) -> Union[PhiTensor, GammaTensor]:
    """
    TODO: Passing in the shape seems unnecessary- it can be inferred from data_subjects.data_subjects_indexed.shape
    output = np.zeros_like(data_subjects.data_subjects_indexed)

    :param shape:
    :param data_subjects:
    :return:
    """
    output = np.zeros(shape)

    # Create Empty DataSubjectSet array
    data_subjects = np.broadcast_to(np.array([DataSubjectArray(set())]), output.shape)

    return GammaTensor(
        child=output,
        data_subjects=data_subjects,
        min_vals=lazyrepeatarray(0, shape),
        max_vals=lazyrepeatarray(1, shape),
    )


def dp_pad(
    input: Union[PhiTensor, GammaTensor],
    width: Union[int, Sequence],
    padding_mode: str = "reflect",
    **kwargs: Dict[Any, Any],
) -> Union[PhiTensor, GammaTensor]:

    data = input.child
    output_data: NDArray = np.pad(data, width, mode=padding_mode, **kwargs)
    min_v = lazyrepeatarray(
        data=min(input.min_vals.data.min(), output_data.min()), shape=output_data.shape
    )
    max_v = lazyrepeatarray(
        data=min(input.max_vals.data.max(), output_data.max()), shape=output_data.shape
    )

    output_data_subjects: NDArray = np.pad(
        input.data_subjects, width, mode=padding_mode, **kwargs
    )

    if isinstance(input, PhiTensor):
        return PhiTensor(
            child=output_data,
            data_subjects=output_data_subjects,
            min_vals=min_v,
            max_vals=max_v,
        )
    elif isinstance(input, GammaTensor):
        return GammaTensor(
            child=output_data,
            data_subjects=output_data_subjects,
            min_vals=min_v,
            max_vals=max_v,
        )
    else:
        raise NotImplementedError(
            f"Padding is not implemented for Input Type: {type(input)}"
        )


def dp_add_at(
    a: Union[PhiTensor, GammaTensor], indices: Tuple, b: Union[PhiTensor, GammaTensor]
) -> Union[PhiTensor, GammaTensor]:
    data_a = a.child
    data_b = b.child

    data_subject_a = a.data_subjects
    # data_subject_b = b.data_subjects

    np.add.at(data_a, indices, data_b)

    # TODO: @Shubham Find np.add.at can be implemented at data_subject_level
    # data_subject_a[indices] = data_subject_a[indices] + data_subject_b
    # The above is an alternate to np.add.at
    # np.add.at(data_subject_a, indices, data_subject_b)

    if isinstance(a, PhiTensor):
        return PhiTensor(
            child=data_a,
            data_subjects=data_subject_a,
            min_vals=data_a.min(),
            max_vals=data_a.max(),
        )
    else:
        return GammaTensor(
            child=data_a,
            data_subjects=data_subject_a,
            min_vals=lazyrepeatarray(data_a.min(), shape=data_a.shape),
            max_vals=lazyrepeatarray(data_a.min(), shape=data_a.shape),
        )


def get_im2col_indices(
    x_shape: Tuple[int, ...],
    field_height: int,
    field_width: int,
    padding: int = 1,
    stride: int = 1,
) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
    # First figure out what the size of the output should be
    N, C, H, W = x_shape
    if (H + 2 * padding - field_height) % stride != 0:
        raise ValueError(
            f"(H + 2 * padding - field_height) H:{H} , padding:{padding} , field_height:{field_height} \
        should be a multiple of stride:{stride}"
        )
    if (W + 2 * padding - field_height) % stride != 0:
        raise ValueError(
            f"(W + 2 * padding - field_height) W:{W} , padding:{padding} , field_height:{field_height} \
        should be a multiple of stride:{stride}"
        )
    out_height = int((H + 2 * padding - field_height) / stride + 1)
    out_width = int((W + 2 * padding - field_width) / stride + 1)

    i0 = np.repeat(np.arange(field_height), field_width)
    i0 = np.tile(i0, C)
    i1 = stride * np.repeat(np.arange(out_height), out_width)
    j0 = np.tile(np.arange(field_width), field_height * C)
    j1 = stride * np.tile(np.arange(out_width), out_height)
    i = i0.reshape(-1, 1) + i1.reshape(1, -1)
    j = j0.reshape(-1, 1) + j1.reshape(1, -1)
    k = np.repeat(np.arange(C), field_height * field_width).reshape(-1, 1)
    return (k.astype(int), i.astype(int), j.astype(int))


def col2im_indices(
    cols: Union[PhiTensor, GammaTensor],
    x_shape: Tuple[int, ...],
    field_height: int = 3,
    field_width: int = 3,
    padding: int = 1,
    stride: int = 1,
) -> Union[PhiTensor, GammaTensor]:
    """An implementation of col2im based on fancy indexing and np.add.at"""
    N, C, H, W = x_shape
    H_padded, W_padded = H + 2 * padding, W + 2 * padding
    x_padded = dp_zeros(shape=(N, C, H_padded, W_padded))

    k, i, j = get_im2col_indices(x_shape, field_height, field_width, padding, stride)
    cols_reshaped = cols.reshape((C * field_height * field_width, -1, N))
    cols_reshaped = cols_reshaped.transpose((2, 0, 1))
    x_padded = dp_add_at(x_padded, (slice(None), k, i, j), cols_reshaped)
    if padding == 0:
        return x_padded
    return x_padded[:, :, padding:-padding, padding:-padding]


def im2col_indices(
    x: Union[PhiTensor, GammaTensor],
    field_height: int,
    field_width: int,
    padding: int = 1,
    stride: int = 1,
) -> Union[PhiTensor, GammaTensor]:
    """An implementation of im2col based on some fancy indexing"""
    # Zero-pad the input
    p = padding

    width: Tuple[Any, ...]
    if len(x.shape) == 4:
        width = ((0, 0), (0, 0), (p, p), (p, p))
    elif len(x.shape) == 3:
        width = ((0, 0), (p, p), (p, p))
    elif len(x.shape) == 2:
        width = ((p, p), (p, p))
    else:
        raise NotImplementedError
    x_padded = dp_pad(x, width)

    k, i, j = get_im2col_indices(x.shape, field_height, field_width, padding, stride)

    cols = x_padded[:, k, i, j]
    C = x.shape[1]
    cols = cols.transpose((1, 2, 0)).reshape((field_height * field_width * C, -1))

    # Not sure why this happens but sometimes this gets a shape of (n, -1)
    if cols.min_vals.shape != cols.shape:
        cols.min_vals.shape = cols.shape
        cols.max_vals.shape = cols.shape
    return cols
