#include <armadillo>
#include <iostream>
#include <fstream>
#include <iomanip>
#include "include/helpers.hpp"

using namespace std;
using namespace arma;

int main(int argc, char** argv) {
    int max_N = 200; // Maximum matrix size
    if (argc > 1){ // Allow user to specify max matrix size during call
        max_N = std::stoi(argv[1]);
    }

    string filename;
    mat (*matrix_creator_func)(int, double);
    if (argc > 2 && string(argv[2]) == "dense") { // Use dense matrix if specified
        filename = "scaling_data_dense.txt";
        matrix_creator_func = create_dense_matrix;
    } else { // Default to tridiagonal matrix
        filename = "scaling_data.txt";
        matrix_creator_func = create_tridiagonal_matrix;
    }
    ofstream outfile;
    outfile.open(filename);
    outfile << "N,Arma_time,Analytical_time,Jacobi_time,Jacobi_iterations,Jacobi_converged\n";

    for(int N = 10; N <= max_N; N += 10) { // Iterate over matrix sizes in steps of 10
        double h = 1.0 / (N + 1);
        mat A = matrix_creator_func(N, h);

        // Armadillo solver
        vec arma_eigenvalues;
        mat arma_eigenvectors;
        auto start_arma = chrono::high_resolution_clock::now();
        eig_sym(arma_eigenvalues, arma_eigenvectors, A);
        auto end_arma = chrono::high_resolution_clock::now();
        chrono::duration<double> arma_duration = end_arma - start_arma;

        // Analytical solver
        vec analytical_eigenvalues;
        mat analytical_eigenvectors;
        auto start_analytical = chrono::high_resolution_clock::now();
        analytical_solution(N, h, analytical_eigenvalues, analytical_eigenvectors);
        auto end_analytical = chrono::high_resolution_clock::now();
        chrono::duration<double> analytical_duration = end_analytical - start_analytical;

        // Jacobi solver
        vec jacobi_eigenvalues;
        mat jacobi_eigenvectors;
        int iterations;
        bool converged;
        double eps = 1e-8;
        int maxiter = 1e6;
        auto start_jacobi = chrono::high_resolution_clock::now();
        jacobi_eigensolver(A, eps, jacobi_eigenvalues, jacobi_eigenvectors, maxiter, iterations, converged);
        auto end_jacobi = chrono::high_resolution_clock::now();
        chrono::duration<double> jacobi_duration = end_jacobi - start_jacobi;

        // Write results to file
        outfile << N << ","
                << arma_duration.count() << ","
                << analytical_duration.count() << ","
                << jacobi_duration.count() << ","
                << iterations << ","
                << (converged ? "true" : "false") << "\n";

        cout << "Completed N = " << N << endl;
    }
    outfile.close();
    return 0;
}