from dataclasses import dataclass
from typing import Union, Tuple, Optional

import numpy as np
import torch
from elias.config import Config

from dreifus.vector.vector_base import unpack_nd_params


@dataclass
class Dimensions(Config, tuple):
    w: int
    h: int

    def __init__(self, w: Union[int, 'Dimensions', Tuple[int, int]], h: Optional[int] = None):
        if isinstance(w, dict):
            # TODO: Is there a better way to match a dictionary against variable names?
            h = w['h']
            w = w['w']

        w, h = unpack_nd_params(2, w, h)
        assert w is not None and h is not None, f"Both w and h have to be given, got: {w}, {h}"

        self.w = w
        self.h = h
        # TODO: Is this a proper dataclass if the super constructor is not called?
        super(Dimensions, self).__init__()  # TODO: Do we need that?

    def __new__(cls, w: Union[int, 'Dimensions', Tuple[int, int]], h: Optional[int] = None):
        if isinstance(w, dict):
            h = w['h']
            w = w['w']

        w, h = unpack_nd_params(2, w, h)
        return tuple.__new__(Dimensions, (w, h))

    # def __getitem__(self, idx: int) -> int:
    #     assert 0 <= idx <= 1, "idx of dimension has to be 0 or 1"
    #
    #     if idx == 0:
    #         return self.w
    #     elif idx == 1:
    #         return self.h

    def __add__(self, other: 'Dimensions') -> 'Dimensions':
        return Dimensions(self.w + other.w, self.h + other.h)

    @property
    def x(self) -> int:
        return self.w

    @property
    def y(self) -> int:
        return self.h


def homogenize(tensor_or_array: Union[torch.Tensor, np.ndarray]) -> Union[torch.Tensor, np.ndarray]:
    """
    Homogenizes the given tensor or array. Works for arbitrary numbers of spatial dimensions.
    Assumes that the spatial dimension is the last in the given array.

    Parameters
    ----------
        tensor_or_array:
            A [..., D] array with arbitrary number of leading dimensions and arbitrary size D of the spatial dimension

    Returns
    -------
        A [..., (D + 1)] array with the last (spatial) dimension extended by ones
    """

    if isinstance(tensor_or_array, np.ndarray):
        ones = np.ones((*tensor_or_array.shape[:-1], 1),
                       dtype=tensor_or_array.dtype)
        return np.concatenate([tensor_or_array, ones], axis=-1)
    elif isinstance(tensor_or_array, torch.Tensor):
        ones = torch.ones((*tensor_or_array.shape[:-1], 1),
                          dtype=tensor_or_array.dtype,
                          device=tensor_or_array.device)
        return torch.concat([tensor_or_array, ones], dim=-1)
    else:
        raise ValueError(f"Unsupported array type: {tensor_or_array.__class__}")