#include "matrix.h"
#include "solver.h"
#include "argh.h"

#include <iostream>
#include <fstream>
#include <chrono>
#include <utility>

std::pair<Matrix<double>, Matrix<double>> prepare(double a1, size_t n, size_t index) 
{
    double a2, a3;
    a2 = a3 = -1;

    auto M = Matrix<double>::diag(n, a1);
    for(size_t i = 0; i < n; i++) {
        if (i >= 1) {
            M(i, i-1) = a2;
            M(i-1, i) = a2;
        }

        if (i >= 2) {
            M(i, i-2) = a3;
            M(i-2, i) = a3;
        }
    }

    Matrix<double> b(n, 1);
    for (size_t i = 0; i < n; ++i) {
        b(i, 0) = sin((i + 1)*((index / 1000) % 10 + 1));
    }

    return std::make_pair(M, b);
}

void help(const char* program)
{
    std::cout 
        << "Usage: " << program << " <N> <index> <method> [options]" << std::endl
        << "    <N>      - size of generated matrix" << std::endl
        << "    <index>  - index used for matrix generation" << std::endl
        << "    <method> - one of: gauss, jacobi, lu" << std::endl;

    std::cout << std::endl
        << "Options:" << std::endl
        << "    -x, --solution - displays result (x)" << std::endl
        << "    -M, --matrix   - displays matrix (M)" << std::endl
        << "    -b, --result   - displays result vector (b)" << std::endl
        << "    -r, --residuum - displays residuum norm" << std::endl 
        << std::endl
        << "    -a a1               - sets diagonal element" << std::endl
        << "    -o, --output output - sets residuum output file" << std::endl
    ;
}

int main(int argc, const char* argv[])
{
    std::ofstream output;
    JacobiSolver jacobi;
    GaussSolver gauss;
    LUSolver lu;

    Solver* solver;

    if (argc < 4) {
        help(argv[0]);
        return -1;
    }

    argh::parser args;
    args.add_param("a");
    args.add_params({"o", "output"});
    args.parse(argc, argv);

    size_t N, index;
    std::string method, csv;
    double a1;

    args(1) >> N;
    args(2) >> index;
    args(3) >> method;
    args({"a"}, (index / 100) % 10 + 5) >> a1;
    args({"o", "output"}, "") >> csv;

    auto tuple = prepare(a1, N, index);
    auto M = std::get<0>(tuple);
    auto b = std::get<1>(tuple);

    if (args[{"M", "matrix"}]) std::cout << "Matrix:" << std::endl << M << std::endl;
    if (args[{"b", "result"}]) std::cout << "Result:" << std::endl << b << std::endl;

    if (method == "gauss") solver = &gauss;
    else if (method == "jacobi") solver = &jacobi;
    else if (method == "lu") solver = &lu;
    else {
        std::cout << "Method unknown: " << method << std::endl;
        help(argv[0]);
        return -1;
    }

    IterativeSolver* iterative = dynamic_cast<IterativeSolver*>(solver);
    if (iterative && !method.empty()) {
        output.open(csv);
        iterative->output(&output);
    }

    std::cout << "Solving using " << method << " method..." << std::endl;
    auto start = std::chrono::high_resolution_clock::now();

    auto x = solver->solve(M, b);
    if (args[{"x", "solution"}]) std::cout << "Solution:" << std::endl << x << std::endl;
    if (args[{"r", "residuum"}]) std::cout << "Residuum:" << norm(M*x - b, 2) << std::endl;
    
    std::chrono::duration<double> duration = std::chrono::high_resolution_clock::now() - start;

    if (iterative) std::cout << "Iterations: " << iterative->iterations() << std::endl;
    
    std::cout << "Solving time: " << duration.count() << "s" << std::endl;
}