#include "../include/helpers.hpp"

arma::mat create_tridiagonal_matrix(int N, double h) {
    arma::mat A = arma::mat(N, N, arma::fill::zeros);

    double diag_value = 2.0 / (h * h);
    double off_diag_value = -1.0 / (h * h);

    for (int i = 0; i < N; ++i) {
        A(i, i) = diag_value;
        if (i > 0) { // There is now lower diagnonal element in the first row
            A(i, i - 1) = off_diag_value;
        }
        if (i < N - 1) { // There is no upper diagonal element in the last row
            A(i, i + 1) = off_diag_value;
        }
    }

    return A;
}

arma::mat create_dense_matrix(int N, double h) {
    arma::mat A = arma::mat(N, N).randn();
    A = arma::symmatu(A); // Make it symmetric
    return A;
}

void normalize_columns(arma::mat& A) { // Useful for eigenvector normalization
    A = arma::normalise(A, 2, 0); // 2-norm along columns (0)
    
    // Flip sign to ensure first element is positive
    for (size_t col = 0; col < A.n_cols; ++col) {
        if (A(0, col) < 0) {
            A.col(col) = -A.col(col);
        }
    }
}

void analytical_solution(int N, double h, arma::vec& eigenvalues, arma::mat& eigenvectors) {
    eigenvalues = arma::vec(N);
    eigenvectors = arma::mat(N, N);
    double d = 2.0 / (h * h);
    double a = -1.0 / (h * h);

    for (int j = 0; j < N; ++j) {
        eigenvalues(j) = d + 2 * a * std::cos((j + 1) * arma::datum::pi / (N + 1)); // Eigenvalues

        for (int i = 0; i < N; ++i) {
            eigenvectors(i, j) = std::sin((i + 1) * (j + 1) * arma::datum::pi / (N + 1)); // Eigenvectors
        }
    }

    normalize_columns(eigenvectors); // Normalize the eigenvectors
}

double max_offdiag_symmetric(const arma::mat& A, int& k, int &l) {
    int N = A.n_rows;
    double max_val = 0.0;

    for (int i = 0; i < N; ++i) {
        for (int j = i + 1; j < N; ++j) { // Only check upper triangle
            if (std::abs(A(i, j)) > max_val) {
                max_val = std::abs(A(i, j));
                k = i;
                l = j;
            }
        }
    }

    return max_val;
}

void jacobi_rotate(arma::mat& A, arma::mat& R, int k, int l) {
    if (A(k, l) != 0.0) {
        double tau = (A(l, l) - A(k, k)) / (2.0 * A(k, l));
        double t;
        if (tau >= 0) {
            t = 1.0 / (tau + std::sqrt(1.0 + tau * tau));
        } else {
            t = -1.0 / (-tau + std::sqrt(1.0 + tau * tau));
        }
        double c = 1.0 / std::sqrt(1 + t * t);
        double s = t * c;
        double temp;

        // Update matrix A
        double a_kk = A(k, k);
        double a_ll = A(l, l);
        A(k, k) = a_kk * c * c - 2.0 * A(k, l) * c * s + a_ll * s * s;
        A(l, l) = a_ll * c * c + 2.0 * A(k, l) * c * s + a_kk * s * s;
        A(k, l) = 0.0; // Hard-coding to zero
        A(l, k) = 0.0; // Hard-coding to zero

        for (int i = 0; i < A.n_rows; ++i) {
            if (i != k && i != l) {
                temp = A(i, k);
                A(i, k) = temp * c - A(i, l) * s;
                A(k, i) = A(i, k); // Since A is symmetric
                A(i, l) = temp * s + A(i, l) * c;
                A(l, i) = A(i, l); // Since A is symmetric
            }

            // Update the eigenvector matrix R
            temp = R(i, k);
            R(i, k) = temp * c - R(i, l) * s;
            R(i, l) = temp * s + R(i, l) * c;
        }
    }
}

void jacobi_eigensolver(const arma::mat& A, double eps, arma::vec& eigenvalues, arma::mat& eigenvectors, 
                        const int maxiter, int& iterations, bool& converged) {
    int N = A.n_rows;
    arma::mat A_copy = A; // Make a copy to avoid modifying the original matrix
    eigenvectors = arma::eye<arma::mat>(N, N); // Initialize eigenvector matrix as identity
    iterations = 0;
    converged = false;
    int k, l;
    double max_offdiag;
    while (iterations < maxiter) {
        max_offdiag = max_offdiag_symmetric(A_copy, k, l);
        if (max_offdiag < eps) {
            converged = true;
            break;
        }
        jacobi_rotate(A_copy, eigenvectors, k, l);
        iterations++;
    }
    eigenvalues = A_copy.diag(); // Extract eigenvalues from the diagonal of A_copy
}

void sort_eigenpairs(arma::vec& eigenvalues, arma::mat& eigenvectors) {
    arma::uvec indices = arma::sort_index(eigenvalues, "ascend");
    eigenvalues = eigenvalues(indices);
    eigenvectors = eigenvectors.cols(indices);
}