#!/usr/bin/env python3.10

"""Compare Jacobi and Gauss-Seidel for diagonally dominant linear systems."""

import numpy as np


def generate_diagonally_dominant_matrix(matrix_size, rng):
    """
    Generate a column-wise strictly diagonally dominant matrix.

    Parameters
    ----------
    matrix_size : int
        Matrix order.
    rng : numpy.random.Generator
        Random number generator.
    """
    matrix = rng.uniform(-1.0, 1.0, size=(matrix_size, matrix_size))
    for column_index in range(matrix_size):
        off_diagonal_sum = np.sum(np.abs(matrix[:, column_index])) - abs(matrix[column_index, column_index])
        matrix[column_index, column_index] = off_diagonal_sum + 1.0 + rng.uniform(0.0, 1.0)
    return matrix


def jacobi_solver(system_matrix, right_hand_side, initial_guess, tolerance, max_iterations):
    """Solve a linear system with the Jacobi method."""
    diagonal = np.diag(system_matrix)
    remainder = system_matrix - np.diag(diagonal)
    current_solution = initial_guess.copy()

    for iteration in range(1, max_iterations + 1):
        next_solution = (right_hand_side - remainder @ current_solution) / diagonal
        if np.linalg.norm(next_solution - current_solution, ord=np.inf) <= tolerance:
            return next_solution, iteration
        current_solution = next_solution

    return current_solution, max_iterations


def gauss_seidel_solver(system_matrix, right_hand_side, initial_guess, tolerance, max_iterations):
    """Solve a linear system with the Gauss-Seidel method."""
    current_solution = initial_guess.copy()
    matrix_size = system_matrix.shape[0]

    for iteration in range(1, max_iterations + 1):
        next_solution = current_solution.copy()
        for row_index in range(matrix_size):
            left_sum = np.dot(system_matrix[row_index, :row_index], next_solution[:row_index])
            right_sum = np.dot(system_matrix[row_index, row_index + 1 :], current_solution[row_index + 1 :])
            next_solution[row_index] = (
                right_hand_side[row_index] - left_sum - right_sum
            ) / system_matrix[row_index, row_index]

        if np.linalg.norm(next_solution - current_solution, ord=np.inf) <= tolerance:
            return next_solution, iteration
        current_solution = next_solution

    return current_solution, max_iterations


def main():
    """Run the numerical experiment described in the project statement."""
    rng = np.random.default_rng(42)
    tolerance = 1.0e-10
    max_iterations = 10000

    for matrix_size in range(10, 101, 10):
        jacobi_iterations = []
        gauss_seidel_iterations = []

        for _ in range(100):
            system_matrix = generate_diagonally_dominant_matrix(matrix_size, rng)
            exact_solution = rng.uniform(-1.0, 1.0, size=matrix_size)
            right_hand_side = system_matrix @ exact_solution
            initial_guess = right_hand_side.copy()

            _, jacobi_count = jacobi_solver(
                system_matrix, right_hand_side, initial_guess, tolerance, max_iterations
            )
            _, gauss_seidel_count = gauss_seidel_solver(
                system_matrix, right_hand_side, initial_guess, tolerance, max_iterations
            )

            jacobi_iterations.append(jacobi_count)
            gauss_seidel_iterations.append(gauss_seidel_count)

        print(
            f"n = {matrix_size:3d} | "
            f"Jacobi mean iterations = {np.mean(jacobi_iterations):8.3f} | "
            f"Gauss-Seidel mean iterations = {np.mean(gauss_seidel_iterations):8.3f}"
        )


if __name__ == "__main__":
    main()