import math

import numpy
from numpy import linalg


class Matrix:
    """
    This Matrix class represents a real 2D-matrix.
    """
    __data__: list
    __shape__: (int, int)

    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.

        - ``Matrix(list)``: will create a new matrix with the given data in the list and its shape.
        - ``Matrix(numpy.ndarray)``: will create a new matrix with the given data in ndarray and its shape.
        - ``Matrix(list, (int,int))``: will create a new nxm matrix with the given rows and columns and data in list.
        - ``Matrix(list, str, int, int)``: will create a new square matrix of given size and structure of \"diagonal\"
        - ``Matrix(list, str, int)``: will create a new square matrix of given size and structure of either \"unity\", \"diagonal\" or \"tridiagonal\"
        - ``Matrix(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: list | numpy.ndarray
        :type shape: (int, int)
        :type structure: str
        :type model: str
        :type offset: int
        :type n: int

        :rtype: Matrix
        """
        # Case Matrix(str, int)
        if n is not None and model is not None:
            ...  # TODO: what shall one do here?
        # Case: Matrix(list, str, int, int)
        elif n is not None and offset is not None and structure == "diagonal" and data is not None:
            diag = numpy.diag(data * (n - abs(offset)), offset)
            self.__data__ = diag.tolist()
            self.__shape__ = diag.shape
        # Case: Matrix(list, str, int)
        elif n is not None and structure is not None and data is not None:
            if structure == "unity":
                ...  # TODO: what does it mean?
            elif structure == "tridiagonal":
                if len(data) != 3:
                    raise ValueError("If structure is tridiagonal, then the given data must be of length 3")
                tridiag = numpy.diag([data[0]] * (n - 1), -1) + numpy.diag([data[1]] * n, 0) + numpy.diag(
                    [data[2]] * (n - 1), 1)
                self.__data__ = tridiag.tolist()
                self.__shape__ = tridiag.shape
        # Case: Matrix(list, (int,int))
        elif shape is not None and data is not None:
            self.__shape__ = shape
            self.__data__ = numpy.array(data).reshape(shape).tolist()
        # Case: Matrix(numpy.ndarray) or Matrix(list)
        elif data is not None:
            if isinstance(data, numpy.ndarray):
                try:
                    data.shape[1]
                except IndexError:
                    self.__shape__ = (data.shape[0], 1)
                else:
                    self.__shape__ = (data.shape[0], data.shape[1])
            elif isinstance(data, list):
                self.__shape__ = (len(data), len(data[0]))
            self.__data__ = data
        else:
            raise ValueError(
                "Only following signatures are allowed: "
                "(list), (numpy.ndarray), (list, tuple), (list, str, int), (list, str, int, int), (str, int)")

    def get_data(self):
        """
        :return: the data of the matrix as a ``list``
        """
        return self.__data__

    def shape(self):
        """
        :return: the shape of the matrix, which is ``(rows, columns)``
        """
        return self.__shape__

    def transpose(self):
        """
        :return: the transpose of the matrix
        """
        rows = self.__shape__[0]
        cols = self.__shape__[1]

        transposed_data = [[0 for _ in range(rows)] for _ in range(cols)]

        for i in range(rows):
            for j in range(cols):
                transposed_data[j][i] = self.__data__[i][j]

        return Matrix(transposed_data, (cols, rows))

    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 ``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, Matrix):
            data_to_compare = other.__data__
            if self.__shape__ != other.__shape__:
                return False
        elif isinstance(other, list):
            data_to_compare = other
            if self.__shape__[0] != len(other) or self.__shape__[1] != len(other[0]):
                return False
        elif isinstance(other, numpy.ndarray):
            data_to_compare = other.tolist()
        else:
            raise ValueError("Matrix type is not comparable to type of given ``other``")

        for i in range(len(self.__data__)):
            for j in range(len(self.__data__[i])):
                if self.__data__[i][j] != data_to_compare[i][j]:
                    return False
        return True

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

    def __neg__(self):
        rows = range(self.__shape__[0])
        cols = range(self.__shape__[1])
        return Matrix([[-(self.__data__[i][j]) for j in cols] for i in rows], self.__shape__)

    def __add_matrix_internal__(self, other):
        rows = self.__shape__[0]
        cols = self.__shape__[1]
        return [[(self.__data__[i][j] + other.__data__[i][j]) for j in range(cols)] for i in range(rows)]

    def __add_scalar_internal__(self, other):
        rows = self.__shape__[0]
        cols = self.__shape__[1]
        return [[(self.__data__[i][j] + other) for j in range(cols)] for i in range(rows)]

    def __add__(self, other):
        if isinstance(other, Matrix):
            if self.__shape__ != other.__shape__:
                raise ValueError("The shape of the operands must be the same")
            return Matrix(self.__add_matrix_internal__(other), self.__shape__)
        elif isinstance(other, int) or isinstance(other, float):
            return Matrix(self.__add_scalar_internal__(other), self.__shape__)
        else:
            raise ValueError("Only a number or another ``Matrix`` can be added to a ``Matrix``")

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

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

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

    def __truediv_scalar_internal__(self, other):
        rows = self.__shape__[0]
        cols = self.__shape__[1]
        return [[(self.__data__[i][j] / other) for j in range(cols)] for i in range(rows)]

    def __truediv__(self, other):
        if isinstance(other, int) or isinstance(other, float):
            return Matrix(self.__truediv_scalar_internal__(other), self.__shape__)
        else:
            raise ValueError("A ``Matrix`` can only be divided ba a number")

    def __mul_matrix_internal__(self, other):
        rows = self.__shape__[0]
        cols = other.__shape__[1]

        new_data = [[0 for _ in range(rows)] for _ in range(cols)]

        for i in range(rows):
            for k in range(cols):
                new_data[i][k] = sum([self.__data__[i][j] * other.__data__[j][k] for j in range(self.__shape__[1])])
        return new_data

    def __mul_scalar_internal__(self, other):
        cols = range(self.__shape__[1])
        rows = range(self.__shape__[0])
        return [[(self.__data__[i][j] * other) for j in cols] for i in rows]

    def __mul__(self, other):
        if isinstance(other, Matrix):
            if self.__shape__[1] != other.__shape__[0]:
                raise ValueError(
                    "The amount of columns of the first operand must match the amount of rows of the second operand")
            return Matrix(self.__mul_matrix_internal__(other), (self.__shape__[0], other.__shape__[1]))
        elif isinstance(other, int) or isinstance(other, float):
            return Matrix(self.__mul_scalar_internal__(other), self.__shape__)
        else:
            raise ValueError("Only a number or another ``Matrix`` can be multiplied to a ``Matrix``")

    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", "rowsum" or "colsum"

        :return: the norm as a number
        """
        norm = 0

        rows = self.__shape__[0]
        cols = self.__shape__[1]

        if f == "frobenius":
            abs_sum = 0
            for i in range(rows):
                for j in range(cols):
                    abs_sum += abs(self.__data__[i][j])**2
            norm = math.sqrt(abs_sum)
        elif f == "col sum":
            row_sum = [0 for _ in range(cols)]
            for j in range(cols):
                for i in range(rows):
                    row_sum[j] += abs(self.__data__[i][j])
            norm = max(row_sum)
        elif f == "row sum":
            col_sum = [0 for _ in range(rows)]
            for i in range(rows):
                for j in range(cols):
                    col_sum[i] += abs(self.__data__[i][j])
            norm = max(col_sum)
        return norm

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

    def __setitem__(self, key, value):
        manipulated_data = numpy.array(self.__data__)
        manipulated_data[key] = value
        self.__data__ = manipulated_data.tolist()