import numpy as np


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


def diag_strict_dominante(A) -> bool:
    """
    Pour chaque ligne, il faut que la somme des coefs non-diagonaux soit inférieure au coef diagonal
    TOUT EST EN VALEUR ABSOLUE.
    :param A:
    :return:
    """
    A = np.array(A)
    n = len(A)
    for i in range(n):
        diag = abs(A[i,i])
        row_sum = np.sum(np.abs(A[i,:])) - diag # On fait la somme de toute la ligne puis on soustrait le coef diagonal
        if diag <= row_sum:
            return False
    return True


def jacobi(A, b, epsilon=1e-6, max_iter=100_000):
    if not diag_strict_dominante(A):
        raise ValueError("A doit être à diagonale strictement dominante")

    x = np.zeros_like(b, dtype=float)
    L = np.tril(A, k=-1)     # Partie triangulaire inférieure (sans diagonale)
    U = np.triu(A, k=1)      # Partie triangulaire supérieure (sans diagonale)

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

    return x_new


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

    D = np.diag(np.diag(A))
    L = np.tril(A, k=-1)
    U = np.triu(A, k=1)
    epsilon = 1e-6
    max_iter = 100_000

    # Pré-calcul de l'inverse de (D - L) pour éviter de le recalculer à chaque itération
    inv_D_minus_L = np.linalg.inv(D - L)

    x = x0

    for _ in range(max_iter):
        x_new = inv_D_minus_L @ (-U @ x) + inv_D_minus_L @ b

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

    raise RuntimeError('Gauss-Seidel')


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 = np.diag(A)
    L = np.tril(A)
    U = np.triu(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')