/* * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * *
 * Filename    : algorithms.cpp                                                                        *
 * Project     : Planewalker - Schnorr-Euchner sphere decoder simulation for space-time lattice codes  *
 * Authors     : Pasi Pyrrö, Oliver Gnilke                                                             *
 * Version     : 1.0                                                                                *
 * Copyright   : Aalto University ~ School of Science ~ Department of Mathematics and Systems Analysis *
 * Date        : 9.1.2017                                                                              *
 * Language    : C++ (2011 or newer standard)                                                          *
 * Description : All custom mathematical algorithms not found in Armadillo or STL are collected here   *
 * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * */

#define ARMA_NO_DEBUG /* disable Armadillo bound checks for addiotional speed */
#define FPLLL_WITH_ZDOUBLE
#define FPLLL_WITH_LONG_DOUBLE

#include <iostream>
#include <armadillo> /* linear algebra library */
#include <algorithm>
#include <complex>
#include <vector>
#include <string>
#include <set>
#include <random>
#include <omp.h>
#include <tuple>

/* library for LLL reduction (optional) */
#ifdef USE_LLL
#include <fplll.h>
#endif
/* required packages to install:
 *   libmpfr-dev
 *   libgmp-dev
 */

#include "algorithms.hpp"
#include "misc.hpp"

using namespace std;
using namespace arma;

/* Signum function for schorr-euchnerr sphere decoder (sesd)
 * ---------------------------------------------------------
 * Note that sesd_sign(0) = -1
 */
int sesd_sign(double x) {
    return (x <= 0) ? -1 : 1;
}

/* compute greatest common divisor of a and b
 * credits go to: https://codereview.stackexchange.com/a/66735
 */
int gcd(int a, int b) {
    return b == 0 ? a : gcd(b, a % b);
}

/* Rounds x to nearest integer in S
 * --------------------------------
 * E.g. let S = 4-PAM and x = 1.8 then
 * nearest_symbol(1.8) = 3
 * cf. round(1.8) = 2
 */
double nearest_symbol(double x, const vector<int> &S){
    double min = 10e6;
    double d = 0.0;
    double nearest = 0.0;

    for (const int symbol : S) {
        d = fabs(x - symbol);
        if (d < min) {
            nearest = symbol;
            min = d;
        }
    }
    return nearest;
}

double nearest_symbol_floor(double x, const vector<int> &S){

    // for (const int symbol : S) {
    //     if (symbol < x) {
    //         return symbol;
    //     }
    // }
    for (int i = S.size()-1; i >= 0; i--) {
        if (S[i] <= x) {
            return S[i];
        }
    }
    // return S[S.size()-1];
    return S[0];
}

double nearest_symbol_ceil(double x, const vector<int> &S){

    for (auto i = 0u; i < S.size()-1; i++) {
        if (S[i] >= x) {
            return S[i];
        }
    }
    return S[S.size()-1];
}

/* A simple heuristic function based on the idea that the volume of the hypersphere divided by the volume of the fundamental region of the lattice
 * should roughly equal to the number of lattice points inside the hypersphere
 * i.e. this function estimates the squared radius required for the hypersphere in order to have 2^s codewords (lattice points) inside it
 * (for q-PAM we need to multiply the generator matrix by 2 to get the correct volume for the fundamental region)
 */
double estimate_squared_radius(const mat &G, int s){
    int n = params["no_of_matrices"];
    double pi = 3.1415926535897;
    // cout << "lattice constant: " << to_string(sqrt(det(4*G.t()*G))) << endl;
    // cout << "volume of the sphere: " << to_string(pow(pi, n/2.0)*pow(dparams["spherical_shaping_max_power"], n/2)/tgamma(n/2.0 + 1.0)) << endl;

    /* The basic idea: 2^s = vol(Sphere(R))/vol(Lambda)
     * --> Solve for R when s and Lambda (G) are given
     */

    /* This algorithm uses this equation for the volume of the n-ball:
     * https://en.wikipedia.org/wiki/Volume_of_an_n-ball
     */
    return pow(pow(2.0, s)*tgamma(n/2.0 + 1.0)*sqrt(det(4*G.t()*G))*pow(pi, n/-2.0), 2.0/n);
}

/* Generates a new symbolset where elements are between given lower and upper bounds */
vector<int> slice_symbset(const vector<int> &S, double lb, double ub){
    vector<int> retval;
    for (const int q : S){
        if (lb <= (double)q && (double)q <= ub) {
            retval.push_back(q);
        }
    }
    return retval;
}

/* Takes the squared Frobenius norm from a complex matrix A */
double frob_norm_squared(const cx_mat &A){
    double sum = 0;
    for (auto i = 0u; i < A.n_rows; i++)
        for (auto j = 0u; j < A.n_cols; j++)
            sum += norm(A(i,j));
    return sum;
}

/* Makes sure the upper triangular matrix R has only positive diagonal elements */
void process_qr(mat &Q, mat &R){
    for (auto i = 0u; i < R.n_cols; i++){
        if (R(i,i) < 0.0){
            // R.col(i) *= -1;
            // Q.row(i) *= -1;
            R.row(i) *= -1;
            Q.col(i) *= -1;
        }
    }
}

/* Create the lattice generator matrix G out of basis matrices B_i
 * i.e. G = vec(B_i), i = 1,...,k
 */
cx_mat create_generator_matrix(const vector<cx_mat> &bases){
    int m = params["no_of_transmit_antennas"];
    int t = params["time_slots"];
    int k = params["no_of_matrices"];
    cx_mat G(k, m*t);
    for(int i = 0; i < k; i++){
        G.row(i) = vectorise(bases[i], 1);
    }
    return G.st();
}

/***** DEPRECATED *****
 * Use
 * G_real = to_real_matrix(create_generator_matrix(bases));
 * instead
 **********************/

/* Create the real valued lattice generator matrix G out of basis matrices */
// mat create_real_generator_matrix(const vector<cx_mat> &bases){
//     int m = params["no_of_transmit_antennas"];
//     int t = params["time_slots"];
//     int k = params["no_of_matrices"];
//     mat G(2*m*t,k);
//     for(int i = 0; i < k; i++){
//         G.col(i) = to_real_vector(bases[i]);
//     }
//     return G;
// }

/* Inverse of the create_generator_matrix() function
 * Useful for generating basis matrix files in conjunction with
 * output_complex_matrix()
 * See main.cpp for an example of this use case
 */
vector<cx_mat> generator_to_bases(const cx_mat &G){
    int m = params["no_of_transmit_antennas"];
    int t = params["time_slots"];
    int idx = 0;

    vector<cx_mat> bases;
    cx_mat X(m, t);
    cx_vec col(m*t);

    for (auto i = 0u; i < G.n_cols; i++){
        X.zeros();
        col = G.col(i);
        for (int j = 0; j < m; j++){
            for (int s = 0; s < t; s++){
                X(j,s) = col[idx];
                idx++;
            }
        }
        idx = 0;
        bases.push_back(X);
    }
    return bases;
}

/* Vectorizes a complex matrix A to real vector (each row is concatenated) */
vec to_real_vector(const cx_mat &A, bool row_wise){
    vec a(2*A.n_elem);
    // as the imaginary parts are now independent real elements in the vector,
    // we need to double the number of elements in the returned vector
    int index = 0;
    if (row_wise){
        for (auto i = 0u; i < A.n_rows; i++){
            for (auto j = 0u; j < A.n_cols; j++){
                auto z = A(i,j);
                a[index++] = z.real();
                a[index++] = z.imag();
            }
        }
    } else { /* concatenate columns instead */
        for (auto i = 0u; i < A.n_cols; i++){
            for (auto j = 0u; j < A.n_rows; j++){
                auto z = A(j,i);
                a[index++] = z.real();
                a[index++] = z.imag();
            }
        }
    }
    return a;
}

/* Converts a complex matrix A to its real representation */
mat to_real_matrix(const cx_mat &A){
    mat B(2*A.n_rows, A.n_cols);
    for(auto i = 0u; i < A.n_cols; i++){
        B.col(i) = to_real_vector(A.col(i));
    }
    return B;
}

/* Converts a real matrix constructed out of a complex matrix back to complex */
cx_mat to_complex_matrix(const mat &A){
    if (A.n_rows % 2 != 0){
        log_msg("to_complex_matrix: the input matrix needs to have an even number of rows!", "Error");
        exit(1);
    }

    cx_mat B(A.n_rows/2, A.n_cols);
    // int index = 0;
    for (auto j = 0u; j < A.n_cols; j++){
        for (auto i = 0u; i < B.n_rows; i++){
            B(i,j) = complex<double>(A(2*i,j), A(2*i+1,j));
        }
        // index = 0;
    }
    return B;
}

/* generates a random n x m 'full' complex matrix from normal distribution */
cx_mat create_random_matrix(int n, int m, double mean, double variance){
    normal_distribution<double> distr(mean, sqrt(variance));
    cx_mat A(n,m);
    return A.imbue([&]() { /* Armadillo magic with lambda function */
        return complex<double>(distr(mersenne_twister), distr(mersenne_twister));
    });
}

/* generates a random n x n diagonal complex matrix from normal distribution
 * Used for generating the SISO channel matrix
 */
cx_mat create_random_diag_matrix(int n, double mean, double variance){
    normal_distribution<double> distr(mean, sqrt(variance));
    cx_vec diag(n);
    cx_mat A(n,n);

    A.zeros();
    diag.imbue([&]() { /* Armadillo magic with lambda function */
        return complex<double>(distr(mersenne_twister), 0.0);
    });
    A.diag() = diag;
    return A;
}

/* Creates q-PAM symbol set (integer vector)
 * q is usually some power of two, i.e. q = 2, 4, 8, 16...
 * Example: 4-PAM = {-3, -1, 1, 3}
 */
vector<int> create_symbolset(int q){
    vector<int> symbset(q);
    for (int u = 0; u < q; u++) {
        symbset[u] = 2*u - q + 1;
    }
    return symbset;
}

/* Calculates all possible combinations of elements for code vector _a_
   given the set of feasible symbols (element values)
   Used for generating the whole codebook, rarely needed though */
void combinations(parallel_set< vector<int> > &comblist, const vector<int> &symbset, vector<int> comb, int dim){
    comblist.par_insert(comb);
    if (dim >= 0){
        #pragma omp parallel
        {
            vector<int> new_comb = comb;
            #pragma omp for
            for (size_t i = 0; i < symbset.size(); i++){
                new_comb[dim] = symbset[i];
                combinations(comblist, symbset, new_comb, dim-1);
            }
        }
    }
}

/* Helper function for above combinations algorithm */
set< vector<int> > comb_wrapper(const vector<int> &symbset, int vector_len){
    parallel_set< vector<int> > comblist;
    vector<int> init(vector_len);
    for (int i = 0; i < vector_len; i++)
        init[i] = symbset[0];
    combinations(comblist, symbset, init, vector_len-1);
    return comblist;
}

/* Creates a random codeword from basis matrices B_i and symbolset x-PAM */
pair<vector<int>, cx_mat> create_random_codeword(const vector<cx_mat> &bases, const vector<int> &symbolset){
    int m = params["no_of_transmit_antennas"];
    int t = params["time_slots"];
    int k = params["no_of_matrices"];
    int q = params["x-PAM"];
    int random_index = 0;

    cx_mat X(m, t);
    vector<int> coeffs(k);
    X.zeros();

    uniform_int_distribution<int> dist(0, q-1);

    for (int i = 0; i < k; i++) {
        random_index = dist(mersenne_twister);
        coeffs[i] = symbolset[random_index];
        X = X + symbolset[random_index]*bases[i];
    }

    return make_pair(coeffs, X);
}

/* Attempts to take account the probability bias when picking vector coefficients uniform randomly
   from hyperplane intervals within a hypersphere */
// vector<int> create_unbiased_subset(const mat &R, const vector<int> &subset, vec xt,
//                                    vec ener, vec curr, double radius, int i){
//     if (i <= 0) return subset; /* can't perform lookahead for the lowest dimension */

//     vector<int> unbiased_subset; //, delta_vec;
//     unbiased_subset.reserve(4*subset.size());

//     double xiener = 0.0;
//     int ub = 0, lb = 0, delta = 0; //, delta_sum = 0, bias_correction_amount = 0;
//     int k = params["no_of_matrices"];

//     // cout << i << endl;
//     // cout << vec2str(xt, xt.size()) << endl;
//     // cout << vec2str(ener, ener.size()) << endl;
//     // cout << vec2str(curr, curr.size()) << endl << endl;

//     for (const int elem : subset) {
//         xt[i] = elem;
//         xiener = pow(R(i,i)*xt[i] + curr[i], 2);

//         for (int j = i; j < k; j++)
//             curr[i-1] += xt[j]*R(i-1,j);
//         ener[i-1] = ener[i] + xiener;

//         lb = nearest_symbol_ceil(-(sqrt(radius - ener[i-1]) + curr[i-1])/R(i-1,i-1), subset);
//         ub = nearest_symbol_floor((sqrt(radius - ener[i-1]) - curr[i-1])/R(i-1,i-1), subset);
//         delta = (ub - lb)/2 + 1;

//         // cout << lb << ", " << ub << ", " << delta << endl;

//         for (int d = 0; d < delta; d++)
//             unbiased_subset.push_back(elem);
//         // delta_vec.push_back(delta);
//         // delta_sum += delta;
//     }
//     // cout << vec2str(xt, xt.size()) << endl;
//     // cout << vec2str(ener, ener.size()) << endl;
//     // cout << vec2str(curr, curr.size()) << endl << endl;
//     // cout << endl;

//     // for (auto j = 0u; j < subset.size(); j++){
//     //     bias_correction_amount = delta_sum/gcd(delta_sum, delta_vec[j]);
//     //     for (int s = 0; s < bias_correction_amount; s++)
//     //         unbiased_subset.push_back(subset[j]);
//     // }
//     // cout << vec2str(unbiased_subset, unbiased_subset.size()) << endl << endl;
//     return unbiased_subset;
// }

/* Attempts to pick q-PAM coefficients which correspond to large subhyperspheres more often
   This should eliminate some the bias present (picking border points too often) in naive uniform random picking */
int unbiased_random_select(const mat &R, const vector<int> &subset, vec xt,
                                   vec ener, vec curr, double radius, int i){
    if (i <= 0) return 0; /* can't perform lookahead for the lowest dimension */

    vector<int> unbiased_subset, feasible_coeffs;
    vector<double> radiuses;
    vector<pair<double, double>> intervals;
    unbiased_subset.reserve(4*subset.size());
    radiuses.reserve(subset.size());
    intervals.reserve(subset.size());
    feasible_coeffs.reserve(subset.size());

    double xiener = 0.0, radius_sum = 0.0, r = 0.0, numer = 0.0, ub = 0, lb = 0;
    int k = params["no_of_matrices"];

    /* calculate subsphere radiuses for each coefficient */
    for (auto j = 0u; j < subset.size(); j++) {
        xt[i] = subset[j];
        xiener = pow(R(i,i)*xt[i] + curr[i], 2);
        r = radius - xiener - ener[i];
        if (r > 0){ /* We don't want negative radiuses, do we? */
            r = pow(r, k-1);
            radiuses.push_back(r);
            radius_sum += r;
            feasible_coeffs.push_back(subset[j]);
        }
    }

    /* Calculate intervals from [0,1] for each coefficient */
    for (auto j = 0u; j < radiuses.size(); j++) {
        lb = numer/radius_sum;
        ub = (numer+radiuses[j])/radius_sum;
        intervals.push_back(make_pair(lb, ub));
        numer += radiuses[j];
    }

    // intervals.back().second = 1.0;

    // cout << vec2str(radiuses, radiuses.size()) << endl;
    // cout << vec2str(feasible_coeffs, feasible_coeffs.size()) << endl;
    // cout << vec2str(intervals, intervals.size()) << endl << endl;
    // cout << endl;

    // uniform_real_distribution<double> dist(0.0, 1.0);
    double X = uniform_real_dist(mersenne_twister);
    // cout << X << endl << endl;

    /* Select the coefficient from distribution created earlier (weights large radiuses) */
    for (auto j = 0u; j < feasible_coeffs.size(); j++) {
        // cout << "[" << intervals[j].first << ", " << intervals[j].second << "]" << endl;
        if (intervals[j].first <= X && X <= intervals[j].second)
            return feasible_coeffs[j];
    }
    // cout << "------" << endl;

    return 0;
}

/* Creates a random codeword from basis matrices B_i and symbolset x-PAM within given radius (spherical shaping) */
pair<vector<int>, cx_mat> create_random_spherical_codeword(const vector<cx_mat> &bases, const mat &R, const vector<int> &S, double radius){
    int m = params["no_of_transmit_antennas"];
    int t = params["time_slots"];
    int k = params["no_of_matrices"];

    cx_mat X(m, t);
    vector<int> coeffs(k), subset;
    X.zeros();

    vec xt(k), curr(k), ener(k);
    xt.zeros(); curr.zeros(); ener.zeros();
    int i = k-1;
    double xiener = 0.0;
    double lb = 0.0; // lower bound
    double ub = 0.0; // upper bound

    // radius = sqrt(radius);

    /* Randomly pick each component of the codeword and check if it's still within the energy bound in each dimension */
    while (i >= 0){
        lb = -(sqrt(radius - ener[i]) + curr[i])/R(i,i);
        ub = (sqrt(radius - ener[i]) - curr[i])/R(i,i);
        // uniform_real_distribution<double> xirange(lb, ub);
        // xt[i] = nearest_symbol(xirange(mersenne_twister), S); // probability bias fix this
        subset = slice_symbset(S, lb, ub);
        // if (i >= k-10 /*max(k-13, 3)*/ && !subset.empty())
        //     subset = create_unbiased_subset(R, subset, xt, ener, curr, radius, i);

        if (subset.empty()) { /* No lattice points are within energy bounds in this dimension */
            i = k-1;
            xt.zeros(); curr.zeros(); ener.zeros();
            continue;
        }

        // if (i > 0)
        //     xt[i] = unbiased_random_select(R, subset, xt, ener, curr, radius, i);
        // else
        xt[i] = pick_uniform(subset);

        if (xt[i] == 0) {
            // cout << "something went wrong" << endl;
            i = k-1;
            xt.zeros(); curr.zeros(); ener.zeros();
            continue;
        }

        xiener = pow(R(i,i)*xt[i] + curr[i], 2);

        // if (xiener + ener[i] < radius + 1e-6) {
        if (i > 0) {
            for (int j = i; j < k; j++)
                curr[i-1] += xt[j]*R(i-1,j);
            ener[i-1] = ener[i] + xiener;
        }
        i--;
        // } else { /* we're outside the energy bound, start over */
        //     log_msg("test!", "Error");
        //     i = k-1;
        //     xt.zeros(); curr.zeros(); ener.zeros();
        // }
    }
    coeffs = conv_to<vector<int>>::from(xt);
    for (int i = 0; i < k; i++) {
        X = X + coeffs[i]*bases[i];
    }
    return make_pair(coeffs, X);
}

/* Creates a codebook from basis matrices B_i and symbolset x-PAM within given radius (spherical shaping) */
// vector<pair<vector<int>, cx_mat>> create_spherical_codebook(const vector<cx_mat> &bases, const mat &R, const vector<int> &S, double radius){
vector<pair<vector<int>, cx_mat>> create_spherical_codebook(const vector<cx_mat> &bases, const mat &R, const vector<int> &S,
                                                            double radius, vec xt, int dim, double dist) {
    int m = params["no_of_transmit_antennas"];
    int t = params["time_slots"];
    int k = params["no_of_matrices"];

    cx_mat X(m, t, fill::zeros);
    vector<int> coeffs(k); //, subset;
    vector<pair<vector<int>,cx_mat>> codebook, part;

    // vec xt(k), curr(k), ener(k), indices(k);
    // xt.zeros(); curr.zeros(); ener.zeros(); indices.zeros();
    // int i = k-1;
    int i = dim-1;
    double xidist = 0;
    // double xiener = 0.0;
    // double lb = 0.0; // lower bound
    // double ub = 0.0; // upper bound
    // bool nextlevel = true;

    // while (true) {
    //  if (nextlevel) {
       //      lb = -(sqrt(radius - ener[i]) + curr[i])/R(i,i);
       //      ub = (sqrt(radius - ener[i]) - curr[i])/R(i,i);
       //      subset = slice_symbset(S, lb, ub);
       //      if (subset.empty()) { /* No lattice points are within energy bounds in this dimension */
       //       cout << "empty interval!" << endl;
       //       if (i == k - 1) break;
       //       else {
       //           indices[i] = 0;
       //           // indices[i+1]++;
          //           i++;
          //           continue;
          //       }
       //      }
       //      nextlevel = false;
       //  }
    //      xt[i] = subset[indices[i]];

    //     xiener = pow(R(i,i)*xt[i] + curr[i], 2);

    //     cout << i << ": " << vec2str(xt, k) << ", xiener = " << xiener + ener[i] << " interval: [" << lb << ", " << ub << "], xt[i] = " << xt[i] << endl;
    //     cout << "subset: " << vec2str(subset, subset.size()) << endl;

    //     // if (xiener + ener[i] < radius + 1e-6) {
    //     if (i > 0) {
    //      curr[i-1] = 0;
    //         for (int j = i; j < k; j++)
    //             curr[i-1] += xt[j]*R(i-1,j);
    //         ener[i-1] = ener[i] + xiener;
    //         i--;
    //         nextlevel = true;
    //         cout << "lowering dimension!" << endl;
    //         // continue;
    //     } else {
    //      coeffs = conv_to<vector<int>>::from(xt);
    //      X.zeros();
    //      for (int i = 0; i < k; i++) {
          //       X = X + coeffs[i]*bases[i];
          //   }
          //   cout << "found point!" << endl;
          //   codebook.push_back(make_pair(coeffs, X));
    //     }
    //     // }
    //     if (indices[i] >= subset.size() - 1){
    //      cout << "subset clear!" << endl;
    //      if (i == k - 1) break;
       //   else {
       //       indices[i] = 0;
       //       // indices[i+1]++;
       //       i++;
       //       nextlevel = true;
       //      }
    //     } else indices[i]++;
    // }
    // return codebook;

    // #pragma omp parallel for
    for (auto j = 0u; j < S.size(); j++) {
        vec tmp = xt;
        tmp[i] = S[j];
        xidist = pow(dot(R.row(i).subvec(i, k-1), tmp.subvec(i, k-1)), 2) + dist;
        if (xidist <= radius) {
            if (i > 0) {
                part = create_spherical_codebook(bases, R, S, radius, tmp, dim-1, xidist);
                codebook.insert(codebook.end(), part.begin(), part.end());
            } else {
                coeffs = conv_to<vector<int>>::from(tmp);
                X.zeros();
                for (int s = 0; s < k; s++) {
                    X = X + coeffs[s]*bases[s];
                }
                // cout << vec2str(coeffs, k) << endl;
                codebook.push_back(make_pair(coeffs, X));
            }
        }
    }
    return codebook;
}

/* Creates a codebook (set of X matrices) from basis matrices B_i and symbolset q-PAM */
vector<pair<vector<int>,cx_mat>> create_codebook(const vector<cx_mat> &bases, const mat &R, const vector<int> &symbolset){
    int m = params["no_of_transmit_antennas"];
    int t = params["time_slots"];
    int k = params["no_of_matrices"];
    int cs = dparams["codebook_size"];
    int samples = params["energy_estimation_samples"];
    double P = dparams["spherical_shaping_max_power"];

    vector<pair<vector<int>,cx_mat>> codebook;
    pair<vector<int>,cx_mat> code;
    vector<int> tmp;
    cx_mat X(m, t);
    X.zeros();

    if(samples > 0){
        for (int j = 0; j < samples; j++){
            if (P > 0)
                code = create_random_spherical_codeword(bases, R, symbolset, P);
            else
                code = create_random_codeword(bases, symbolset);
            codebook.push_back(code);
        }
    } else {
        if (P > 0)
            return create_spherical_codebook(bases, R, symbolset, P, vec(k, fill::zeros), k, 0);

        if (cs > 1e6)
            log_msg("create_codebook: Generating a codebook of size " + to_string(cs) + ", this can be really slow!", "Warning");

        /* all possible combinations of code words */
        auto c = comb_wrapper(symbolset, k);

        // log_msg("All possible data vector combinations:");
        for (const auto &symbols : c){
            X.zeros();
            for (int i = 0; i < k; i++)
                X = X + symbols[i]*bases[i];
            // log_msg(vec2str(symbols, symbols.size()));
            // if (frob_norm_squared(X) > P + 1e-6 && P > 0) continue;
            codebook.push_back(make_pair(symbols, X));
        }
    }
    return codebook;
}

/* Computes the average and maximum energy of given codebook X */
pair<double,double> code_energy(const vector<pair<vector<int>,cx_mat>> &X){
    double sum = 0, max = 0, tmp = 0, average = 0;
    int cs = (int) X.size();

    for (int i = 0; i < cs; i++){
        tmp = frob_norm_squared(X[i].second);
        sum += tmp;
        if (tmp > max)
            max = tmp;
    }
    average = sum / cs;
    return make_pair(average, max);
}

/* Returns the shortest lattice basis vector from the lattice generator matrix G*/
cx_vec shortest_basis_vector(const cx_mat &G){
    cx_vec shortest(G.n_rows, fill::zeros);
    double len = 10e9;
    for (auto i = 0u; i < G.n_cols; i++){
        if (len > norm(G.col(i), 2)) {
            shortest = G.col(i);
            len = norm(shortest, 2);
        }
    }
    return shortest;
}

/* Check if x belongs to the right coset (when using coset encoding)
 * i.e. check if the difference vector sent_x - decoded_x belongs in
 * the given sublattice of the codebook carved out of lattice G_b
 */
bool coset_check(const cx_mat &Gb, const cx_mat &invGe, const Col<int> diff){
    cx_vec x = Gb*diff;
    cx_vec lambda = invGe*x;

    vec lambda_real = to_real_vector(lambda);

    /* check if lambda_real has only integer entries */
    for (const auto &l : lambda_real)
        if (fabs(l - round(l)) > 10e-5) /* tolerance */
            return false;

    return true;
}

/* calculates different rates related to the coset encoded block code */
tuple<double, double, double> code_rates(const cx_mat &Gb, const cx_mat &Ge){
    mat RGb = to_real_matrix(Gb);
    mat RGe = to_real_matrix(Ge);
    // cout << "Vol(Ge): " << sqrt(det(RGe.t()*RGe)) << endl;
    // cout << "Vol(Gb): " << sqrt(det(RGb.t()*RGb)) << endl;
    double r = log2(dparams["codebook_size"]);                          // overall rate
    double rt = log2(sqrt(det(RGe.t()*RGe))/sqrt(det(RGb.t()*RGb)));   // transmission rate
    double rc = r - rt;                                                // confusion rate
    return make_tuple(r, rt, rc);
}

/* algorithm for counting the lattice points inside a _dim_ dimensional hypersphere of radius _radius_
 * considers only a single radius, but probably faster than the function below thanks to parallelisation
 */
int count_points(const mat &R, const vector<int> &S, double radius, vec xt, int dim, double dist){

    int k = params["no_of_matrices"];
    int i = dim - 1;

    vector<int> counters(omp_get_max_threads());

    #pragma omp parallel for
    for (auto j = 0u; j < S.size(); j++) {
        vec tmp = xt;
        tmp[i] = S[j];
        double xidist = pow(dot(R.row(i).subvec(i, k-1), tmp.subvec(i, k-1)), 2) + dist;
        if (xidist <= radius) {
            if (i > 0) {
                counters[omp_get_thread_num()] += count_points(R, S, radius, tmp, dim-1, xidist);
            } else {
                counters[omp_get_thread_num()]++;
            }
        }
    }
    return accumulate(counters.begin(), counters.end(), 0);
}

/* algorithm for counting the lattice points inside a _dim_ dimensional hypersphere of radiuses r_1, r_2, ..., r_n */
vector<int> count_points_many_radiuses(const mat &R, const vector<int> &S, vector<double> radiuses, vec xt, int dim, double dist){

    int k = params["no_of_matrices"];
    int i = dim - 1;

    vector<int> counters(radiuses.size());

    for (auto j = 0u; j < S.size(); j++){
        vec tmp = xt;
        tmp[i] = S[j];
        double xidist = pow(dot(R.row(i).subvec(i, k-1), tmp.subvec(i, k-1)), 2) + dist;

        if (xidist <= radiuses[0]){
            if (i > 0){
                vector<int> counts = count_points_many_radiuses(R, S, radiuses, tmp, dim-1, xidist);
                for (auto s = 0u; s < radiuses.size(); s++)
                    counters[s] += counts[s];
            } else {
                for (auto k = 0u; k < radiuses.size(); k++){
                    if (xidist <= radiuses[k])
                        counters[k]++;
                }
            }
        }
    }
    return counters;
}

/*
 * Should return the LLL reduced (as close to orthogonal as possible) basis for lattice generated by G
 */
cx_mat LLL_reduction(const cx_mat &G) {
    #ifdef USE_LLL

        mat G_real = to_real_matrix(G);

        // ZZ_mat<double> fpG(G_real.n_rows, G_real.n_cols);
        // ZZ_mat<double> fpG(G_real.n_cols, G_real.n_rows);
        FP_mat<double> fpG(G_real.n_cols, G_real.n_rows);
        for (auto i = 0u; i < G_real.n_rows; i++) {
            for (auto j = 0u; j < G_real.n_cols; j++) {
                // fpG(i, j) = G_real(i, j);
                fpG(j, i) = G_real(i, j);
            }
        }
        cout << fpG << endl << endl;
        int status = lll_reduction(fpG, LLL_DEF_DELTA, LLL_DEF_ETA, LM_PROVED, FT_DEFAULT, 0, LLL_DEFAULT);
        if (status != RED_SUCCESS) {
            log_msg("LLL reduction failed!", "Error"); //with error '" + to_string(get_red_status_str(status)) + "'", "Error");
            exit(1);
        }
        cout << fpG << endl << endl;
        for (auto i = 0u; i < G_real.n_rows; i++) {
            for (auto j = 0u; j < G_real.n_cols; j++) {
                // G_real(i, j) = fpG(i, j).get_d();
                G_real(i, j) = fpG(j, i).get_d();
            }
        }
        cout << G_real << endl;
        return to_complex_matrix(G_real);

    #else /* Do nothing */

        return G;

    #endif
}


#include <cstdio>
#include <cstdlib>
#include <iostream>
#include <vector>
#include <list>

#include <boost/numeric/ublas/io.hpp>
#include <boost/numeric/bindings/lapack/geqrf.hpp>
#include <boost/numeric/bindings/lapack/orgqr.hpp>
#include <boost/numeric/bindings/traits/ublas_matrix.hpp>
#include <boost/numeric/ublas/matrix.hpp>
#include <boost/numeric/ublas/matrix_proxy.hpp>
#include <boost/numeric/ublas/matrix_sparse.hpp>
#include <boost/numeric/ublas/banded.hpp>
#include <boost/numeric/ublas/operation.hpp>
#include <boost/numeric/ublas/operation_sparse.hpp>

#include <boost/numeric/ublas/vector.hpp>
#include <boost/numeric/ublas/vector_proxy.hpp>
#include <boost/numeric/ublas/vector_sparse.hpp>
#include <boost/numeric/ublas/vector_of_vector.hpp>
#include <boost/numeric/ublas/vector_expression.hpp>

#ifdef OPENMP
#include <omp.h>
#endif

void Perform_Alpert_transform(boost::numeric::ublas::vector<double> f, std::vector<int> ki, int J,int dir, boost::numeric::ublas::vector<double>& w, std::list<std::list<boost::numeric::ublas::matrix<double,boost::numeric::ublas::column_major> > > Uj, std::list<std::vector<int> > part, std::vector<int> G);
void Compute_Uj(std::vector<std::vector<double> >p, std::vector<int> ki, int &J, std::list<std::list<boost::numeric::ublas::matrix<double,boost::numeric::ublas::column_major> > > &Uj, std::list<std::vector<int> > &part, std::vector<int> &G);
void Optimal_Alpert_transform(boost::numeric::ublas::vector<double> f, std::vector<std::vector<double> >p, std::vector<int> ki, int J, boost::numeric::ublas::vector<double>& w, std::list<std::list<boost::numeric::ublas::matrix<double,boost::numeric::ublas::column_major> > > Uj, std::list<std::vector<int> > part, std::vector<int> G, double NRMSE, double &R, double mx, double mn, double mxg, double mng, double alpha);


//
// Created by jachu on 27.02.18.
//

#include <iostream>

#include <Eigen/Dense>

#include <g2o/types/slam3d/se3quat.h>

#include "PlaneEstimator.hpp"

using namespace std;

PlaneEstimator::PlaneEstimator() {

}

PlaneEstimator::PlaneEstimator(const Eigen::MatrixXd &pts) {
    init(pts);
}

PlaneEstimator::PlaneEstimator(const Eigen::Vector3d &icentroid,
                               const Eigen::Matrix3d &icovar,
                               int inpts)
{
    init(icentroid,
         icovar,
         inpts);
}

void PlaneEstimator::init(const Eigen::MatrixXd &pts) {
    Eigen::Vector3d icentroid;
    Eigen::Matrix3d icovar;
    compCentroidAndCovar(pts, icentroid, icovar);
    init(icentroid, icovar, pts.cols());
}

void
PlaneEstimator::init(const Eigen::Vector3d &icentroid, const Eigen::Matrix3d &icovar, int inpts) {
    centroid = icentroid;
    covar = icovar;
    npts = inpts;
    compPlaneParams(centroid,
                    covar,
                    evecs,
                    evals,
                    planeEq);
}

void
PlaneEstimator::update(const Eigen::Vector3d &ucentroid, const Eigen::Matrix3d &ucovar, int unpts) {
//    cout << "centroid = " << centroid.transpose() << endl;
//    cout << "covar = " << covar << endl;
//    cout << "npts = " << npts << endl;
//    cout << "ucentroid = " << ucentroid.transpose() << endl;
//    cout << "ucovar = " << ucovar << endl;
//    cout << "unpts = " << unpts << endl;

    updateCentroidAndCovar(centroid,
                           covar,
                           npts,
                           ucentroid,
                           ucovar,
                           unpts,
                           centroid,
                           covar,
                           npts);

//    cout << "centroid = " << centroid.transpose() << endl;
//    cout << "covar = " << covar << endl;
//    cout << "npts = " << npts << endl;

    static constexpr int ptsLimit = 500000;
    if(npts > ptsLimit){
        double scale = (double)npts/ptsLimit;
        covar /= scale;
        npts = ptsLimit;
    }
//    cout << "centroid = " << centroid.transpose() << endl;
//    cout << "covar = " << covar << endl;
//    cout << "npts = " << npts << endl;

    compPlaneParams(centroid,
                    covar,
                    evecs,
                    evals,
                    planeEq);
}

double PlaneEstimator::distance(const PlaneEstimator &other) const {
    const Eigen::Vector3d &centroid1 = centroid;
    const Eigen::Vector3d &centroid2 = other.centroid;
    Eigen::Matrix3d covar1 = covar / npts;
    Eigen::Matrix3d covar2 = other.covar / other.npts;
    int npts1 = npts;
    int npts2 = other.npts;

//    cout << "centroid1 = " << centroid1.transpose() << endl;
//    cout << "covar1 = " << covar1 << endl;
//    cout << "centroid2 = " << centroid2.transpose() << endl;
//    cout << "covar2 = " << covar2 << endl;
//
//    Eigen::Vector3d centrDiff = centroid1 - centroid2;
//    Eigen::Matrix3d infComb = (covar1 + covar2).inverse();
//    Eigen::Matrix3d covarProd = covar1 * infComb * covar2;
//    Eigen::Vector3d meanProd = covar2 * infComb * centroid1 +
//                               covar1 * infComb * centroid2;
//    double detVal = (2*M_PI*(covar1 + covar2)).determinant();
//    double expVal = -0.5 * centrDiff.transpose() * infComb * centrDiff;
//    double normFactor = 1.0 / sqrt(detVal) * exp(expVal);
//    cout << "detVal = " << detVal << endl;
//    cout << "expVal = " << -2.0*expVal << endl;
//    cout << "normFactor = " << normFactor << endl;
//    cout << "dist = " << 1.0/normFactor << endl;

    Eigen::Vector3d centrDiff = centroid1 - centroid2;
    // covariance of the second plane relative to the centroid of the first plane
    Eigen::Matrix3d relCovar = covar2 + (centrDiff * centrDiff.transpose());

    const Eigen::Vector3d &normal = evecs.col(2);
    double varNorm = normal.transpose() * relCovar * normal;
//    cout << "covar2 = " << covar2 << endl;
//    cout << "relCovar = " << relCovar << endl;
//    cout << "normal = " << normal.transpose() << endl;
//    cout << "varNorm = " << varNorm << endl;
    return varNorm;
}

void PlaneEstimator::compCentroidAndCovar(const Eigen::MatrixXd &pts,
                                         Eigen::Vector3d &centroid,
                                         Eigen::Matrix3d &covar)
{
    Eigen::Vector4d mean = Eigen::Vector4d::Zero();
    for(int i = 0; i < pts.cols(); ++i){
        mean += pts.col(i);
    }
    mean /= pts.cols();

    centroid = mean.head<3>();

    Eigen::MatrixXd demeanPts = pts;
    for(int i = 0; i < demeanPts.cols(); ++i){
        demeanPts.col(i) -= mean;
    }

    covar = demeanPts.topRows<3>() * demeanPts.topRows<3>().transpose();
}

void PlaneEstimator::compPlaneParams(const Eigen::Vector3d centroid,
                                     const Eigen::Matrix3d &covar,
                                     Eigen::Matrix3d &evecs,
                                     Eigen::Vector3d &evals,
                                     Eigen::Vector4d &planeEq)
{
    Eigen::SelfAdjointEigenSolver<Eigen::Matrix3d> evd(covar);

    for(int i = 0; i < 3; ++i){
        evecs.col(i) = evd.eigenvectors().col(2 - i);
        evals(i) = evd.eigenvalues()(2 - i);
    }

    planeEq.head<3>() = evecs.col(2).cast<double>();
    // distance is the dot product of normal and point lying on the plane
    planeEq(3) = -planeEq.head<3>().dot(centroid);
}

void PlaneEstimator::transform(const Vector7d &transform) {
    g2o::SE3Quat transformSE3Quat(transform);
    Eigen::Matrix4d transformMat = transformSE3Quat.to_homogeneous_matrix();
    Eigen::Matrix3d R = transformMat.block<3, 3>(0, 0);
    Eigen::Vector3d t = transformMat.block<3, 1>(0, 3);
    Eigen::Matrix4d Tinvt = transformMat.inverse();
    Tinvt.transposeInPlace();

    // Eigen::Vector3d centroid;
    centroid = R * centroid + t;

    // Eigen::Matrix3d covar;
    covar = R * covar * R.transpose();

    // Eigen::Matrix3d evecs;
    evecs = R * evecs;

    // Eigen::Vector3d evals;
    // no need to transform

    // Eigen::Vector4d planeEq;
    planeEq = Tinvt * planeEq;

    // int npts;
    // no need to transform
}

void PlaneEstimator::updateCentroidAndCovar(const Eigen::Vector3d &centroid1,
                                            const Eigen::Matrix3d &covar1,
                                            const int &npts1,
                                            const Eigen::Vector3d &centroid2,
                                            const Eigen::Matrix3d &covar2,
                                            const int &npts2,
                                            Eigen::Vector3d &ocentroid,
                                            Eigen::Matrix3d &ocovar,
                                            int &onpts)
{
    ocentroid = (npts1 * centroid1 + npts2 * centroid2)/(npts1 + npts2);
    Eigen::Vector3d centrDiff = centroid1 - centroid2;
    double fact = ((double)npts1 * npts2)/(npts1 + npts2);
//    cout << "(npts1 * npts2)/(npts1 + npts2) = " << fact << endl;
//    cout << "(centrDiff * centrDiff.transpose()) = " << (centrDiff * centrDiff.transpose()) << endl;
    ocovar = covar1 + covar2 + fact*(centrDiff * centrDiff.transpose());
    onpts = npts1 + npts2;
}


// Software License for MTL
//
// Copyright (c) 2007 The Trustees of Indiana University.
//               2008 Dresden University of Technology and the Trustees of Indiana University.
//               2010 SimuNova UG (haftungsbeschränkt), www.simunova.com.
// All rights reserved.
// Authors: Peter Gottschling and Andrew Lumsdaine
//
// This file is part of the Matrix Template Library
//
// See also license.mtl.txt in the distribution.

#ifndef MTL_REDUCTION_FUNCTORS_INCLUDE
#define MTL_REDUCTION_FUNCTORS_INCLUDE

#include <cmath>
#include <boost/numeric/linear_algebra/identity.hpp>
#include <boost/numeric/mtl/operation/squared_abs.hpp>

namespace mtl { namespace vector {

struct one_norm_functor
{
    template <typename Value>
    static inline void init(Value& value)
    {
	using math::zero;
	value= zero(value);
    }

    template <typename Value, typename Element>
    static inline void update(Value& value, const Element& x)
    {
	using std::abs;
	value+= abs(x);
    }

    template <typename Value>
    static inline void finish(Value& value, const Value& value2)
    {
	value+= value2;
    }

    template <typename Value>
    static inline Value post_reduction(const Value& value)
    {
	return value;
    }
};

// sub-optimal if abs is not needed
struct two_norm_functor
{
    template <typename Value>
    static inline void init(Value& value)
    {
	using math::zero;
	value= zero(value);
    }

    template <typename Value, typename Element>
    static inline void update(Value& value, const Element& x)
    {
	using mtl::squared_abs;
	value+= squared_abs(x);
    }

    template <typename Value>
    static inline void finish(Value& value, const Value& value2)
    {
	value+= value2;
    }

    // After reduction compute square root
    template <typename Value>
    static inline Value post_reduction(const Value& value)
    {
	using std::sqrt;
	return sqrt(value);
    }
};

// same as two-norm without the root at the end
struct unary_dot_functor
  : two_norm_functor
{
    template <typename Value>
    static inline Value post_reduction(const Value& value)
    {
	return value;
    }
};

struct infinity_norm_functor
{
    template <typename Value>
    static inline void init(Value& value)
    {
	using math::zero;
	value= zero(value);
    }

    template <typename Value, typename Element>
    static inline void update(Value& value, const Element& x)
    {
	using std::abs; using std::max;
	value= max(value, Value(abs(x)));
    }

    template <typename Value>
    static inline void finish(Value& value, const Value& value2)
    {
	using std::abs; using std::max;
	value= max(value, Value(abs(value2)));
    }

    template <typename Value>
    static inline Value post_reduction(const Value& value)
    {
	return value;
    }
};

struct sum_functor
{
    template <typename Value>
    static inline void init(Value& value)
    {
	using math::zero;
	value= zero(value);
    }

    template <typename Value, typename Element>
    static inline void update(Value& value, const Element& x)
    {
	value+= x;
    }

    template <typename Value>
    static inline void finish(Value& value, const Value& value2)
    {
	value+= value2;
    }

    template <typename Value>
    static inline Value post_reduction(const Value& value)
    {
	return value;
    }
};

struct product_functor
{
    template <typename Value>
    static inline void init(Value& value)
    {
	using math::one;
	value= one(value);
    }

    template <typename Value, typename Element>
    static inline void update(Value& value, const Element& x)
    {
	value*= x;
    }

    template <typename Value>
    static inline void finish(Value& value, const Value& value2)
    {
	value*= value2;
    }

    template <typename Value>
    static inline Value post_reduction(const Value& value)
    {
	return value;
    }
};

struct max_functor
{
    template <typename Value>
    static inline void init(Value& value)
    {
	using math::identity;
	value= math::identity(math::max<Value>(), value); // ADL doesn't work here in g++ 4.4
    }

    template <typename Value, typename Element>
    static inline void update(Value& value, const Element& x)
    {
	value= math::max<Value>()(value, x);
    }

    template <typename Value>
    static inline void finish(Value& value, const Value& value2)
    {
	value= math::max<Value>()(value, value2);
    }

    template <typename Value>
    static inline Value post_reduction(const Value& value)
    {
	return value;
    }
};

struct min_functor
{
    template <typename Value>
    static inline void init(Value& value)
    {
	using math::identity;
	value= math::identity(math::min<Value>(), value); // ADL doesn't work here in g++ 4.4
    }

    template <typename Value, typename Element>
    static inline void update(Value& value, const Element& x)
    {
	value= math::min<Value>()(value, x);
    }

    template <typename Value>
    static inline void finish(Value& value, const Value& value2)
    {
	value= math::min<Value>()(value, value2);
    }

    template <typename Value>
    static inline Value post_reduction(const Value& value)
    {
	return value;
    }
};

}} // namespace mtl::vector

#endif // MTL_REDUCTION_FUNCTORS_INCLUDE


#include <armadillo>
#include <ForwardBackward.hpp>
#include <HSMM.hpp>
#include <cmath>
#include <json.hpp>
#include <iostream>
#include <memory>
#include <vector>

using namespace arma;
using namespace std;
using json = nlohmann::json;

namespace hsmm {

    // TODO: Make sure there is not any bias here.
    int sampleFromCategorical(rowvec pmf) {
        rowvec prefixsum(pmf);
        for(int i = 1; i < pmf.n_elem; i++)
            prefixsum(i) += prefixsum(i - 1);
        assert(abs(prefixsum(pmf.n_elem - 1) - 1.0) < 1e-7);
        return lower_bound(prefixsum.begin(), prefixsum.end(), randu()) -
                prefixsum.begin();
    }

    /**
     * HSMM implementation.
     */
    HSMM::HSMM(shared_ptr<AbstractEmission> emission, mat transition,
            vec pi, mat duration, int min_duration) : emission_(emission),
            duration_learning_choice_("histogram"), debug_(false),
            learning_transitions_(true), learning_pi_(true) {
        nstates_ = emission_->getNumberStates();
        ndurations_ = duration.n_cols;
        min_duration_ = min_duration;
        assert(min_duration_ >= 1);
        assert(nstates_ >= 1);
        assert(ndurations_ >= 1);
        setDuration(duration);
        setPi(pi);
        setTransition(transition);
        mat default_dirichlet_parameters = ones<mat>(nstates_, ndurations_);
        setDurationDirichletPrior(default_dirichlet_parameters);
    }

    HSMM::HSMM(shared_ptr<AbstractEmission> emission, int nstates, int ndurations,
            int min_duration) : HSMM(emission,
            ones<mat>(nstates, nstates) * (1.0/nstates),
            ones<vec>(nstates) * (1.0/nstates),
            ones<mat>(nstates, ndurations) * (1.0/ndurations), min_duration)
    {}

    void HSMM::setDuration(mat duration) {
        assert(duration.n_rows == nstates_);
        assert(duration.n_cols == ndurations_);
        duration_ = duration;
    }

    void HSMM::setEmission(shared_ptr<AbstractEmission> emission) {
        assert(emission->getNumberStates() == nstates_);
        emission_ = emission;
    }

    void HSMM::setPi(vec pi) {
        assert(pi.n_elem == nstates_);
        pi_ = pi;
    }

    void HSMM::setTransition(mat transition) {
        assert(transition.n_rows == transition.n_cols);
        assert(transition.n_rows == nstates_);
        transition_ = transition;
    }

    void HSMM::setPiFromLabels(field<ivec> seqs) {
        vec new_pi(nstates_, fill::zeros);
        for(auto& s: seqs)
            new_pi(s(0))++;
        new_pi = new_pi / accu(new_pi);
        setPi(new_pi);
    }

    void HSMM::setTransitionFromLabels(field<ivec> seqs) {
        mat new_transition(nstates_, nstates_, fill::zeros);
        for(auto& s: seqs) {
            const ivec& hs = computeViterbiStateDurationSequenceFromLabels(
                    s).first;
            for(int i = 0; i < hs.n_elem - 1; i++)
                new_transition(hs(i), hs(i + 1))++;
        }
        for(int i = 0; i < nstates_; i++) {
            auto row = new_transition.row(i);
            new_transition.row(i) = row / accu(row);
        }
        setTransition(new_transition);
    }

    void HSMM::setDurationFromLabels(field<ivec> seqs) {
        mat new_duration(nstates_, ndurations_, fill::zeros);
        for(auto& s: seqs) {
            pair<ivec,ivec> p = computeViterbiStateDurationSequenceFromLabels(
                    s);
            const ivec& hs = p.first;
            const ivec& dur = p.second;
            for(int i = 0; i < dur.n_elem; i++) {
                if (hs(i)<nstates_ && dur(i)<=min_duration_ + ndurations_)
                    new_duration(hs(i), dur(i) - min_duration_)++;
                else
                    cout << "There is an invalid (hs,d) entry" << endl;
            }
        }

        // Taking into account the Dirichlet prior.
        for(int i = 0; i < nstates_; i++)
            for(int j = 0; j < ndurations_; j++)
                new_duration(i, j) += dirichlet_alphas_(i, j) - 1;

        for(int i = 0; i < nstates_; i++) {
            auto row = new_duration.row(i);
            new_duration.row(i) = row / accu(row);
        }

        setDuration(new_duration);
    }

    void HSMM::setDurationDirichletPrior(mat alphas) {
        assert(alphas.n_rows == nstates_);
        assert(alphas.n_cols == ndurations_);
        for(auto alpha: alphas)
            assert(alpha > 1.0 - 1e-7);
        dirichlet_alphas_ = alphas;
    }

    void HSMM::setDurationLearningChoice(string choice) {
        duration_learning_choice_ = choice;
    }

    field<mat> HSMM::sampleSegments(int nsegments, ivec& hiddenStates,
            ivec& hiddenDurations) {
        assert(nsegments >= 1);

        // Generating states sequence.
        ivec states(nsegments);
        states(0) = sampleFromCategorical(pi_.t());
        for(int i = 1; i < nsegments; i++) {
            rowvec nstateDist = transition_.row(states(i - 1));
            states(i) = sampleFromCategorical(nstateDist);
        }

        // Generating durations from states.
        int sampleSequenceLength = 0;
        ivec durations(nsegments);
        for(int i = 0; i < nsegments; i++) {
            int currentDuration = sampleFromCategorical(
                    duration_.row(states(i))) + min_duration_;
            durations(i) = currentDuration;
            sampleSequenceLength += currentDuration;
        }

        // Generating samples
        field<mat> samples(sampleSequenceLength);
        int idx = 0;
        for(int i = 0; i < nsegments; i++) {
            field<mat> currSample = emission_->sampleFromState(
                    states(i), durations(i));
            samples.rows(idx, idx + durations(i) - 1) = currSample;
            idx += durations(i);
        }
        hiddenStates = states;
        hiddenDurations = durations;
        return samples;
    }

    field<field<mat>> HSMM::sampleMultipleSequences(int nsequences,
            int nsegments, field<ivec>& seqsHiddenStates,
            field<ivec>& seqsHiddenDurations) {
        field<field<mat>> mobs(nsequences);
        field<ivec> seqsHS(nsequences);
        field<ivec> seqsDur(nsequences);
        for(int s = 0; s < nsequences; s++) {
            ivec hs, dur;
            mobs(s) = sampleSegments(nsegments, hs, dur);
            seqsHS(s) = hs;
            seqsDur(s) = dur;
        }
        seqsHiddenStates = seqsHS;
        seqsHiddenDurations = seqsDur;
        return mobs;
    }

    void HSMM::init_params_from_data(const field<field<mat>> &mobs) {
        emission_->init_params_from_data(min_duration_, ndurations_, mobs);
    }

    bool HSMM::fit(field<field<mat>> mobs, field<Labels> mobserved_segments,
            int max_iter, double tol) {

        // Array initializations.
        int nseq = mobs.n_elem;
        assert(nseq >= 1 && nseq == mobserved_segments.n_elem);
        field<mat> malpha(nseq);
        field<mat> mbeta(nseq);
        field<mat> malpha_s(nseq);
        field<mat> mbeta_s(nseq);
        field<vec> mbeta_s_0(nseq);
        field<cube> meta(nseq);
        field<cube> mzeta(nseq);
        for(int i = 0; i < nseq; i++) {
            int nobs = mobs(i).n_elem;
            assert(nobs >= min_duration_);
            malpha(i) = zeros<mat>(nstates_, nobs);
            mbeta(i) = zeros<mat>(nstates_, nobs);
            malpha_s(i) = zeros<mat>(nstates_, nobs);
            mbeta_s(i) = zeros<mat>(nstates_, nobs);
            mbeta_s_0(i) = zeros<vec>(nstates_);
            meta(i) = zeros<cube>(nstates_, ndurations_, nobs);
            mzeta(i) = zeros<cube>(nstates_, nstates_, nobs - 1);
        }

        mat log_estimated_transition = log(transition_);
        vec log_estimated_pi = log(pi_);
        mat log_estimated_duration = log(duration_);
        double marginal_llikelihood = -datum::inf;
        bool convergence_reached = false;
        for(int i = 0; i < max_iter && !convergence_reached; i++) {
            for(int s = 0; s < nseq; s++) {
                const field<mat>& obs = mobs(s);
                const Labels& observed_segments = mobserved_segments(s);

                // Assertions if labels are provided.
                if (!observed_segments.empty()) {
                    for(const auto &segment : observed_segments.getLabels())
                        assert(segment.getDuration() >= min_duration_ &&
                                segment.getDuration() < min_duration_ +
                                ndurations_);
                    int start_time = observed_segments.getFirstSegment(
                            ).getStartingTime();
                    int end_time = observed_segments.getLastSegment(
                            ).getEndingTime();
                    if (start_time > 0)
                        assert(start_time >= min_duration_);
                    if (end_time < obs.n_elem - 1)
                        assert(obs.n_elem - end_time > min_duration_);
                }

                mat& alpha = malpha(s);
                mat& beta = mbeta(s);
                mat& alpha_s = malpha_s(s);
                mat& beta_s = mbeta_s(s);
                vec& beta_s_0 = mbeta_s_0(s);
                cube& eta = meta(s);
                cube& zeta = mzeta(s);

                // Recomputing the emission likelihoods.
                cube logpdf = computeEmissionsLogLikelihood(obs);

                logsFB(log_estimated_transition, log_estimated_pi,
                        log_estimated_duration, logpdf, observed_segments,
                        alpha, beta, alpha_s, beta_s, beta_s_0, eta, zeta,
                        min_duration_, obs.n_elem);
            }
            vec sequences_llikelihood(nseq);
            for(int s = 0; s < nseq; s++) {
                int nobs = mobs(s).n_elem;

                // Computing the marginal likelihood (aka observation
                // likelihood).
                sequences_llikelihood(s) = logsumexp(malpha(s).col(nobs - 1));
            }
            double current_llikelihood = sum(sequences_llikelihood);
            cout << "EM iteration " << i << " marginal log-likelihood: " <<
                    current_llikelihood << ". Diff: " <<
                    current_llikelihood - marginal_llikelihood << endl;
            if (current_llikelihood < marginal_llikelihood)
                cout << "Warning: The log-likelihood decreased probably due" <<
                        " to numerical errors." << endl;
            if (current_llikelihood - marginal_llikelihood < tol) {
                convergence_reached = true;
                marginal_llikelihood = current_llikelihood;
                break;
            }
            marginal_llikelihood = current_llikelihood;

            // M step.
            // Lower bound before M step.
            double lb_pi, lb_transition, lb_duration;
            if (debug_) {
                lb_pi = lower_bound_term_pi(meta, log_estimated_pi);
                lb_transition = lower_bound_term_transition(mzeta,
                        log_estimated_transition);
                lb_duration = lower_bound_term_duration(meta,
                        log_estimated_duration);
            }

            mat tmp_transition(log_estimated_transition);
            for(int i = 0; i < nstates_; i++) {
                vector<double> den;
                for(int j = 0; j < nstates_; j++) {
                    vector<double> num;
                    for(int s = 0; s < nseq; s++) {
                        const cube& zeta = mzeta(s);
                        for(int t = 0; t < mobs(s).n_elem - 1; t++) {
                            num.push_back(zeta(i, j, t));
                            den.push_back(zeta(i, j, t));
                        }
                    }
                    vec num_v(num);
                    if (num_v.n_elem > 0)
                        tmp_transition(i, j) = logsumexp(num_v);
                }
                vec den_v(den);
                if (den_v.n_elem > 0) {
                    double denominator = logsumexp(den_v);

                    // Handling the case when the transition probability is 0.
                    if (denominator != -datum::inf) {
                        for(int j = 0; j < nstates_; j++)
                            tmp_transition(i, j) -= denominator;
                    }
                }
            }
            if (learning_transitions_)
                log_estimated_transition = tmp_transition;

            // Reestimating the initial state pmf.
            vec tmp_pi(size(pi_), fill::zeros);
            for(const cube& eta : meta) {
                vec current_log_estimated_pi(nstates_);
                for(int i = 0; i < nstates_; i++) {
                    vector<double> terms;
                    for(int d = 0; d < ndurations_; d++)
                        terms.push_back(eta(i, d, min_duration_ + d - 1));
                    vec vterms(terms);
                    current_log_estimated_pi(i) = logsumexp(vterms);
                }
                vec current_pi = exp(current_log_estimated_pi);
                assert(abs(sum(current_pi) - 1) < 1e-7);
                tmp_pi += current_pi;
            }
            if (learning_pi_)
                log_estimated_pi = log(tmp_pi / nseq);

            // Reestimating durations.
            // D(j, d) represents the expected number of times that state
            // j is visited with duration d (non-normalized).
            mat D(size(duration_), fill::zeros);
            for(int i = 0; i < nstates_; i++) {
                vector<double> den;
                for(int d = 0; d < ndurations_; d++) {
                    vector<double> ts;
                    for(int s = 0; s < nseq; s++) {
                        int nobs = mobs(s).n_rows;
                        const cube& eta = meta(s);
                        for(int t = 0; t < nobs; t++) {
                            ts.push_back(eta(i, d, t));
                            den.push_back(eta(i, d, t));
                        }
                    }

                    // Taking into account the Dirichlet prior.
                    double log_alpha_minus_one = log(dirichlet_alphas_(
                                i, d) - 1);
                    ts.push_back(log_alpha_minus_one);
                    den.push_back(log_alpha_minus_one);

                    vec ts_v(ts);
                    D(i, d) = logsumexp(ts_v);
                }
                vec den_v(den);
                double denominator = logsumexp(den_v);

                // Handling the case when the transition probability mass is 0.
                if (denominator != -datum::inf) {
                    for(int d = 0; d < ndurations_; d++)
                        D(i, d) -= denominator;
                }
            }
            if (!duration_learning_choice_.compare("histogram"))
                log_estimated_duration = D;
            else if (!duration_learning_choice_.compare("momentmatching"))
                log_estimated_duration = gaussianMomentMatching(D);
            else if (duration_learning_choice_.compare("nodur")) {
                cout << "Duration learning not supported or invalid." << endl;
                assert(false);
            }

            // Lower bound after M step.
            if (debug_) {
                double lb_pi_after = lower_bound_term_pi(meta, log_estimated_pi);
                double lb_transition_after = lower_bound_term_transition(mzeta,
                        log_estimated_transition);
                double lb_duration_after = lower_bound_term_duration(meta,
                        log_estimated_duration);
                cout << "Diff lower bound before and after M step" << endl;
                cout << "pi: " << lb_pi_after - lb_pi << endl;
                cout << "transition: " << lb_transition_after - lb_transition << endl;
                cout << "duration: " << lb_duration_after - lb_duration << endl;
                assert(lb_pi_after - lb_pi > -tol);
                assert(lb_transition_after - lb_transition > -tol);
                assert(lb_duration_after - lb_duration > -tol);
            }

            // Reestimating emissions.
            // NOTE: the rest of the HSMM parameters are updated out of
            // this loop.
            emission_->reestimate(min_duration_, meta, mobs);
        }

        cout << "Stopped because of " << ((convergence_reached) ?
                "convergence." : "max iter.") << endl;

       // Updating the model parameters.
       setTransition(exp(log_estimated_transition));
       setPi(exp(log_estimated_pi));
       setDuration(exp(log_estimated_duration));
       return convergence_reached;
    }

    bool HSMM::fit(field<field<mat>> mobs, int max_iter, double tol) {
        field<Labels> mobserved_segments(mobs.n_elem);  // empty.
        return fit(mobs, mobserved_segments, max_iter, tol);
    }

    double HSMM::loglikelihood(const field<field<mat>>& mobs) {
        field<Labels> no_labels(mobs.n_elem);  // empty.
        return loglikelihood(mobs, no_labels);
    }

    // TODO: make this a const function.
    double HSMM::loglikelihood(const field<field<mat>>& mobs,
            const field<Labels>& labels) {
        int nseq = mobs.n_elem;
        assert(nseq >= 1);
        double ll = 0.0;
        for(int i = 0; i < nseq; i++) {
            const field<mat>& obs = mobs(i);
            int nobs = obs.n_elem;
            mat alpha =  zeros<mat>(nstates_, nobs);
            mat beta = zeros<mat>(nstates_, nobs);
            mat alpha_s = zeros<mat>(nstates_, nobs);
            mat beta_s = zeros<mat>(nstates_, nobs);
            vec beta_s_0 =  zeros<vec>(nstates_);
            cube eta = zeros<cube>(nstates_, ndurations_, nobs);
            cube zeta = zeros<cube>(nstates_, nstates_, nobs - 1);
            cube logpdf = computeEmissionsLogLikelihood(obs);
            mat ltransition = log(transition_);
            vec lpi = log(pi_);
            mat ldur = log(duration_);
            logsFB(ltransition, lpi, ldur, logpdf, labels(i), alpha, beta,
                    alpha_s, beta_s, beta_s_0, eta, zeta, min_duration_, nobs);
            ll += logsumexp(alpha.col(nobs - 1));
        }
        return ll;
    }

    // Computes the likelihoods w.r.t. the emission model.
    cube HSMM::computeEmissionsLikelihood(const field<mat>& obs) {
        return emission_->likelihoodCube(min_duration_, ndurations_, obs);
    }

    // Computes the loglikelihoods w.r.t. the emission model.
    cube HSMM::computeEmissionsLogLikelihood(const field<mat>& obs) {
        return emission_->loglikelihoodCube(min_duration_, ndurations_,
            obs);
    }

    // Returns a json representation of the model.
    nlohmann::json HSMM::to_stream() const {
        nlohmann::json ret;
        ret["nstates"] = nstates_;
        ret["min_duration"] = min_duration_;
        ret["ndurations"] = ndurations_;
        ret["initial_pmf"] = pi_;
        ret["emission_params"] = emission_->to_stream();

        // Taking care of the serialization of armadillo matrices.
        vector<vector<double>> transition_v, duration_v;
        for(int i = 0; i < nstates_; i++) {
            transition_v.push_back(conv_to<vector<double>>::from(
                    transition_.row(i)));
            duration_v.push_back(conv_to<vector<double>>::from(
                    duration_.row(i)));
        }
        ret["transition"] = transition_v;
        ret["duration"] = duration_v;
        return ret;
    }

    // Reads the parameters from a json file.
    void HSMM::from_stream(const nlohmann::json &params) {
        nstates_ = params.at("nstates");
        min_duration_ = params.at("min_duration");
        ndurations_ = params.at("ndurations");
        const nlohmann::json& emission_params = params.at("emission_params");
        emission_->from_stream(emission_params);

        // Parsing the armadillo matrices (transition, pi, duration).
        vector<double> initial_pmf_v = params.at("initial_pmf");
        vector<vector<double>> transition_v = params.at("transition");
        vector<vector<double>> duration_v = params.at("duration");
        vec pi = conv_to<vec>::from(initial_pmf_v);
        mat transition = zeros<mat>(nstates_, nstates_);
        mat duration = zeros<mat>(nstates_, ndurations_);
        for(int i = 0; i < nstates_; i++) {
            transition.row(i) = conv_to<rowvec>::from(transition_v.at(i));
            duration.row(i) = conv_to<rowvec>::from(duration_v.at(i));
        }
        setPi(pi);
        setTransition(transition);
        setDuration(duration);
    }

    pair<ivec,ivec> HSMM::computeViterbiStateDurationSequenceFromLabels(
            ivec seq) {
        assert(seq.n_elem > 0);
        int last_state = seq(0);
        int current_dur = 1;
        vector<int> hs, dur;
        for(int i = 1; i < seq.n_elem; i++) {
            if (seq(i) == last_state)
                current_dur++;
            else {
                hs.push_back(last_state);
                dur.push_back(current_dur);
                last_state = seq(i);
                current_dur = 1;
            }
        }
        hs.push_back(last_state);
        dur.push_back(current_dur);
        ivec hs_v = conv_to<ivec>::from(hs);
        ivec dur_v = conv_to<ivec>::from(dur);
        return make_pair(hs_v, dur_v);
    }

    mat HSMM::gaussianMomentMatching(mat duration) const {
        duration = exp(duration);
        mat ret = zeros(size(duration));
        for(int i = 0; i < duration.n_rows; i++) {
            double mean = 0;
            for(int j = 0; j < duration.n_cols; j++)
                mean += j * duration(i, j);
            double var = 0;
            for(int j = 0; j < duration.n_cols; j++)
                var += (j - mean) * (j - mean) * duration(i, j);
            for(int j = 0; j < duration.n_cols; j++) {
                double tmp = (j - mean);
                tmp = (tmp * tmp) / var;
                tmp = -0.5 * tmp;
                tmp = exp(tmp);
                ret(i, j) = tmp;
            }

            // Renormalizing.
            ret.row(i) = ret.row(i) / sum(ret.row(i));
        }
        return log(ret);
    }

    // Debugging functions.
    double HSMM::lower_bound_term_transition(const field<cube>& zetas,
            const mat& log_transition) const {
        double ret = 0;
        for(const cube& zeta : zetas)
            for(int i = 0; i < nstates_; i++)
                for(int j = 0; j < nstates_; j++)
                    for(int t = 0; t < zeta.n_slices; t++)
                        if (log_transition(i, j) > -datum::inf)
                            ret += exp(zeta(i, j, t)) * log_transition(i, j);
       return ret;
    }

    double HSMM::lower_bound_term_pi(const field<cube>& etas,
            const vec& log_pi) const {
        double ret = 0;
        for(const cube& eta : etas)
            for(int i = 0; i < nstates_; i++)
                for(int d = 0; d < ndurations_; d++)
                    if (log_pi(i) > -datum::inf)
                        ret += exp(eta(i, d, min_duration_ + d - 1)) *
                            log_pi(i);
        return ret;
    }

    double HSMM::lower_bound_term_duration(const field<cube>& etas,
            const mat& log_duration) const {
        double ret = 0;
        for(const cube& eta : etas)
            for(int i = 0; i < nstates_; i++)
                for(int d = 0; d < ndurations_; d++)
                    for(int t = 0; t < eta.n_slices; t++)
                        if (log_duration(i, d) > -datum::inf)
                            ret += exp(eta(i, d, t)) * log_duration(i, d);
        return ret;
    }

    /*
     * Online HSMM implementation
     */
    OnlineHSMM::OnlineHSMM(shared_ptr<AbstractEmissionOnlineSetting> emission,
            mat transition, vec pi, mat duration, int min_duration) : HSMM(
            static_pointer_cast<AbstractEmission>(emission), transition,
            pi, duration, min_duration),
            last_log_posterior_(ndurations_, min_duration_ + ndurations_,
            nstates_) {}

    shared_ptr<AbstractEmissionOnlineSetting> OnlineHSMM::getOnlineEmission(
            ) const {
        return static_pointer_cast<AbstractEmissionOnlineSetting>(emission_);
    }

    void OnlineHSMM::addNewObservation(const mat& obs) {
        if (observations_.empty())
             alpha_posteriors_.push_back(log(pi_));
        observations_.push_back(obs);
        mat log_duration = log(duration_);
        mat log_transition = log(transition_);
        last_log_posterior_.fill(-datum::inf);
        vector<double> normalization_terms;
        for(int d = min_duration_; d < min_duration_ + ndurations_; d++) {
            for(int s = 0; s < d; s++) {

                // Making sure the offset is consistent with the number of
                // observations so far.
                if (s >= observations_.size())
                    break;

                // Building the current segment. Notice that it is padded with
                // empty matrices.
                field<mat> current_segment(d);
                for(int idx = s, obs_idx = observations_.size()-1; idx >= 0;
                        idx--, obs_idx--) {
                    current_segment(idx) = observations_[obs_idx];
                }
                const vec& relevant_alpha_posterior = alpha_posteriors_.at(
                        alpha_posteriors_.size() - 1 - s);
                for(int i = 0; i < nstates_; i++) {
                    double log_pdf_seg = emission_->loglikelihood(i,
                            current_segment) + log_duration(i,d-min_duration_);
                    double log_unnormalized_value = log_pdf_seg +
                            relevant_alpha_posterior(i);
                    last_log_posterior_(d - min_duration_, s, i) =
                            log_unnormalized_value;
                    normalization_terms.push_back(log_unnormalized_value);
                }
            }
        }

        // Normalizing the current joint posterior.
        last_normalization_c_ = logsumexp(normalization_terms);
        assert(last_normalization_c_ > -datum::inf);
        last_log_posterior_ -= last_normalization_c_;

        // Computing the marginal posterior over the next hidden state assuming
        // there is a change point right after the current input observation.
        vec current_marginal_posterior(nstates_);
        for(int i = 0; i < nstates_; i++) {
            vector<double> terms;
            for(int j = 0; j < nstates_; j++) {
                for(int dur = 0; dur < ndurations_; dur++) {
                    double term = last_log_posterior_(dur,
                            min_duration_ + dur - 1, j) + log_transition(j, i);
                    terms.push_back(term);
                }
            }
            current_marginal_posterior(i) = logsumexp(terms);
        }

        // Pushing back the computed posterior. NOTE: it isn't a pmf.
        alpha_posteriors_.push_back(current_marginal_posterior);
    }

    void OnlineHSMM::sampleFromPosterior(int & dur, int & offset,
            int & hs) const {
        dur = offset = hs = -1;
        cube current_posterior = exp(last_log_posterior_);
        int nrows = current_posterior.n_rows;
        int ncols = current_posterior.n_cols;
        int nslices = current_posterior.n_slices;
        double accum = 0.0;
        double usample = randu();
        int found = 0;
        mat sample;
        for(int i = 0; i < nslices; i++) {
            for(int d = 0; d < nrows; d++) {
                for(int s = 0; s < ncols; s++) {
                    double naccum = accum + current_posterior(d, s, i);
                    if (accum < usample && usample < naccum) {
                        found++;
                        dur = d;
                        offset = s;
                        hs = i;
                    }
                    accum = naccum;
                }
            }
        }
        assert(found == 1);
        assert(abs(accum - 1.0) < 1e-7);
        return;
    }

    field<mat> OnlineHSMM::sampleNextObservations(int nobs) const {
        assert(!observations_.empty());  // TODO: handle empty observations.

        // Sampling a triplet (d,s,i) from the current posterior.
        int d, s, i;
        sampleFromPosterior(d, s, i);
        if (debug_) {
            cout << "Posterior sample (" << d << ", " << s << ", " <<
                    i << ")" << endl;
        }
        field<mat> last_obs(s + 1);
        int tam = observations_.size();
        for(int j = s, idx = tam - 1; j >= 0; j--, idx--)
            last_obs(j) = observations_.at(idx);
        field<mat> ret(nobs);
        int idx = 0;
        field<mat> seg;
        int curr_state;
        if (s == min_duration_ + d - 1) {

            // Handling the case when there is a transition
            // right after the current observation. Sampling
            // next state and duration.
            int next_state = sampleFromCategorical(
                    transition_.row(i));
            int next_duration = sampleFromCategorical(
                    duration_.row(next_state)) + min_duration_;
            seg = getOnlineEmission()->sampleFirstSegmentObsGivenLastSegment(
                    next_state, next_duration, last_obs, i);
            curr_state = next_state;
            for(int j = 0; j < seg.n_elem && idx < ret.n_elem; j++)
                ret(idx++) = seg(j);
        }
        else {
            seg = getOnlineEmission()->sampleNextObsGivenPastObs(i,
                    min_duration_ + d, last_obs);
            curr_state = i;
            field<mat> suffix = seg.rows(last_obs.n_elem, seg.n_elem - 1);
            idx = appendToField(ret, idx, suffix);
        }

        while(idx < ret.n_elem) {
             int next_state = sampleFromCategorical(transition_.row(
                         curr_state));
             int next_duration =  sampleFromCategorical(duration_.row(
                         next_state)) + min_duration_;
             if (debug_) {
                cout << "Next state and dur: " << next_state << " " <<
                        next_duration << endl;
             }
             seg = getOnlineEmission()->sampleFirstSegmentObsGivenLastSegment(
                     next_state, next_duration, seg, curr_state);
             curr_state = next_state;
             idx = appendToField(ret, idx, seg);
        }
        assert(idx == nobs);
        return ret;
    }

    void OnlineHSMM::printTopKFromPosterior(int k) const {
        cube current_posterior = exp(last_log_posterior_);
        set<pair<double,string>> pq;
        for(int d = 0; d < ndurations_; d++)
            for(int s = 0; s < min_duration_ + d; s++)
                for(int i = 0; i < nstates_; i++) {
                    double p = current_posterior(d, s, i);
                    string str = "(" + to_string(d) + ", " + to_string(s) +
                            ", " + to_string(i) + ")";
                    pq.insert(make_pair(p, str));
                    if (pq.size() > k)
                        pq.erase(pq.begin());
                }
        for(auto p: pq)
            cout << p.first << " " << p.second << endl;
    }

    vec OnlineHSMM::getStateMarginal() const {
        cube current_posterior = exp(last_log_posterior_);
        mat tmp = sum(current_posterior, 0);
        vec ret = vectorise(sum(tmp, 0));
        assert(abs(sum(ret) - 1.0) < 1e-7);
        assert(ret.n_elem == nstates_);
        return ret;
    }

    // TODO: the output len shoud be actually min_duration_ + ndurations_ - 1.
    vec OnlineHSMM::getRunlengthMarginal() const {
        cube current_posterior = exp(last_log_posterior_);
        mat tmp = sum(current_posterior, 0);
        vec ret = vectorise(sum(tmp, 1));
        assert(abs(sum(ret) - 1.0) < 1e-7);
        assert(ret.n_elem == min_duration_ + ndurations_);
        return ret;
    }

    vec OnlineHSMM::getDurationMarginal() const {
        cube current_posterior = exp(last_log_posterior_);
        mat tmp = sum(current_posterior, 2);
        vec ret = vectorise(sum(tmp, 1));
        assert(abs(sum(ret) - 1.0) < 1e-7);
        assert(ret.n_elem == ndurations_);
        return ret;
    }

    vec OnlineHSMM::getResidualTimeMarginal() const {
        cube current_posterior = exp(last_log_posterior_);
        mat tmp = sum(current_posterior, 2);
        vec ret(min_duration_ + ndurations_, fill::zeros);
        for(int d = 0; d < ndurations_; d++) {
            for(int r = 0; r < min_duration_ + ndurations_; r++) {
                if (r < min_duration_ + d) {
                    int residual_time = min_duration_ + d - r - 1;
                    ret(residual_time) += tmp(d, r);
                }
            }
        }
        assert(abs(sum(ret) - 1.0) < 1e-7);
        return ret;
    }

    vec OnlineHSMM::getImplicitDurationMarginal() const {
        cube current_posterior = exp(last_log_posterior_);
        mat runlength_state_posterior = sum(current_posterior, 0);

        vec ret(ndurations_, fill::zeros);
        mat dur_suffix_sum = getDurationSuffixSum();
        for(int i = 0; i < nstates_; i++) {
            for(int r = 0; r < min_duration_ + ndurations_; r++) {
                for(int d = 0; d < ndurations_; d++) {
                    if (r < min_duration_ + d) {
                        double value = runlength_state_posterior(r, i) *
                                duration_(i, d) /
                                dur_suffix_sum(i, max(0,r - min_duration_+ 1));
                        ret(d) += value;
                    }
                }
            }
        }
        assert(abs(sum(ret) - 1.0) < 1e-7);
        return ret;
    }

    vec OnlineHSMM::getImplicitResidualTimeMarginal() const {
        cube current_posterior = exp(last_log_posterior_);
        mat duration_runlength_posterior = sum(current_posterior, 2);
        vec ret(min_duration_ + ndurations_, fill::zeros);
        for(int r = 0; r < min_duration_ + ndurations_; r++) {
            for(int d = 0; d < ndurations_; d++) {
                int current_dur = min_duration_ + d;
                if (r < current_dur)
                    ret(current_dur-r-1) += duration_runlength_posterior(d,r);
            }
        }
        assert(abs(sum(ret) - 1.0) < 1e-7);
        return ret;
    }

    double OnlineHSMM::getLastOneStepAheadLoglikelihood() const {
        assert(observations_.size() > 0);
        return last_normalization_c_;
    }

    // TODO: This could be cached.
    mat OnlineHSMM::getDurationSuffixSum() const {
        mat suffix_sum(duration_);
        for(int i = suffix_sum.n_cols - 2; i >= 0; i--)
            suffix_sum.col(i) += suffix_sum.col(i + 1);
        return suffix_sum;
    }

    int OnlineHSMM::appendToField(field<mat> &current_obs, int idx,
            const field<mat> new_obs) const {
        for(int j = 0; j < new_obs.n_elem && idx < current_obs.n_elem; j++)
            current_obs(idx++) = new_obs(j);
        return idx;
    }

    /*
     * Online OnlineHSMMRunlengthBased implementation. The emission doesn't
     * take into account the total segment duration.
     */
    OnlineHSMMRunlengthBased::OnlineHSMMRunlengthBased(
            shared_ptr<AbstractEmissionObsCondIIDgivenState> emission,
            mat transition, vec pi, mat duration, int min_duration) : HSMM(
            static_pointer_cast<AbstractEmission>(emission), transition,
            pi, duration, min_duration),
            last_posterior_(min_duration + duration.n_cols - 1, pi.n_elem),
            last_residualtime_posterior_(min_duration + duration.n_cols - 1,
                    pi.n_elem) {
        init();
    }

    OnlineHSMMRunlengthBased::OnlineHSMMRunlengthBased(
            shared_ptr<AbstractEmissionObsCondIIDgivenState> emission,
            int nstates, int ndurations, int min_duration) : HSMM(
            static_pointer_cast<AbstractEmission>(emission), nstates,
            ndurations, min_duration),
            last_posterior_(min_duration + ndurations - 1, nstates),
            last_residualtime_posterior_(min_duration + ndurations - 1,
                    nstates) {
        init();
    }

    void OnlineHSMMRunlengthBased::init() {
        last_posterior_.zeros();
        last_residualtime_posterior_.zeros();
        observations_.clear();
    }

    shared_ptr<AbstractEmissionObsCondIIDgivenState>
            OnlineHSMMRunlengthBased::getOnlineDurationAgnosticEmission(
            ) const {
        return static_pointer_cast<AbstractEmissionObsCondIIDgivenState>(
                emission_);
    }

    void OnlineHSMMRunlengthBased::addNewObservation(const mat& obs) {
        observations_.push_back(obs);
        mat new_posterior(size(last_posterior_), fill::zeros);
        mat new_residualtime_posterior(size(last_residualtime_posterior_),
                fill::zeros);
        int max_duration = min_duration_ + ndurations_ - 1;
        mat hazard = getHazardFunction_();
        if (observations_.size() == 1) {

            // Base case.
            for(int i = 0; i < nstates_; i++)
                new_posterior(0, i) = pi_(i) * exp(loglikelihood_(i, obs));

            // Residential time base case.
            for(int i = 0; i < nstates_; i++) {
                double tmp = pi_(i) * exp(loglikelihood_(i, obs));
                for(int d = 0; d < ndurations_; d++)
                    new_residualtime_posterior(min_duration_ + d - 1, i) =
                            tmp * duration_(i, d);
            }
        }
        else {

            // Runlength posterior updates.
            for(int r = 0; r < max_duration; r++) {

                // Segment transition.
                for(int i = 0; i < nstates_; i++) {
                    if (!is_finite(hazard(i, r)))
                        continue;
                    for(int j = 0; j < nstates_; j++) {
                        new_posterior(0, j) += hazard(i, r)*transition_(i, j) *
                            exp(loglikelihood_(j, obs)) * last_posterior_(r, i);
                    }
                }

                // Segment continuation.
                if (r + 1 < max_duration) {
                    for(int i = 0; i < nstates_; i++)
                        if (is_finite(hazard(i, r))) {
                            new_posterior(r + 1, i) = last_posterior_(r, i) *
                                    (1 - hazard(i, r)) *
                                    exp(loglikelihood_(i, obs));
                        }
                }
            }

            // Residualtime posterior updates.
            // Segment continuation.
            for(int r = 1; r < max_duration; r++)
                for(int i = 0; i < nstates_; i++)
                    new_residualtime_posterior(r - 1, i) = exp(loglikelihood_(
                            i, obs)) * last_residualtime_posterior_(r, i);

            // Segment transition
            vec prob_of_transition_to(nstates_, fill::zeros);
            for(int i = 0; i < nstates_; i++)
                for(int j = 0; j < nstates_; j++)
                    prob_of_transition_to(i) += transition_(j, i) *
                            last_residualtime_posterior_(0, j);

            for(int i = 0; i < nstates_; i++)
                for(int d = 0; d < ndurations_; d++)
                    new_residualtime_posterior(min_duration_ + d - 1, i) +=
                            prob_of_transition_to(i) *
                            exp(loglikelihood_(i, obs)) * duration_(i, d);
        }

        // Normalizing the current posteriors.
        new_posterior = new_posterior / accu(new_posterior);
        new_residualtime_posterior = new_residualtime_posterior /
                accu(new_residualtime_posterior);

        assert(abs(accu(new_posterior) - 1.0) < 1e-7);
        assert(abs(accu(new_residualtime_posterior) - 1.0) < 1e-7);
        last_posterior_ = new_posterior;
        last_residualtime_posterior_ = new_residualtime_posterior;
        return;
    }

    vec OnlineHSMMRunlengthBased::getRunlengthMarginal() const {
        vec marginal = conv_to<vec>::from(sum(last_posterior_, 1));
        assert(abs(accu(marginal) - 1.0) < 1e-7);
        assert(marginal.n_elem == min_duration_ + ndurations_ - 1);
        return marginal;
    }

    vec OnlineHSMMRunlengthBased::getStateMarginal() const {
        vec marginal = conv_to<vec>::from(sum(last_posterior_, 0));
        assert(abs(accu(marginal) - 1.0) < 1e-7);
        assert(marginal.n_elem == nstates_);
        return marginal;
    }

    vec OnlineHSMMRunlengthBased::getResidualTimeMarginal() const {
        vec marginal = conv_to<vec>::from(sum(last_residualtime_posterior_,
                    1));
        assert(abs(accu(marginal) - 1.0) < 1e-7);
        assert(marginal.n_elem == min_duration_ + ndurations_ - 1);
        return marginal;
    }

    vec OnlineHSMMRunlengthBased::getStateMarginal2() const {
        vec marginal = conv_to<vec>::from(sum(last_residualtime_posterior_,
                    0));
        assert(abs(accu(marginal) - 1.0) < 1e-7);
        assert(marginal.n_elem == nstates_);
        return marginal;
    }

    double OnlineHSMMRunlengthBased::oneStepAheadLoglikelihood(
            const arma::mat& obs) const {
        if (observations_.size() == 0) {
            vec ret(nstates_);
            for(int i = 0; i < nstates_; i++)
                ret(i) = loglikelihood_(i, obs) + log(pi_(i));
            return logsumexp(ret);
        }
        int max_duration = min_duration_ + ndurations_ - 1;
        mat hazard = getHazardFunction_();
        vector<double> v;
        for(int r = 0; r < max_duration; r++) {

            // Segment transition.
            for(int i = 0; i < nstates_; i++)
                if (is_finite(hazard(i, r)))
                    for(int j = 0; j < nstates_; j++)
                        v.push_back(log(last_posterior_(r, i)) + log(
                                hazard(i, r)) + log(transition_(i, j)) +
                                loglikelihood_(j, obs));

            // Segment continuation.
            if (r > 0) {
                for(int i = 0; i < nstates_; i++)
                    if (is_finite(hazard(i, r - 1)))
                        v.push_back(log(last_posterior_(r - 1, i)) +
                            loglikelihood_(i, obs) + log(1 - hazard(i, r - 1)));
            }
        }
        vec ret = conv_to<vec>::from(v);
        return logsumexp(ret);
    }

    double OnlineHSMMRunlengthBased::oneStepAheadLoglikelihood2(
            const arma::mat& obs) const {
        if (observations_.size() == 0) {
            vec ret(nstates_);
            for(int i = 0; i < nstates_; i++)
                ret(i) = loglikelihood_(i, obs) + log(pi_(i));
            return logsumexp(ret);
        }
        int max_duration = min_duration_ + ndurations_ - 1;
        mat hazard = getHazardFunction_();
        vector<double> v;

        // Segment continuation.
        for(int r = 1; r < max_duration; r++)
            for(int i = 0; i < nstates_; i++)
                v.push_back(log(last_residualtime_posterior_(r, i)) +
                        loglikelihood_(i, obs));

        // Segment transition. Notice that the next bit is not log scaled.
        vec prob_of_transition_to(nstates_, fill::zeros);
        for(int i = 0; i < nstates_; i++)
            for(int j = 0; j < nstates_; j++)
                prob_of_transition_to(i) += transition_(j, i) *
                        last_residualtime_posterior_(0, j);
        for(int i = 0; i < nstates_; i++)
            v.push_back(log(prob_of_transition_to(i)) + loglikelihood_(i, obs));
        vec ret = conv_to<vec>::from(v);
        return logsumexp(ret);
    }

    double OnlineHSMMRunlengthBased::loglikelihood_(int state,
            const mat& obs) const {

        // TODO: assuming for now that the input can be converted into a vec.
        vec obs_v = conv_to<vec>::from(obs);
        getOnlineDurationAgnosticEmission()->loglikelihood(state, obs_v);
    }

    mat OnlineHSMMRunlengthBased::getDurationSuffixSum_() const {
        mat suffix_sum(duration_);
        for(int i = suffix_sum.n_cols - 2; i >= 0; i--)
            suffix_sum.col(i) += suffix_sum.col(i + 1);
        return suffix_sum;
    }

    mat OnlineHSMMRunlengthBased::getHazardFunction_() const {
        mat suffix_sum = getDurationSuffixSum_();
        mat hazard = duration_ / suffix_sum;
        if (min_duration_ > 1)
            hazard = join_horiz(zeros<mat>(nstates_,min_duration_ -1), hazard);
        assert(hazard.n_rows == nstates_ && hazard.n_cols == min_duration_
                + ndurations_ - 1);
        return hazard;
    }

};


#include <iostream>
#include <Eigen/Dense>
#include <Eigen/Geometry>

int main()
{
	Eigen::Vector3f acc, acc_result;

	Eigen::Matrix3f rot_x, rot_y, rot_z;
	rot_x = Eigen::AngleAxisf(0.1, Eigen::Vector3f(1.0, 0.0, 0.0));
	rot_y = Eigen::AngleAxisf(0.1, Eigen::Vector3f(0.0, 1.0, 0.0));
	rot_z = Eigen::AngleAxisf(0.1, Eigen::Vector3f(0.0, 0.0, 1.0));

	acc_result = ( rot_z * (rot_y * (rot_x * acc)));

	std::cout << acc_result << std::endl;

	return 0;
}