/** @file jacobi.cpp
 *  @brief Implementation of the jacobi rotation method
 *
 *  @author Cory Alexander Balaton (coryab)
 *  @author Janita Ovidie Sandtrøen Willumsen (janitaws)
 *  @bug The eigenvalues fail the test. When comparing them to arma::eigsym
 *  there is one value that is way off when testing with a 6x6 matrix.
 */
#include <algorithm>
#include <cmath>
#include <utility>

#include "utils.hpp"
#include "jacobi.hpp"
#include "matrix.hpp"

void jacobi_rotate(arma::mat& A, arma::mat& R, int k, int l)
{
    double tau, t, c, s, sr, a_kk, a_ik, r_ik;

    tau = (A(l,l) - A(k,k))/(2.*A(k,l));
    
    sr = sqrt(1 + tau*tau);
    t = - tau + (tau > 0. ? sr : -sr);

    c = 1. / std::sqrt(1. + t*t);

    s = c*t;

    // Calculating the diagonal elements at index k an l
    a_kk = A(k,k);
    A(k,k) = a_kk*c*c - 2.*A(k,l)*c*s + A(l,l)*s*s;
    A(l,l) = A(l,l)*c*c + 2.*A(k,l)*c*s + a_kk*s*s;
    A(k,l) = A(l,k) = 0.; // These are 0 by definition

    // Calculating the other values on rows/cols k and l of A
    for (int i=0; i < A.n_rows; i++) {
        if (i == k || i == l) {
            continue;
        }
        a_ik = A(i,k);
        A(i,k) = a_ik*c - A(i,l)*s;
        A(k,i) = A(i,k);
        A(i,l) = A(i,l)*c + a_ik*s;
        A(l,i) = A(i,l);
    }

    // Calculate the values for R
    for (int i=0; i < R.n_rows; i++) {
        r_ik = R(i,k);
        R(i,k) = r_ik*c - R(i,l)*s;
        R(i,l) = R(i,l)*c + r_ik*s;
    }
}

void jacobi_eigensolver(const arma::mat& A, 
                        double eps, 
                        arma::vec& eigenvalues, 
                        arma::mat& eigenvectors, 
                        const int maxiter, 
                        int& iterations, 
                        bool& converged)
{
    int k, l;
    double max_offdiag;

    // Initialize matrices
    arma::mat A_m = A;
    arma::mat R = arma::mat(A.n_rows, A.n_cols, arma::fill::eye);
    
    iterations = 0;

    // Use a do while so that max_offdiag gets a value before doing a comparison
    do {
    max_offdiag = max_offdiag_symmetric(A_m, k, l);
        jacobi_rotate(A_m, R, k, l);

    } while (max_offdiag >= eps && ++iterations < maxiter);

    // Get the diagonal of A
    arma::vec res = arma::diagvec(A_m);
    // Get the indeces of res in sorted order
    arma::uvec sorted_indeces = arma::sort_index(res);

    eigenvalues.resize(res.n_elem);
    eigenvectors.resize(res.n_elem, res.n_elem);
    
    // Set the eigenvalues and their corresponding vectors in order
    for (int i=0; i < res.n_elem; i++) {
        eigenvalues[i] = res[sorted_indeces[i]]; 
        eigenvectors.insert_cols(i, R.col(sorted_indeces[i]));
    }

    converged = max_offdiag < eps;
}