/*=======================================================================
  ch08_thomas.cpp

  Solve a tridiagonal system C*x = b by the Thomas algorithm.

  Origin:
    C++ translation of ch08_thomas.f90 by Alexander Godunov, based on
    Thomas.f90. The forward elimination and back substitution match
    the Fortran version, including the corrected right-hand-side
    update that uses the elimination multiplier.

  Example:
    Eight equations with diagonal 4, subdiagonal -1, and
    superdiagonal -1. The right-hand side is zero except b(8)=16.
    The matrix is strictly diagonally dominant, so this example
    does not require pivoting.

  The program prints the three stored diagonals, the solution,
  and the maximum residual computed from the original system.
=======================================================================*/

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

// In each row: subdiagonal, diagonal, superdiagonal.
using Tridiagonal = std::vector<std::array<double, 3>>;

int thomas(Tridiagonal& c, std::vector<double>& b, std::vector<double>& x);

int main()
{
    // c[0][0] and c[n-1][2] are unused, as in the Fortran program.
    const Tridiagonal c = {
        {{ 0.0, 4.0, -1.0}},
        {{-1.0, 4.0, -1.0}},
        {{-1.0, 4.0, -1.0}},
        {{-1.0, 4.0, -1.0}},
        {{-1.0, 4.0, -1.0}},
        {{-1.0, 4.0, -1.0}},
        {{-1.0, 4.0, -1.0}},
        {{-1.0, 4.0,  0.0}}
    };
    const std::vector<double> b = {0.0, 0.0, 0.0, 0.0,
                                   0.0, 0.0, 0.0, 16.0};
    const int n = static_cast<int>(c.size());

    std::cout << "Thomas algorithm for a tridiagonal system\n"
              << "\nSubdiagonal, diagonal, superdiagonal, and b:\n"
              << std::fixed << std::setprecision(6);
    for (int i = 0; i < n; ++i) {
        for (int j = 0; j < 3; ++j)
            std::cout << std::setw(12) << c[i][j];
        std::cout << std::setw(12) << b[i] << '\n';
    }

    // The Thomas routine changes the diagonal and right-hand side.
    // Keep the originals to check the final solution independently.
    Tridiagonal c_work = c;
    std::vector<double> b_work = b;
    std::vector<double> x;
    const int info = thomas(c_work, b_work, x);

    if (info != 0) {
        if (info < 0)
            std::cerr << "Invalid tridiagonal matrix or right-hand-side size.\n";
        else
            std::cerr << "Zero pivot at row " << info << ".\n"
                      << "The Thomas algorithm cannot solve this system "
                      << "without pivoting.\n";
        return 1;
    }

    std::cout << "\nSolution x:\n";
    for (int i = 0; i < n; ++i)
        std::cout << "x(" << i + 1 << ") = " << std::setw(12) << x[i] << '\n';

    // Each original equation contains at most three terms.
    double max_residual = 0.0;
    for (int i = 0; i < n; ++i) {
        double row_residual = c[i][1] * x[i] - b[i];
        if (i > 0)
            row_residual += c[i][0] * x[i-1];
        if (i < n-1)
            row_residual += c[i][2] * x[i+1];
        max_residual = std::max(max_residual, std::abs(row_residual));
    }
    std::cout << "\nMaximum residual |C*x-b| = "
              << std::scientific << std::setprecision(4)
              << max_residual << '\n';
    return 0;
}

int thomas(Tridiagonal& c, std::vector<double>& b, std::vector<double>& x)
/*=======================================================================
  Solve a tridiagonal system by elimination without row pivoting.

  Input/output:
    c  - three diagonals by row; c[0][0] and c[n-1][2] are unused
         the diagonal c[i][1] is overwritten by its eliminated value
    b  - right-hand side, overwritten during elimination

  Output:
    x  - solution when the return value is zero
  Return value:
    0  - success
   -1  - empty system or mismatched right-hand-side size
    i  - zero pivot at row i, numbered from 1

  The elimination multiplier is c[i][0]/c[i-1][1]. It updates both
  the diagonal and the right-hand side. Storage and work grow as O(n).
  No pivoting is performed; a nonsingular system can still have a
  zero pivot or produce an inaccurate result without pivoting.
=======================================================================*/
{
    const int n = static_cast<int>(c.size());
    if (n == 0 || static_cast<int>(b.size()) != n)
        return -1;

    x.assign(n, 0.0);
    if (c[0][1] == 0.0)
        return 1;

    // Eliminate the subdiagonal, using the multiplier also for b[i].
    for (int i = 1; i < n; ++i) {
        const double coeff = c[i][0] / c[i-1][1];
        c[i][1] -= coeff * c[i-1][2];
        b[i] -= coeff * b[i-1];
        if (c[i][1] == 0.0)
            return i + 1;
    }

    // Solve the upper triangular system from bottom to top.
    x[n-1] = b[n-1] / c[n-1][1];
    for (int i = n-2; i >= 0; --i)
        x[i] = (b[i] - c[i][2] * x[i+1]) / c[i][1];
    return 0;
}