import numpy as np
from mpi4py import MPI


class Vector:
    start_idx = 0  # Nullter Eintrag des Vektors auf dem aktuellen Rang
    end_idx = 0  # Letzer Eintrag des Vektors auf dem aktuellen Rang
    rank_size = 0  # Dimension des Vektors der auf dem aktuellen Rang gespeichert wird
    kind = ''  # Art des Vektors, Zeilen oder Spaltenvektor
    vec = np.arange(rank_size)
    vshape = np.arange(2)  # Array mit Länge 2, um die Shape des Vektors zu speichern
    dim = 0  # Gesamtdimension des Vektors, Länge des Vektors

    comm = MPI.COMM_WORLD
    size = comm.Get_size()
    rank = comm.Get_rank()

    # Konstruktor
    def __init__(self, array):

        if isinstance(array, np.ndarray):
            form = array.shape
            if len(array) < self.size:
                raise ValueError("ERROR_3: Die Dimension des Vektors ist kleiner als die Anzahl der benutzten Ränge.")

            if len(form) > 2:
                raise ValueError("ERROR_2: Falsche Dimension, kann kein 1 x n oder n x 1 Vektor sein.")
            # Array ist Zeilenvektor:
            if len(form) == 1 or (len(form) == 2 and form[0] == 1):
                self.vshape[0] = 1
                self.vshape[1] = len(array)  # Shape des Vectors
                self.start_idx = int(self.rank * len(array) / self.size)
                self.end_idx = int(len(array) / self.size + self.rank * len(array) / self.size) - 1
                self.rank_size = (
                                         self.end_idx - self.start_idx) + 1  # Größe des Teilvektors auf dem akt. Rang: Differenz zw. Start- und Endindex + 1
                self.vec = array[
                           self.start_idx: self.end_idx + 1]  # Auf jedem Rang werden die Einträge vom Start bis zum Endindex gespeichert
                self.kind = 'row'
                self.dim = len(array)

            # Array ist Spaltenvektor
            if len(form) == 2 and form[1] == 1:
                self.vshape[0] = len(array)
                self.vshape[1] = 1
                self.start_idx = int(self.rank * len(array) / self.size)
                self.end_idx = int(len(array) / self.size + self.rank * len(array) / self.size) - 1
                self.rank_size = (
                                         self.end_idx - self.start_idx) + 1  # Größe des Teilvektors auf dem akt. Rang: Differenz zw. Start- und Endindex + 1
                self.vec = array[
                           self.start_idx: self.end_idx + 1]  # Auf jedem Rang werden die Einträge vom Start bis zum Endindex gespeichert
                self.kind = 'column'
                self.dim = len(array)

        elif isinstance(array, list):
            self.vshape[0] = 1
            self.vshape[1] = len(array)
            self.start_idx = int(self.rank * len(array) / self.size)
            self.end_idx = int(len(array) / self.size + self.rank * len(array) / self.size) - 1
            self.rank_size = (self.end_idx - self.start_idx) + 1
            self.vec = np.array(array[self.start_idx:self.end_idx + 1])
            self.kind = 'row'
            self.dim = len(array)
        else:
            raise ValueError(
                "ERROR_1: Die übergebene Variable ist kein Numpy-Array, Keine Initialisierung der Vector-Klasse möglich.")

    def __add__(self, other):  # Überschreibung der Addition
        if isinstance(self, Vector) and isinstance(other, Vector):
            Add_Vec = Vector(np.arange(self.dim))
            if self.vshape[0] == other.vshape[0] and self.vshape[1] == other.vshape[1]:
                for i in range(0, self.rank_size):
                    Add_Vec.vec[i] = self.vec[i] + other.vec[i]
            else:
                raise ValueError("Die Dimensionen der Vektoren stimmen nicht überein, Addition nicht möglich.")
        elif isinstance(self, Vector) and isinstance(other, (int, float, complex)):
            Add_Vec = Vector(np.arange(self.dim))
            for i in range(0, self.rank_size):
                Add_Vec.vec[i] = self.vec[i] + other
        elif isinstance(self, (int, float, complex)) and isinstance(other, Vector):
            Add_Vec = Vector(np.arange(other.dim))
            for i in range(0, other.rank_size):
                Add_Vec.vec[i] = other.vec[i] + self
        else:
            raise ValueError("Ungeeigneter Datentyp für die Addition mit einem Vektor.")
        return Add_Vec

    def __radd__(self, other):  # Überschreibung der Addition eines Vektors von rechts
        Vector(np.arange(self.dim))
        if isinstance(self, Vector) and isinstance(other, (int, float, complex)):
            Add_Vec = Vector(np.arange(self.dim))
            for i in range(0, self.rank_size):
                Add_Vec.vec[i] = self.vec[i] + other
        else:
            raise ValueError("Ungeeigneter Datentyp für die Addition mit einem Vektor.")
        return Add_Vec

    def __sub__(self, other):  # Überschreibung der Subtraktion
        if isinstance(self, Vector) and isinstance(other, Vector):
            Sub_Vec = Vector(np.arange(self.dim))
            if self.vshape[0] == other.vshape[0] and self.vshape[1] == other.vshape[1]:
                for i in range(0, self.rank_size):
                    Sub_Vec.vec[i] = self.vec[i] - other.vec[i]
            else:
                raise ValueError("Die Dimension der Vektoren stimmen nicht überein, Subtraktion nicht möglich.")
        elif isinstance(self, Vector) and isinstance(other, (int, float, complex)):
            Sub_Vec = Vector(np.arange(self.dim))
            for i in range(0, self.rank_size):
                Sub_Vec.vec[i] = self.vec[i] - other
        elif isinstance(self, (int, float, complex)) and isinstance(other, Vector):
            Sub_Vec = Vector(np.arange(self.dim))
            for i in range(0, other.rank_size):
                Sub_Vec.vec[i] = other.vec[i] - self
        else:
            raise ValueError("Ungeeigneter Datentyp für die Subtraktion mit einem Vektor.")
        return Sub_Vec

    def __rsub__(self, other):  # Subtraktion einer Zahl von einem Vektor
        Sub_Vec = Vector(np.arange(self.dim))
        if isinstance(self, Vector) and isinstance(other, (float, int, complex)):
            for i in range(0, self.rank_size):
                Sub_Vec.vec[i] = self.vec[i] - other
        else:
            raise ValueError("Ungeeigneter Datentyp für die Subtraktion von einem Vektor.")
        return Sub_Vec

    def __mul__(self, other):  # Überschreibung der Multiplikation
        if isinstance(self, Vector) and isinstance(other, Vector):
            Mult_Vec = Vector(np.arange(self.dim))
            if (self.vshape[0] == other.vshape[0] and self.vshape[1] == other.vshape[
                1]):  # Elementweise Multiplikation
                for i in range(0, self.rank_size):
                    Mult_Vec.vec[i] = self.vec[i] * other.vec[i]
            elif self.vshape[1] == other.vshape[0] and self.vshape[0] == 1:  # Inneres Produkt (Skalarprodukt)
                skal_prod = 0
                for i in range(0, self.rank_size):
                    skal_prod = skal_prod + self.vec[i] * other.vec[i]
                return skal_prod
            elif self.vshape[0] == other.vshape[1] and self.vshape[1] == 1:
                raise ValueError("Kann erst implementiert werden, wenn Matrix-Klasse existiert.")
            else:
                raise ValueError("Die Dimensionen der Vektoren stimmen nicht überein, Multiplikation nicht möglich.")
        elif isinstance(self, Vector) and isinstance(other, (int, float, complex)):
            Mult_Vec = Vector(np.arange(self.dim))
            for i in range(0, self.rank_size):
                Mult_Vec.vec[i] = self.vec[i] * other
        elif isinstance(self, (int, float, complex)) and isinstance(other, Vector):
            Mult_Vec = Vector(np.arange(self.dim))
            for i in range(0, other.rank_size):
                Mult_Vec.vec[i] = other.vec[i] * self
        else:
            raise ValueError("Ungeeigneter Datentyp für die Multiplikation mit einem Vektor.")
        return Mult_Vec

    def __rmul__(self, other):  # Rechtsseitige Multiplikation von einer Zahl an einen Vektor
        Mult_Vec = Vector(np.arange(self.dim))
        if isinstance(self, Vector) and isinstance(other, (int, float, complex)):
            for i in range(0, self.rank_size):
                Mult_Vec.vec[i] = self.vec[i] * other
        else:
            raise ValueError("Ungeeigneter Datentyp für die Multiplikation mit einem Vektor.")
        return Mult_Vec

    def __truediv__(self, other):
        Div_Vec = Vector(np.arange(self.dim, dtype=np.double))
        if isinstance(self, Vector) and isinstance(other, Vector):
            if self.vshape[0] == other.vshape[0] and self.vshape[1] == other.vshape[1]:
                for i in range(0, self.rank_size):
                    if other.vec[i] == 0:
                        raise ValueError("Ein Eintrag des Divisor-Vektors ist 0, Divion nicht möglich.")
                    Div_Vec.vec[i] = self.vec[i] / other.vec[i]
            else:
                raise ValueError("Die Dimensionen der Vektoren stimmen nicht überein, Division nicht möglich.")
        elif isinstance(self, (int, float, complex)) and isinstance(other, Vector):
            for i in range(0, other.rank_size):
                if other.vec[i] == 0:
                    raise ValueError("Ein Eintrag des Divisor-Vektors ist 0, Divion nicht möglich.")
                Div_Vec.vec[i] = self / other.vec[i]
        elif isinstance(self, Vector) and isinstance(other, (float, int, complex)):
            if other == 0:
                raise ValueError("Division durch Null ist nicht möglich.")
            else:
                for i in range(0, self.rank_size):
                    Div_Vec.vec[i] = self.vec[i] / other
        else:
            raise ValueError("ERROR 11: Ungeeigneter Datentyp für die Division mit einem Vektor.")
        return Div_Vec

    def __rtruediv__(self, other):
        Div_Vec = Vector(np.arange(self.dim, dtype=np.double))
        if isinstance(self, Vector) and isinstance(other, (float, int, complex)):
            if other == 0:
                raise ValueError("Division durch Null ist nicht möglich.")
            else:
                for i in range(0, self.rank_size):
                    Div_Vec.vec[i] = self.vec[i] / other
        else:
            raise ValueError("ERROR 10: Uneignete Datentyp, um einen Vektor durch diesen zu dividieren")
        return Div_Vec

    def __neg__(self):
        Neg_Vec = -1 * self
        return Neg_Vec

    def shape(self):
        if self.rank == 0:
            return self.vshape

    def T(self):
        Transpose = self
        if self.kind == 'row':
            Transpose.kind = 'column'
        else:
            Transpose.kind = 'row'

        # Tauschen der Dimensionen
        var_shift = self.vshape[0]
        Transpose.vshape[0] = self.vshape[1]
        Transpose.vshape[1] = var_shift
        return Transpose

    def str(self):  # Rückgabe des gesamten Vektors als string
        str_rep = ''

        if self.rank == 0:
            if self.kind == 'row':
                str_rep = '[' + ','.join(map(str, self.vec))
            if self.kind == 'column':
                str_rep = '[' + '\n'.join(map(str, self.vec))
            if self.size > 1:
                self.comm.send(str_rep, dest=self.rank + 1)

        elif self.rank == self.size - 1:
            if self.kind == 'row':
                str_rep = self.comm.recv(source=self.rank - 1) + ',' + ','.join(map(str, self.vec))
            if self.kind == 'column':
                str_rep = self.comm.recv(source=self.rank - 1) + '\n' + '\n'.join(map(str, self.vec))

        else:
            if self.kind == 'row':
                str_rep = self.comm.recv(source=self.rank - 1) + ',' + ','.join(map(str, self.vec))
            if self.kind == 'column':
                str_rep = self.comm.recv(source=self.rank - 1) + '\n' + '\n'.join(map(str, self.vec))
            self.comm.send(str_rep, dest=self.rank + 1)

        str_rep = self.comm.bcast(str_rep, root=self.size - 1)
        if self.rank == 0:
            return str_rep + ']'

    def string(self, limit_entry):  # Gibt den Vektor als String zurück bis zum Eintrag limit_entry
        str_rep = ''

        if limit_entry > self.vec_size:
            raise ValueError("ERROR_4: Die eingegebene Zahl ist größer, als der größte Index des Vectors.")

        # Rank 0
        if self.rank == 0 and limit_entry <= self.end_idx:  # Limit_entry befindet sich im Rang 0
            if self.kind == 'row':
                str_rep = '[' + ','.join(map(str, self.vec[:limit_entry]))
            if self.kind == 'column':
                str_rep = '[' + '\n'.join(map(str, self.vec[:limit_entry]))
            if self.size > 1:
                self.comm.send(str_rep, dest=self.rank + 1)
        if self.rank == 0 and limit_entry > self.end_idx:  # Limit_entry befindet sich nicht im Rang 0
            if self.kind == 'row':
                str_rep = '[' + ','.join(map(str, self.vec))
            if self.kind == 'column':
                str_rep = '[' + '\n'.join(map(str, self.vec))
            if self.size > 1:
                self.comm.send(str_rep, dest=self.rank + 1)

        # Rank im Intervall [1,size-1]
        if (
                0 < self.rank < self.size - 1 and limit_entry <= self.start_idx):  # wenn lim_ent == start_idx, dann wurden bereits alle relevanten Indizes im String gespeichert, da Vector nullinitialisiert ist
            str_rep = self.comm.recv(source=self.rank - 1)
            self.comm.send(str_rep, dest=self.rank + 1)
        if (
                0 < self.rank < self.size - 1 and self.start_idx < limit_entry <= self.end_idx):
            if self.kind == 'row':
                str_rep = self.comm.recv(source=self.rank - 1) + ',' + ','.join(
                    map(str, self.vec[:(limit_entry - self.start_idx)]))
            if self.kind == 'column':
                str_rep = self.comm.recv(source=self.rank - 1) + '\n' + '\n'.join(
                    map(str, self.vec[:(limit_entry - self.start_idx)])) + ']'
            self.comm.send(str_rep, dest=self.rank + 1)
        if 0 < self.rank < self.size - 1 and limit_entry > self.end_idx:
            if self.kind == 'row':
                str_rep = self.comm.recv(source=self.rank - 1) + ',' + ','.join(map(str, self.vec))
            if self.kind == 'column':
                str_rep = self.comm.recv(source=self.rank - 1) + '\n' + '\n'.join(map(str, self.vec))
            self.comm.send(str_rep, dest=self.rank + 1)

        # Rank size-1
        if self.rank == self.size - 1 and limit_entry <= self.start_idx and self.rank > 1:
            str_rep = self.comm.recv(source=self.rank - 1)

        if self.rank == self.size - 1 and limit_entry >= self.start_idx and self.rank > 1:
            if self.kind == 'row':
                str_rep = self.comm.recv(source=self.rank - 1) + ',' + ','.join(
                    map(str, self.vec[:(limit_entry - self.start_idx)]))
            if self.kind == 'column':
                str_rep = self.comm.recv(source=self.rank - 1) + '\n' + '\n'.join(
                    map(str, self.vec[:(limit_entry - self.start_idx)]))

        str_rep = self.comm.bcast(str_rep, root=self.size - 1)
        if self.rank == 0:
            return str_rep + ']'

    def norm(self):  # Berechnung der 2-Norm / euklidischen Norm
        0
        sum_of_squares = 0
        if self.rank == 0:
            for i in range(0, self.rank_size):
                sum_of_squares = sum_of_squares + self.vec[i] ** 2

            if self.size > 1:
                self.comm.send(sum_of_squares, dest=self.rank + 1)

        elif self.rank == self.size - 1:
            sum_of_squares = self.comm.recv(source=self.rank - 1)
            for i in range(0, self.rank_size):
                sum_of_squares = sum_of_squares + self.vec[i] ** 2

        else:
            sum_of_squares = self.comm.recv(source=self.rank - 1)
            for i in range(0, self.rank_size):
                sum_of_squares = sum_of_squares + self.vec[i] ** 2
            self.comm.send(sum_of_squares, dest=self.rank + 1)

        sum_of_squares = self.comm.bcast(sum_of_squares, root=self.size - 1)
        norm = np.sqrt(sum_of_squares)

        return norm

    def normalize(self):  # Normalisierung eines Vectors
        norm = self.norm()
        if norm == 0:
            return self
        normalized_vec = self / norm
        return normalized_vec


# Main-Funktion
x = Vector(np.arange(10))
print(x.str(), x.shape())
print(x.vshape[0], x.vshape[1])
minus = -1 * x
_x = -x
print(_x.str())
print(minus.str())
y = Vector(2 * np.arange(10))
print(y.str())
z = x - y
print(z.str())
ae = x + 5
print(ae.str())
o = x * y
print(o.str())

a = Vector(np.array([[1], [2]]))
b = Vector(np.array([1, 2]))
print(a.shape())
# c = a * b
# print(c.vec)