import numpy as np

def calculate_error(diff, norm=0):
    if norm == 0:  # Infinity norm
        return np.max(np.abs(diff))
    elif norm == 1:
        return np.sum(np.abs(diff))
    else:
        return np.sum(np.abs(diff) ** norm) ** (1 / norm)


def jacobi_iterative_method(A, b, x0=None, max_iterations=1000, termination_error=1e-4, error_norm=0):
    x = x0 if x0 is not None else np.zeros_like(b)  # set initial estimate
    x = x.astype(float)  # Cast the array to float type

    if not (np.size(A, 0) == np.size(A, 1) == np.size(b, 0) == np.size(x, 0)) or np.size(x, 1) != 1:
        raise ValueError("The size of the input matrix does not match")

    current_error = calculate_error(A.dot(x) - b, error_norm)
    iteration = 0

    if current_error <= termination_error:
        return x, current_error, iteration

    D = np.diag(np.diag(A))  # Containing the diagonal component of A
    R = A - D  # Containing component of A except diagonal component
    invD = np.linalg.inv(D)  # Inverse matrix of D
    T = -invD.dot(R)
    C = invD.dot(b)

    # Core iterates
    while iteration < max_iterations:
        iteration = iteration + 1
        x = T.dot(x) + C
        current_error = calculate_error(A.dot(x) - b, error_norm)
        if current_error < termination_error:
            break

    print("Jacobi iteration:")
    print(f"A =\n{A}\nb =\n{b}\n")
    print("Solution:")
    print(f"x =\n{x}\n")
    print(f"current_error = {current_error}\niteration = {iteration}\n")
    return x, current_error, iteration


def gauss_seidel_iterative_method(A, b, x0=None, max_iterations=1000, termination_error=1e-4, error_norm=0):
    x = x0 if x0 is not None else np.zeros_like(b)  # set initial estimate
    x = x.astype(float)  # Cast the array to float type

    if not (np.size(A, 0) == np.size(A, 1) == np.size(b, 0) == np.size(x, 0)) or np.size(x, 1) != 1:
        raise ValueError("The size of the input matrix does not match")

    current_error = calculate_error(A.dot(x) - b, error_norm)
    iteration = 0
    collum = np.size(b, 0)

    if current_error <= termination_error:
        return x, current_error, iteration

    D = np.diag(np.diag(A))  # Containing the diagonal component of A
    R = A - D  # Containing component of A except diagonal component
    invD = np.linalg.inv(D)  # Inverse matrix of D
    T = -invD.dot(R)
    C = invD.dot(b)

    # Core iterates
    while iteration < max_iterations:
        iteration = iteration + 1
        for i in range(collum):
            x[i] = T[i].dot(x) + C[i]
        # x = T.dot(x) + C
        current_error = calculate_error(A.dot(x) - b, error_norm)
        if current_error < termination_error:
            break

    print("Gauss-Seidel iteration:")
    print(f"A =\n{A}\nb =\n{b}\n")
    print("Solution:")
    print(f"x =\n{x}\n")
    print(f"current_error = {current_error}\niteration = {iteration}\n")
    return x, current_error, iteration