import numpy as np

def jacobi(A, b, x0=None, tol=1e-8, max_iter=100):
    """
    Método de Jacobi para resolver Ax = b.

    Parámetros
    ----------
    A : array_like
        Matriz de coeficientes.
    b : array_like
        Vector de términos independientes.
    x0 : array_like, opcional
        Aproximación inicial. Si no se da, se usa el vector cero.
    tol : float, opcional
        Tolerancia para el criterio de parada.
    max_iter : int, opcional
        Número máximo de iteraciones.

    Retorna
    -------
    x : ndarray
        Aproximación de la solución.
    historial : list
        Lista con las aproximaciones de cada iteración.
    """

    A = np.array(A, dtype=float)
    b = np.array(b, dtype=float)

    n = len(b)

    if x0 is None:
        x = np.zeros(n, dtype=float)
    else:
        x = np.array(x0, dtype=float)

    # Verificación básica
    if A.shape[0] != A.shape[1]:
        raise ValueError("La matriz A debe ser cuadrada.")
    if A.shape[0] != n:
        raise ValueError("Las dimensiones de A y b no son compatibles.")

    # Extraer la diagonal
    D = np.diag(np.diag(A))

    # Verificar que no haya ceros en la diagonal
    if np.any(np.diag(D) == 0):
        raise ValueError("La matriz tiene ceros en la diagonal; Jacobi no se puede aplicar directamente.")

    # Construcción de L + U según A = D - L - U
    LU = D - A

    # Matriz iterativa T y vector c
    D_inv = np.linalg.inv(D)
    T = D_inv @ LU
    c = D_inv @ b

    historial = [x.copy()]

    print("Iteración 0:", x)

    for k in range(1, max_iter + 1):
        x_new = T @ x + c
        historial.append(x_new.copy())

        # Error relativo infinito
        norma = np.linalg.norm(x_new, ord=np.inf)
        if norma == 0:
            error = np.linalg.norm(x_new - x, ord=np.inf)
        else:
            error = np.linalg.norm(x_new - x, ord=np.inf) / norma

        print(f"Iteración {k}: {x_new}  Error = {error:.2e}")

        if error < tol:
            return x_new, historial

        x = x_new

    return x, historial


# =========================
# Ejemplo
# =========================
A = np.array([
    [4, -1,  1],
    [4, -8, 1],
    [-2, 1, 5]
], dtype=float)

b = np.array([1, 2, 2], dtype=float)

x_aprox, historial = jacobi(A, b, x0=[0, 0, 0], tol=1e-6, max_iter=50)

print("\nSolución aproximada:")
print(x_aprox)


# =========================
# Verificación
# =========================

Ax = A @ x_aprox
print("\nVerificación Ax:")
print(Ax)

print("\nVector b:")
print(b)

# Residuo
residuo = b - Ax

print("\nResiduo (b - Ax):")
print(residuo)

print("\nNorma del residuo (infinito):")
print(np.linalg.norm(residuo, ord=np.inf))