/*=======================================================================
  ch08_sor.cpp

  Solve a system of linear equations A*x = b by successive
  over-relaxation (SOR).

  Origin:
    C++ translation of ch08_sor.f90 by Alexander Godunov, based on
    GS_sor.f90. The sequential SOR update and the completed-sweep
    residual check match the Fortran program.

  Example:
    A = [ 8  2  4; 2  6  1; 1  1  8 ], b = [ 2, 6, 4 ].
    Exact solution: x = [ -0.2, 1, 0.4 ].

  The program reports the iteration count and the final maximum
  residual |A*x-b|. Change omega and the initial guess to see how
  they affect convergence. SOR does not converge for every matrix
  or every choice of omega.
=======================================================================*/

#include <algorithm>
#include <cmath>
#include <iomanip>
#include <iostream>
#include <limits>
#include <vector>

// Each inner vector holds one row of the coefficient matrix.
using Matrix = std::vector<std::vector<double>>;

int sor(const Matrix& a, const std::vector<double>& b,
        std::vector<double>& x, double omega, double eps, int max_iter,
        int& iterations, double& residual);
double maximum_residual(const Matrix& a, const std::vector<double>& b,
                        const std::vector<double>& x);

int main()
{
    const Matrix a = {
        {8.0, 2.0, 4.0},
        {2.0, 6.0, 1.0},
        {1.0, 1.0, 8.0}
    };
    const std::vector<double> b = {2.0, 6.0, 4.0};
    const int n = static_cast<int>(a.size());
    const int max_iter = 100;

    // Initial estimate, relaxation parameter, and residual tolerance.
    std::vector<double> x(n, 0.0);
    const double omega = 1.2;
    const double eps = 1.0e-8;

    std::cout << "Successive over-relaxation (SOR)\n"
              << "\nMatrix A and right-hand side b:\n"
              << std::fixed << std::setprecision(5);
    for (int i = 0; i < n; ++i) {
        for (int j = 0; j < n; ++j)
            std::cout << std::setw(12) << a[i][j];
        std::cout << "   " << std::setw(12) << b[i] << '\n';
    }

    std::cout << "\nInitial guess x:" << std::setprecision(6);
    for (double value : x)
        std::cout << std::setw(12) << value;
    std::cout << '\n' << "Relaxation parameter omega = "
              << std::setprecision(4) << omega << '\n'
              << "Residual tolerance         = "
              << std::scientific << std::setprecision(4) << eps << '\n';

    int iterations = 0;
    double residual = 0.0;
    const int info = sor(a, b, x, omega, eps, max_iter,
                         iterations, residual);

    if (info == 1) {
        std::cerr << "\nInvalid input: check dimensions, omega, tolerance, "
                  << "iteration limit, and diagonal entries.\n";
        return 1;
    }

    std::cout << "\nIterations = " << iterations << "\nSolution x:\n"
              << std::fixed << std::setprecision(6);
    for (int i = 0; i < n; ++i)
        std::cout << "x(" << i + 1 << ") = " << std::setw(12) << x[i] << '\n';
    std::cout << "Maximum residual |A*x-b| = "
              << std::scientific << std::setprecision(4) << residual << '\n';

    if (info == 2) {
        std::cerr << "Iteration limit reached before the tolerance was met.\n";
        return 1;
    }
    return 0;
}

int sor(const Matrix& a, const std::vector<double>& b,
        std::vector<double>& x, double omega, double eps, int max_iter,
        int& iterations, double& residual)
/*=======================================================================
  Solve A*x = b by sequential successive over-relaxation sweeps.

  Input:
    a         - square coefficient matrix, left unchanged
    b         - right-hand side, left unchanged
    x         - initial guess, overwritten by the final estimate
    omega     - relaxation parameter, 0 < omega < 2
    eps       - requested maximum absolute residual (positive)
    max_iter  - maximum number of complete sweeps (positive)

  Output:
    iterations - number of complete sweeps performed
    residual   - max_i |sum_j a(i,j)*x(j) - b(i)|
  Return value:
    0 - residual <= eps; 1 - invalid input;
    2 - iteration limit reached without convergence

  For row i, compute delta = b(i) - sum_j a(i,j)*x(j). Values x(j)
  with j < i have been updated during this sweep; values with j > i
  are from the previous sweep. The update x(i) += omega*delta/a(i,i)
  is the original GS_sor.f90 step. For omega = 1 it is Gauss-Seidel.

  Check the residual after a complete sweep. A small absolute residual
  alone may be misleading when A or b has a very different scale or
  when the matrix is poorly conditioned.
=======================================================================*/
{
    iterations = 0;
    residual = std::numeric_limits<double>::max();
    const int n = static_cast<int>(a.size());

    // A zero diagonal makes an update undefined. Diagonal dominance
    // is useful to check but is not required as an input condition.
    if (n == 0 || static_cast<int>(b.size()) != n ||
        static_cast<int>(x.size()) != n || max_iter < 1 || eps <= 0.0 ||
        omega <= 0.0 || omega >= 2.0)
        return 1;
    for (const auto& row : a) {
        if (static_cast<int>(row.size()) != n)
            return 1;
    }
    for (int i = 0; i < n; ++i) {
        if (a[i][i] == 0.0)
            return 1;
    }

    for (int i = 0; i < n; ++i) {
        double off_diagonal = 0.0;
        for (int j = 0; j < n; ++j) {
            if (j != i)
                off_diagonal += std::abs(a[i][j]);
        }
        if (std::abs(a[i][i]) <= off_diagonal) {
            std::cout << "Warning: A is not strictly diagonally dominant.\n";
            break;
        }
    }

    // The supplied guess may already solve the equations.
    residual = maximum_residual(a, b, x);
    if (residual <= eps)
        return 0;

    // Update x(i) immediately, so later rows use its new value.
    for (int k = 1; k <= max_iter; ++k) {
        for (int i = 0; i < n; ++i) {
            double delta = b[i];
            for (int j = 0; j < n; ++j)
                delta -= a[i][j] * x[j];
            x[i] += omega * delta / a[i][i];
        }

        iterations = k;
        residual = maximum_residual(a, b, x);
        if (residual <= eps)
            return 0;
    }

    return 2;
}

double maximum_residual(const Matrix& a, const std::vector<double>& b,
                        const std::vector<double>& x)
/*-----------------------------------------------------------------------
  Evaluate the residual of the completed iterate using explicit loops.
  A and b are the original, unchanged system.
-----------------------------------------------------------------------*/
{
    const int n = static_cast<int>(a.size());
    double residual = 0.0;
    for (int i = 0; i < n; ++i) {
        double row_sum = -b[i];
        for (int j = 0; j < n; ++j)
            row_sum += a[i][j] * x[j];
        residual = std::max(residual, std::abs(row_sum));
    }
    return residual;
}