#ifndef P02_SOLVER_H
#define P02_SOLVER_H

#include "matrix.h"

template <typename T>
double residuum(const Matrix<T>& M, const Matrix<T>& b, const Matrix<T>& x)
{
    return norm(M*x - b);
}

class Solver {
public:
    virtual Matrix<double> solve(const Matrix<double>& M, const Matrix<double>& b) = 0;
};

class LUSolver : public Solver {
public: 
    static std::pair<Matrix<double>, Matrix<double>> decompose(const Matrix<double>& M);
    Matrix<double> solve(const Matrix<double>& L, const Matrix<double>& U, const Matrix<double>& b);
    Matrix<double> solve(const Matrix<double>& M, const Matrix<double>& b) override;
};

class IterativeSolver : public Solver {
protected:
    double error;
    size_t _iterations;
    std::ostream* csv;

public:
    IterativeSolver();
    IterativeSolver(double error);   

    virtual Matrix<double> iteration(const Matrix<double>& M, const Matrix<double>& b, const Matrix<double>& old) = 0;
    virtual Matrix<double> solve(const Matrix<double>& M, const Matrix<double>& b) override;

    size_t iterations() const;
    void output(std::ostream* csv);
};

class JacobiSolver : public IterativeSolver {
public:  
    using IterativeSolver::IterativeSolver;
    virtual Matrix<double> iteration(const Matrix<double>& M, const Matrix<double>& b, const Matrix<double>& old) override;
};

class GaussSolver : public IterativeSolver {
public:  
    using IterativeSolver::IterativeSolver;
    virtual Matrix<double> iteration(const Matrix<double>& M, const Matrix<double>& b, const Matrix<double>& old) override;
};

Matrix<double> IterativeSolver::solve(const Matrix<double>& M, const Matrix<double>& b)
{
    assert(M.rows() == M.cols());
    assert(b.rows() == M.cols());

    size_t n = M.cols();

    Matrix<double> x(n, 1);
    for (size_t i = 0; i < n; ++i) x(i, 0) = 1.;

    _iterations = 0;
    double r;
    if (csv) *csv << "iter,residuum" << std::endl;
    while((r = residuum(M, b, x)) > error) {
        x = iteration(M, b, x);
        if (csv) *csv << iterations() << "," << r << std::endl; 

        _iterations ++;
        if (_iterations >= 100) {
            std::cerr << "Giving up after " << _iterations << " iteration, residuum: " << r << std::endl;
            return x;
        }
    }

    return x;
}

void IterativeSolver::output(std::ostream* stream)
{
    this->csv = stream;
}

size_t IterativeSolver::iterations() const
{
    return _iterations;
}

IterativeSolver::IterativeSolver() : IterativeSolver(1e-9) {}
IterativeSolver::IterativeSolver(double error) : error(error), csv(nullptr) {}

Matrix<double> JacobiSolver::iteration(const Matrix<double>& M, const Matrix<double>& b, const Matrix<double>& old)
{
    Matrix<double> x(b.rows(), 1);
    
    for (size_t i = 0; i < b.rows(); ++i) {
        double t = b(i, 0);
        for (size_t j = 0; j < b.rows(); ++j) {
            if (j == i) continue;
            t -= M(i, j)*old(j, 0);
        }
        
        x(i, 0) = t / M(i,i);
    }

    return x;
}

Matrix<double> GaussSolver::iteration(const Matrix<double>& M, const Matrix<double>& b, const Matrix<double>& old)
{
    Matrix<double> x(b.rows(), 1);
    
    for (size_t i = 0; i < b.rows(); ++i) {
        double t = b(i, 0);

        for (size_t j = 0; j < i; ++j)
            t -= M(i, j)*x(j, 0);
        for (size_t j = i+1; j < b.rows(); ++j)
            t -= M(i, j)*old(j, 0);

        x(i, 0) = t / M(i,i);
    }

    return x;
}

Matrix<double> LUSolver::solve(const Matrix<double>& M, const Matrix<double>& b)
{
    auto LU = LUSolver::decompose(M);
    return this->solve(std::get<0>(LU), std::get<1>(LU), b);
}

Matrix<double> LUSolver::solve(const Matrix<double>& L, const Matrix<double>& U, const Matrix<double>& b)
{
    Matrix<double> x(b);
    const size_t n = x.rows();

    for (size_t i = 0; i < n; i++) {
        for (size_t j = 0; j < i; j++) {
            x(i, 0) -= L(i,j)*x(j, 0);
        }
    }

    for (size_t i = n; i != 0; i--) {
        for (size_t j = n-1; j > i-1; j--) {
            x(i-1, 0) -= U(i-1,j)*x(j, 0);
        }
        x(i-1, 0) /= U(i-1,i-1);
    }

    return x;
}

std::pair<Matrix<double>, Matrix<double>> LUSolver::decompose(const Matrix<double>& M)
{
    const size_t n = M.rows();

    Matrix<double> U(M);
    Matrix<double> L = Matrix<double>::diag(n, 1);

    for (size_t i = 0; i < n - 1; ++i) {
        for (size_t j = i+1; j < n; ++j) {
            L(j, i) = U(j, i) / U(i, i);
            for (size_t k = i; k < n; k++) {
                U(j, k) -= L(j, i)*U(i, k);
            }
        }
    }

    return std::make_pair(L, U);
}
#endif