from matrix_mpi import MatrixMPI
from vector import Vector


class VectorMPI(MatrixMPI):
    __data__: Vector = None

    def __init__(self, data=None, shape=None):
        self.__data__ = Vector(data=data, shape=shape)

    @staticmethod
    def of(vector: Vector):
        return VectorMPI(vector.get_data(), vector.shape())

    def get_vector(self):
        return self.__data__

    def shape(self):
        return self.__data__.shape()

    def __eq__(self, other):
        """
        Return ``self==value``

        :param other: The object to compare to; must be either a ``Vector``, a ``list`` or a ``numpy.ndarray``
        :return: True if data in the same-shaped vectors are equal to the given data in other for each component otherwise False
        """
        if isinstance(other, VectorMPI):
            return self.__data__ == other.__data__
        else:
            return self.__data__ == other

    def transpose(self):
        """
        :return: the transpose of the vector
        """
        return VectorMPI.of(self.__data__.transpose())

    def T(self):
        return self.transpose()

    def __neg__(self):
        return VectorMPI.of(-self.__data__)

    def __add__(self, other):
        if isinstance(other, VectorMPI):
            other = other.__data__
        return VectorMPI.of(self.__data__ + other)

    def __mul__(self, other):
        if isinstance(other, VectorMPI):
            other = other.__data__
        result = self.__data__ * other
        return VectorMPI.of(result) if isinstance(result, Vector) else result

    def __rmul__(self, other):
        if isinstance(other, MatrixMPI):
            return VectorMPI.of(other.get_data() * self.get_vector())
        return self * other

    def __truediv__(self, other):
        if isinstance(other, VectorMPI):
            other = other.__data__
        return VectorMPI.of(self.__data__ / other)

    def norm(self, **kwargs):
        """
        Computes the 2-norm of the vector which is the Frobenius-Norm of a nx1 matrix.

        :param kwargs: ignored
        :return: the 2-norm of the vector
        """
        return self.__data__.norm()

    def normalize(self):
        """
        A normalized vector has the length (norm) 1.
        To achieve that the vector is divided by the norm of itself.

        :return: the normalized vector
        """
        return VectorMPI.of(self.__data__ / self.norm())

    def __getitem__(self, key):
        return self.__data__[key]

    def __setitem__(self, key, value):
        self.__data__[key] = value