import math

import numpy

from matrix_mpi import MatrixMPI
from vector import Vector


class VectorMPI(MatrixMPI):
    def __init__(self, data=None, shape=None):
        if isinstance(data, Vector):
            self.__data__ = data
        else:
            self.__data__ = Vector(data=data, shape=shape)

        # Calculate how much rows are delegated to the rank
        total_amount_of_rows = self.__data__.shape()[0]
        chunks = numpy.array_split(list(range(total_amount_of_rows)), self.__mpi_size__)
        self.__chunk__ = chunks[self.__mpi_rank__].tolist()

        # Store the delegated rows explicitly for calculations
        self.__rank_subdata__ = Vector(self.__data__[self.__chunk__])

    @staticmethod
    def of(vector: Vector):
        return VectorMPI(vector)

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

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

    def __str__(self):
        return str(self.__data__)

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

    def __add__(self, other):
        if isinstance(other, VectorMPI):
            other = other.__rank_subdata__
        return VectorMPI.of(Vector.flatten(self.__mpi_comm__.allgather(self.__rank_subdata__ + other)))

    def __mul__(self, other):
        if isinstance(other, VectorMPI):
            other = other.__data__

        if isinstance(other, int) or isinstance(other, float):
            result = Vector.flatten(self.__mpi_comm__.allgather(self.__rank_subdata__ * other))
        else:
            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(Vector.flatten(self.__mpi_comm__.allgather(other.get_rank_subdata() * self.get_data())))
        return self * other

    def __truediv__(self, other):
        if isinstance(other, VectorMPI):
            other = other.__rank_subdata__
        return VectorMPI.of(Vector.flatten(self.__mpi_comm__.allgather(self.__rank_subdata__ / 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 math.sqrt(self.__mpi_comm__.allreduce(self.__rank_subdata__.get_abs_sum_of_squares()))

    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(Vector.flatten(self.__mpi_comm__.allgather(self.__rank_subdata__ / self.norm())))