import numpy as np


# L'inverse d'une diagonale
# D_inv = np.diag(1 / np.diag(A))


def to_D(A: np.array) -> np.array:
    D = np.zeros_like(A)
    for i in range(len(A)):
        D[i, i] = A[i, i]
    return D


def to_L(A: np.array) -> np.array:
    L = np.zeros_like(A)
    for i in range(len(A)):
        for j in range(len(A)):
            if i < j:
                L[i, j] = A[i, j]
    return L


def to_U(A: np.array) -> np.array:
    U = np.zeros_like(A)
    for i in range(len(A)):
        for j in range(len(A)):
            if i > j:
                U[i, j] = A[i, j]
    return U


def diag_strict_dominante(A) -> bool:
    diag_sum = 0
    for i in range(len(A)):
        diag_sum += A[i, i]
    other_sum = 0
    for i in range(len(A)):
        for j in range(len(A)):
            if i != j:
                other_sum += A[i, j]
    return diag_sum > other_sum


def jacobi(A, b):
    if not diag_strict_dominante(A):
        raise Exception('A doit être à diagnonale strictement dominante')

    L = to_L(A)
    U = to_U(A)
    x0 = np.array([0,0,0])
    epsilon = 1e-6
    max_iter = 100_000

    x = x0
    for k in range(max_iter):
        x_new = np.diag(1 / np.diag(A)) @ ((L + U) @ x) + np.diag(1 / np.diag(A)) @ b
        if np.linalg.norm(x_new - x, ord=2) < epsilon or np.linalg.norm(b - A @ x_new, ord=2) < epsilon:
            break
        x = x_new

    return x


def gauss_seidel(A, b):
    x0 = np.array([0, 0, 0])

    D = to_D(A)
    L = to_L(A)
    U = to_U(A)
    epsilon = 1e-6
    done = False

    x = x0

    while not done:
        x_new = np.linalg.inv(D - L) @ U @ x + np.linalg.inv(D - L) @ b

        done: bool = np.linalg.norm(x_new - x, ord=2) < epsilon
        x = x_new

    return x


def relaxation(A, b, omega=1.0, epsilon=1e-6, max_iter=100_000):
    D = np.diag(np.diag(A))
    L = np.tril(A, k=-1)
    U = np.triu(A, k=1)
    x = np.zeros_like(b, dtype=float)

    # Pré-calculer (D - ωL)^(-1) une seule fois
    inv_D_omega_L = np.linalg.inv(D - omega * L)

    if omega == 1:
        return gauss_seidel(A, b)

    for _ in range(max_iter):
        x_new = inv_D_omega_L @ ((1 - omega) * D @ x + omega * (U @ x + b))
        if np.linalg.norm(x_new - x, ord=2) < epsilon:
            return x_new
        x = x_new

    raise RuntimeError("La méthode de relaxation n'a pas convergé.")



if __name__ == '__main__':
    A = np.array([
        [8,4,1],
        [1,6,-5],
        [1,-2,-6]
    ])

    b = np.array([1,0,0])

    D = to_D(A)
    L = to_L(A)
    U = to_U(A)

    res_jacobi = jacobi(A, b)
    print(res_jacobi)

    res_gauss = gauss_seidel(A, b)
    print(res_gauss)

    # je la commente car omega = 1 utilise la ^m fonction que gauss_seidel
    #res_relaxation_1 = relaxation(A, b, 1)
    #print(res_relaxation_1)

    res_relaxation_less_1 = relaxation(A, b, 0)
    print(res_relaxation_less_1)

    res_relaxation_2 = relaxation(A, b, 2)
    print(res_relaxation_2)

    print('fini')