#include <armadillo>
#include <iostream>
#include <fstream>
#include <iomanip>

#include "include/helpers.hpp"

using namespace std;
using namespace arma;

void write_boundary_conditions(ofstream& outfile, int num_eigenvectors, int decimal_precision){
    for(int i = 0; i < num_eigenvectors; i++){
        outfile << setprecision(decimal_precision) << scientific << 0.0 << "," << 0.0; // Boundary condition at x=1
        if (i < num_eigenvectors - 1) {
            outfile << ",";
        }
        else {
            outfile << endl;
        }
    }
}

int main(int argc, char** argv) {
    int n = 10; 
    const int NUM_EIGENVECTORS = 3; // Number of eigenvectors to save
    const int DECIMAL_PRECISION = 12; // Decimal precision for output
    if (argc > 1){
        n = std::stoi(argv[1]);
    }
    int N = n - 1; // Number of interior points
    double h = 1.0 / (n); // Step size
    mat A = create_tridiagonal_matrix(N, h);


    // Generate analytical solution
    vec analytical_eigenvalues;
    mat analytical_eigenvectors;
    analytical_solution(N, h, analytical_eigenvalues, analytical_eigenvectors);
    normalize_columns(analytical_eigenvectors);
    sort_eigenpairs(analytical_eigenvalues, analytical_eigenvectors);
    
    // Using self-made Jacobi solver
    vec jacobi_eigenvalues;
    mat jacobi_eigenvectors;
    int iterations;
    bool converged;
    double eps = 1e-8;
    int maxiter = 1e6;
    jacobi_eigensolver(A, eps, jacobi_eigenvalues, jacobi_eigenvectors, maxiter, iterations, converged);
    if (converged) {
        cout << "Jacobi method converged in " << iterations << " iterations." << endl;
    } else {
        cout << "Jacobi method did not converge within the maximum number of iterations." << endl;
        return 1;
    }
    normalize_columns(jacobi_eigenvectors); // Normalize the eigenvectors
    sort_eigenpairs(jacobi_eigenvalues, jacobi_eigenvectors);

    // Write results to file
    ofstream outfile;
    string filename = "eigen_data_" + to_string(n) + ".txt"; // e.g., eigen_data_10.txt
    outfile.open(filename);
    for(int i = 0; i < NUM_EIGENVECTORS; i++){
        outfile << "analytical_solution_" << i+1 << "," << "jacobi_solution_" << i+1; // Column headers
        if(i < NUM_EIGENVECTORS - 1){
            outfile << ",";
        } else {
            outfile << endl;
        }
    }
    write_boundary_conditions(outfile, NUM_EIGENVECTORS, DECIMAL_PRECISION); // Boundary condition at x=0
    for(int j = 0; j < N; j++){
        for(int i = 0; i < NUM_EIGENVECTORS; i++){
            outfile << setprecision(DECIMAL_PRECISION) << scientific << analytical_eigenvectors(j,i) << "," << jacobi_eigenvectors(j,i);
            if(i < NUM_EIGENVECTORS - 1){
                outfile << ",";
            } else {
                outfile << endl;
            }
        }
    }
    write_boundary_conditions(outfile, NUM_EIGENVECTORS, DECIMAL_PRECISION); // Boundary condition at x=1
    outfile.close();
    return 0;
}