import math

import numpy
from mpi4py import MPI

from matrix import Matrix


class MatrixMPI:
    __mpi_comm__ = MPI.COMM_WORLD
    __mpi_size__ = __mpi_comm__.Get_size()
    __mpi_rank__ = __mpi_comm__.Get_rank()

    __data__ = None
    __rank_subdata__ = None
    __chunk__: list = None

    def __init__(self, data=None, shape=None, structure=None, model=None, offset=None, n=None):
        """
        Creates a new matrix.
        The type of the matrix depends on the signature and arguments.

        - ``MatrixMPI(list)``: will create a new matrix with the given data in the list and its shape.
        - ``MatrixMPI(numpy.ndarray)``: will create a new matrix with the given data in ndarray and its shape.
        - ``MatrixMPI(list, (int,int))``: will create a new nxm matrix with the given rows and columns and data in list.
        - ``MatrixMPI(list, str, int, int)``: will create a new square matrix of given size and structure of \"diagonal\"
        - ``MatrixMPI(list, str, int)``: will create a new square matrix of given size and structure of either \"unity\", \"diagonal\" or \"tridiagonal\"
        - ``MatrixMPI(str, int)``: will create a new square matrix of given size and TODO

        :param data: Either a list or an numpy ndarray
        :param shape: A tuple containing the amount of rows and columns
        :param structure: Either \"unity\", \"diagonal\" or \"tridiagonal\"
        :param model: TODO
        :param offset: Offset to diagonal axis
        :param n: Amount of rows of a square matrix or offset in case of diagonal structure

        :type data: Matrix | list | numpy.ndarray
        :type shape: (int, int)
        :type structure: str
        :type model: str
        :type offset: int
        :type n: int

        :rtype: MatrixMPI
        """
        if isinstance(data, Matrix):
            self.__data__ = data
        else:
            self.__data__ = Matrix(data=data, shape=shape, structure=structure, model=model, offset=offset, n=n)

        # 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
        rows = len(self.__chunk__)
        cols = self.__data__.shape()[1]
        self.__rank_subdata__ = Matrix(self.__data__[self.__chunk__], (rows, cols))

    @staticmethod
    def of(matrix: Matrix):
        return MatrixMPI(matrix)

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

    def get_rank_subdata(self):
        """
        Returns only the delegated rows of the rank as ``Matrix``
        :return: The delegated rows as ``Matrix``
        """
        return self.__rank_subdata__

    def get_data(self):
        """
        Returns the whole ``Matrix`` that is used internally
        :return: The ``Matrix`` that is used internally
        """
        return self.__data__

    def get_internal_data(self):
        """
        Returns the raw data of the internal data structure
        :return: The raw data of the internal data structure
        """
        return self.__data__.get_data()

    def transpose(self):
        """
        :return: the transpose of the matrix
        """
        return MatrixMPI.of(self.__data__.transpose())

    def T(self):
        """
        Same as ``matrix.transpose()``

        :return: see ``matrix.transpose()``
        """
        return self.transpose()

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

        :param other: The object to compare to; must be either a ``MatrixMPI``, ``Matrix``, a ``list`` or a ``numpy.ndarray``
        :return: True if data in the matrix are equal to the given data in other for each component, otherwise False
        """
        if isinstance(other, MatrixMPI):
            return all(self.__mpi_comm__.allgather(self.__rank_subdata__ == other.__rank_subdata__))
        else:
            return self.__data__ == other

    def __neg__(self):
        return MatrixMPI.of(Matrix.flatten(self.__mpi_comm__.allgather(-self.__rank_subdata__)))

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

    def __radd__(self, other):
        return self + other

    def __sub__(self, other):
        return self + (-other)

    def __rsub__(self, other):
        return -self + other

    def __truediv__(self, other):
        return MatrixMPI.of(Matrix.flatten(self.__mpi_comm__.allgather(self.__rank_subdata__ / other)))

    def __mul__(self, other):
        if isinstance(other, MatrixMPI):
            other = other.get_data()
        return MatrixMPI.of(Matrix.flatten(self.__mpi_comm__.allgather(self.__rank_subdata__ * other)))

    def __rmul__(self, other):
        return self * other

    def norm(self, f: str = "frobenius"):
        """
        Calculates the norm of the matrix.

        A norm is a positive definit, absolute homogeneous and subadditive function.
        For Matrices a norm is also sub-multiplicative.

        :param f: The norm to be used, could be either "frobenius", "row sum" or "col sum"

        :return: the norm as a number
        """
        if f == "frobenius":
            return math.sqrt(self.__mpi_comm__.allreduce(self.__rank_subdata__.get_abs_sum_of_squares()))
        elif f == "row sum":
            return max(self.__mpi_comm__.allgather(self.__rank_subdata__.norm(f)))
        return self.__data__.norm(f)

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

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