import numpy as np
import matplotlib.pyplot as plt
import time


def LU_decomposition(A):
    n = len(A[0])
    L = np.zeros([n, n])
    U = np.zeros([n, n])
    for i in range(n):
        L[i][i] = 1
        if i == 0:
            U[0][0] = A[0][0]
            for j in range(1, n):
                U[0][j] = A[0][j]
                L[j][0] = A[j][0] / U[0][0]
        else:
            for j in range(i, n):  # U
                temp = 0
                for k in range(0, i):
                    temp = temp + L[i][k] * U[k][j]
                U[i][j] = A[i][j] - temp
            for j in range(i + 1, n):  # L
                temp = 0
                for k in range(0, i):
                    temp = temp + L[j][k] * U[k][i]
                L[j][i] = (A[j][i] - temp) / U[i][i]
    return L, U


# 生成随机矩阵 A, b, 范围[-10, 10]
def randomAb(m):
    A = np.random.random([m, m]) * 20 - 10
    dia = np.random.random(m) * 10
    for i in range(len(dia)):
        A[i, i] = dia[i]
    return A, np.random.randint(0, 10, [m, 1])


# 生成稀疏矩阵A, b
def sparseMatrix(m):
    A, b = randomAb(m)
    for i in range(m):
        for j in range(m):
            if i != j:
                a = A[i][j]
                if abs(a) < 9:
                    A[i][j] = 0
                elif a > 0:
                    A[i][j] = 10 * (a - 9)
                else:
                    A[i][j] = 10 * (9 + a)

    return A, b


# 生成病态矩阵A, b
def illMatrix(m):
    A, b = randomAb(m)
    return A, b


class Question1:
    """
    求解 Ax=b
    """

    def __init__(self):
        pass

    def solver1(self, A, b):
        """
        LU 分解 求解器
        """
        L, U = LU_decomposition(A)
        # LY=b
        n = len(A)
        y = np.zeros((n, 1))
        for i in range(len(A)):
            t = 0
            for j in range(i):
                t += L[i][j] * y[j][0]
            y[i][0] = b[i][0] - t
        X = np.zeros((n, 1))
        for i in range(len(A) - 1, -1, -1):
            t = 0
            for j in range(i + 1, len(A)):
                t += U[i][j] * X[j][0]
            t = y[i][0] - t
            if t != 0 and U[i][i] == 0:
                return 0
            X[i] = t / U[i][i]
        return X

    def solver2(self, A, b):
        """
        Jacobi 求解器
        """
        x = np.zeros(b.shape)
        Dv = np.diag(A)
        D = np.zeros(A.shape)
        for i in range(len(Dv)):
            D[i, i] = Dv[i]
        R = A - np.diagflat(Dv)

        # Iterate for N times
        print(D)
        D = np.linalg.inv(D)
        print(D)
        while 1:
            x1 = D @ (b - R @ x)
            if np.max(x1 - x) < 1e-6:
                break
            x = x1
        return x
        # return Jacobi(np.zeros(b.shape), A, b)

    def solver3(self, A, b):
        """
        inv(A) * b
        """
        return np.dot(np.linalg.inv(A), b)

    def solver4(self, A, b):
        """
        默认求解器
        """
        return np.linalg.solve(A, b)

    def RMSE(self, solver, n=8, randFunc=randomAb):
        ## 计算方差
        s = time.time()
        A, b = randFunc(n)
        X = (A.dot(solver(A, b)) - b) ** 2
        for i in range(1000):
            A, b = randFunc(n)
            X = X + (A.dot(solver(A, b)) - b) ** 2
        return np.max(X), time.time() - s

    def show(self):
        n = 12
        N = [2 ** i for i in range(n)]
        Y = [[0 for i in range(n)] for _ in range(4)]
        Z = [[0 for i in range(n)] for _ in range(4)]

        randomFuns = [randomAb, sparseMatrix, illMatrix]
        for r in range(3):
            plt.subplot(3, 2, 2*r + 1)
            for i in range(n):
                print("size: %s" % N[i])
                print("LU")
                Y[0][i], Z[0][i] = self.RMSE(self.solver1, N[i])
                print("jacobi")
                Y[1][i], Z[1][i] = self.RMSE(self.solver3, N[i])
                print("inverse")
                Y[2][i], Z[2][i] = self.RMSE(self.solver3, N[i])
                print("default\n")
                Y[3][i], Z[3][i] = self.RMSE(self.solver4, N[i])

            # N = range(12)
            plt.plot(N, Y[0], label="LU")
            plt.plot(N, Y[1], label="Jacobi")
            plt.plot(N, Y[2], label="inverse")
            plt.plot(N, Y[3], label="default solver")
            # plt.xticks([0, 10, 100, 1000], [0, 10, 100, 1000])
            plt.title('Accuracy')
            plt.yscale('symlog')
            plt.xscale('symlog')
            # plt.legend(loc='lower right')
            plt.subplot(3, 2, 2 * r + 2)
            plt.plot(N, Z[0], label="LU")
            plt.plot(N, Z[1], label="Jacobi")
            plt.plot(N, Z[2], label="inverse")
            plt.plot(N, Z[3], label="default solver")
            plt.xscale('symlog')
            plt.yscale('symlog')
            plt.title('time cost')
            # plt.legend(loc='lower right')
        plt.show()


if __name__ == "__main__":
    q = Question1()
    q.show()