#include "lietorch_cpu.h"
#include <Eigen/Dense>

#include <iostream>
#include "common.h"
#include "dispatch.h"

#include "so3.h"
#include "rxso3.h"
#include "se3.h"
#include "sim3.h"

template <typename Group, typename scalar_t>
void exp_forward_kernel(const scalar_t* a_ptr, scalar_t* X_ptr, int batch_size) {
    // exponential map forward kernel
    using Tangent = Eigen::Matrix<scalar_t,Group::K,1>;
    using Data = Eigen::Matrix<scalar_t,Group::N,1>;

    at::parallel_for(0, batch_size, 1, [&](int64_t start, int64_t end) {
        for (int64_t i=start; i<end; i++) {
            Tangent a(a_ptr + i*Group::K);
            Eigen::Map<Data>(X_ptr + i*Group::N) = Group::Exp(a).data();
        }
    });
}

template <typename Group, typename scalar_t>
void exp_backward_kernel(const scalar_t* grad, const scalar_t* a_ptr, scalar_t* da, int batch_size) {
    // exponential map backward kernel
    using Tangent = Eigen::Matrix<scalar_t,Group::K,1>;
    using Grad = Eigen::Matrix<scalar_t,1,Group::K>;

    at::parallel_for(0, batch_size, 1, [&](int64_t start, int64_t end) {
        for (int64_t i=start; i<end; i++) {
            Tangent a(a_ptr + i*Group::K);
            Grad dX(grad + i*Group::N);
            Eigen::Map<Grad>(da + i*Group::K) = dX * Group::left_jacobian(a);
        }
    });
}

template <typename Group, typename scalar_t>
void log_forward_kernel(const scalar_t* X_ptr, scalar_t* a_ptr, int batch_size) {
    // logarithm map forward kernel
    using Tangent = Eigen::Matrix<scalar_t,Group::K,1>;

    at::parallel_for(0, batch_size, 1, [&](int64_t start, int64_t end) {
        for (int64_t i=start; i<end; i++) {
            Tangent a = Group(X_ptr + i*Group::N).Log();
            Eigen::Map<Tangent>(a_ptr + i*Group::K) = a;
        }
    });
}

template <typename Group, typename scalar_t>
void log_backward_kernel(const scalar_t* grad, const scalar_t* X_ptr, scalar_t* dX, int batch_size) {
    // logarithm map backward kernel
    using Tangent = Eigen::Matrix<scalar_t,Group::K,1>;
    using Grad = Eigen::Matrix<scalar_t,1,Group::K>;

    at::parallel_for(0, batch_size, 1, [&](int64_t start, int64_t end) {
        for (int64_t i=start; i<end; i++) {
            Tangent a = Group(X_ptr + i*Group::N).Log();
            Grad da(grad + i*Group::K);
            Eigen::Map<Grad>(dX + i*Group::N) = da * Group::left_jacobian_inverse(a);
        }
    });
}

template <typename Group, typename scalar_t>
void inv_forward_kernel(const scalar_t* X_ptr, scalar_t* Y_ptr, int batch_size) {
    // group inverse forward kernel
    using Tangent = Eigen::Matrix<scalar_t,Group::K,1>;
    using Data = Eigen::Matrix<scalar_t,Group::N,1>;

    at::parallel_for(0, batch_size, 1, [&](int64_t start, int64_t end) {
        for (int64_t i=start; i<end; i++) {
            Group X(X_ptr + i*Group::N);
            Eigen::Map<Data>(Y_ptr + i*Group::N) = X.inv().data();
        }
    });
}

template <typename Group, typename scalar_t>
void inv_backward_kernel(const scalar_t* grad, const scalar_t* X_ptr, scalar_t *dX, int batch_size) {
    // group inverse backward kernel
    using Tangent = Eigen::Matrix<scalar_t,Group::K,1>;
    using Grad = Eigen::Matrix<scalar_t,1,Group::K>;

    at::parallel_for(0, batch_size, 1, [&](int64_t start, int64_t end) {
        for (int64_t i=start; i<end; i++) {
            Group Y = Group(X_ptr + i*Group::N).inv();
            Grad dY(grad + i*Group::N);
            Eigen::Map<Grad>(dX + i*Group::N) = -dY * Y.Adj();
        }
    });
}

template <typename Group, typename scalar_t>
void mul_forward_kernel(const scalar_t* X_ptr, const scalar_t* Y_ptr, scalar_t* Z_ptr, int batch_size) {
    // group multiplication forward kernel
    using Tangent = Eigen::Matrix<scalar_t,Group::K,1>;
    using Data = Eigen::Matrix<scalar_t,Group::N,1>;

    at::parallel_for(0, batch_size, 1, [&](int64_t start, int64_t end) {
        for (int64_t i=start; i<end; i++) {
            Group Z = Group(X_ptr + i*Group::N) * Group(Y_ptr + i*Group::N);
            Eigen::Map<Data>(Z_ptr + i*Group::N) = Z.data();
        }
    });
}

template <class Group, typename scalar_t>
void mul_backward_kernel(const scalar_t* grad, const scalar_t* X_ptr, const scalar_t* Y_ptr, scalar_t* dX, scalar_t* dY, int batch_size) {
    // group multiplication backward kernel
    using Tangent = Eigen::Matrix<scalar_t,Group::K,1>;
    using Grad = Eigen::Matrix<scalar_t,1,Group::K>;
    using Data = Eigen::Matrix<scalar_t,Group::N,1>;

    at::parallel_for(0, batch_size, 1, [&](int64_t start, int64_t end) {
        for (int64_t i=start; i<end; i++) {
            Grad dZ(grad + i*Group::N);
            Group X(X_ptr + i*Group::N);
            Eigen::Map<Grad>(dX + i*Group::N) = dZ;
            Eigen::Map<Grad>(dY + i*Group::N) = dZ * X.Adj();
        }
    });
}

template <typename Group, typename scalar_t>
void adj_forward_kernel(const scalar_t* X_ptr, const scalar_t* a_ptr, scalar_t* b_ptr, int batch_size) {
    // adjoint forward kernel
    using Tangent = Eigen::Matrix<scalar_t,Group::K,1>;
    using Data = Eigen::Matrix<scalar_t,Group::N,1>;

    at::parallel_for(0, batch_size, 1, [&](int64_t start, int64_t end) {
        for (int64_t i=start; i<end; i++) {
            Group X(X_ptr + i*Group::N);
            Tangent a(a_ptr + i*Group::K);
            Eigen::Map<Tangent>(b_ptr + i*Group::K) = X.Adj(a);
        }
    });
}

template <typename Group, typename scalar_t>
void adj_backward_kernel(const scalar_t* grad, const scalar_t* X_ptr, const scalar_t* a_ptr, scalar_t* dX, scalar_t* da, int batch_size) {
    // adjoint backward kernel
    using Tangent = Eigen::Matrix<scalar_t,Group::K,1>;
    using Grad = Eigen::Matrix<scalar_t,1,Group::K>;
    using Data = Eigen::Matrix<scalar_t,Group::N,1>;

    at::parallel_for(0, batch_size, 1, [&](int64_t start, int64_t end) {
        for (int64_t i=start; i<end; i++) {
            Group X(X_ptr + i*Group::N);
            Grad db(grad + i*Group::K);

            Tangent a(a_ptr + i*Group::K);
            Tangent b = X.Adj() * a;

            Eigen::Map<Grad>(da + i*Group::K) = db * X.Adj();
            Eigen::Map<Grad>(dX + i*Group::N) = -db * Group::adj(b);
        }
    });
}

template <typename Group, typename scalar_t>
void adjT_forward_kernel(const scalar_t* X_ptr, const scalar_t* a_ptr, scalar_t* b_ptr, int batch_size) {
    // adjoint forward kernel
    using Tangent = Eigen::Matrix<scalar_t,Group::K,1>;
    using Data = Eigen::Matrix<scalar_t,Group::N,1>;

    at::parallel_for(0, batch_size, 1, [&](int64_t start, int64_t end) {
        for (int64_t i=start; i<end; i++) {
            Group X(X_ptr + i*Group::N);
            Tangent a(a_ptr + i*Group::K);
            Eigen::Map<Tangent>(b_ptr + i*Group::K) = X.AdjT(a);
        }
    });
}

template <typename Group, typename scalar_t>
void adjT_backward_kernel(const scalar_t* grad, const scalar_t* X_ptr, const scalar_t* a_ptr, scalar_t* dX, scalar_t* da, int batch_size) {
    // adjoint backward kernel
    using Tangent = Eigen::Matrix<scalar_t,Group::K,1>;
    using Grad = Eigen::Matrix<scalar_t,1,Group::K>;
    using Data = Eigen::Matrix<scalar_t,Group::N,1>;

    at::parallel_for(0, batch_size, 1, [&](int64_t start, int64_t end) {
        for (int64_t i=start; i<end; i++) {
            Group X(X_ptr + i*Group::N);
            Tangent db(grad + i*Group::K);
            Grad a(a_ptr + i*Group::K);

            Eigen::Map<Tangent>(da + i*Group::K) = X.Adj(db);
            Eigen::Map<Grad>(dX + i*Group::N) = -a * Group::adj(X.Adj(db));
        }
    });
}

template <typename Group, typename scalar_t>
void act_forward_kernel(const scalar_t* X_ptr, const scalar_t* p_ptr, scalar_t* q_ptr, int batch_size) {
    // action on point forward kernel
    using Tangent = Eigen::Matrix<scalar_t,Group::K,1>;
    using Data = Eigen::Matrix<scalar_t,Group::N,1>;
    using Point = Eigen::Matrix<scalar_t,3,1>;

    at::parallel_for(0, batch_size, 1, [&](int64_t start, int64_t end) {
        for (int64_t i=start; i<end; i++) {
            Group X(X_ptr + i*Group::N);
            Point p(p_ptr + i*3);
            Eigen::Map<Point>(q_ptr + i*3) = X * p;
        }
    });
}

template <typename Group, typename scalar_t>
void act_backward_kernel(const scalar_t* grad, const scalar_t* X_ptr, const scalar_t* p_ptr, scalar_t* dX, scalar_t* dp, int batch_size) {
    // adjoint backward kernel
    using Tangent = Eigen::Matrix<scalar_t,Group::K,1>;
    using Grad = Eigen::Matrix<scalar_t,1,Group::K>;
    using Point = Eigen::Matrix<scalar_t,3,1>;
    using PointGrad = Eigen::Matrix<scalar_t,1,3>;
    using Transformation = Eigen::Matrix<scalar_t,4,4>;

    at::parallel_for(0, batch_size, 1, [&](int64_t start, int64_t end) {
        for (int64_t i=start; i<end; i++) {
            Group X(X_ptr + i*Group::N);
            Point p(p_ptr + i*3);
            PointGrad dq(grad + i*3);

            Eigen::Map<PointGrad>(dp + i*3) = dq * X.Matrix().template block<3,3>(0,0);
            Eigen::Map<Grad>(dX + i*Group::N) = dq * Group::act_jacobian(X*p);
        }
    });
}

// template <typename Group, typename scalar_t>
// void tovec_backward_kernel(const scalar_t* grad, const scalar_t* X_ptr, scalar_t* dX, int batch_size) {
//     // group inverse forward kernel
//     using Data = Eigen::Matrix<scalar_t,Group::N,1>;
//     using Grad = Eigen::Matrix<scalar_t,1,Group::N>;

//     at::parallel_for(0, batch_size, 1, [&](int64_t start, int64_t end) {
//         for (int64_t i=start; i<end; i++) {
//             Group X(X_ptr + i*Group::N);
//             Grad g(grad + i*Group::N);
//             Eigen::Map<Grad>(dX + i*Group::N) = g * X.vec_jacobian();
//         }
//     });
// }

// template <typename Group, typename scalar_t>
// void fromvec_backward_kernel(const scalar_t* grad, const scalar_t* X_ptr, scalar_t* dX, int batch_size) {
//     // group inverse forward kernel
//     using Data = Eigen::Matrix<scalar_t,Group::N,1>;
//     using Grad = Eigen::Matrix<scalar_t,1,Group::N>;

//     at::parallel_for(0, batch_size, 1, [&](int64_t start, int64_t end) {
//         for (int64_t i=start; i<end; i++) {
//             Group X(X_ptr + i*Group::N);
//             Grad g(grad + i*Group::N);
//             Eigen::Map<Grad>(dX + i*Group::N) = g * X.vec_jacobian();
//         }
//     });
// }

template <typename Group, typename scalar_t>
void act4_forward_kernel(const scalar_t* X_ptr, const scalar_t* p_ptr, scalar_t* q_ptr, int batch_size) {
    // action on homogeneous point forward kernel
    using Tangent = Eigen::Matrix<scalar_t,Group::K,1>;
    using Data = Eigen::Matrix<scalar_t,Group::N,1>;
    using Point = Eigen::Matrix<scalar_t,4,1>;

    at::parallel_for(0, batch_size, 1, [&](int64_t start, int64_t end) {
        for (int64_t i=start; i<end; i++) {
            Group X(X_ptr + i*Group::N);
            Point p(p_ptr + i*4);
            Eigen::Map<Point>(q_ptr + i*4) = X.act4(p);
        }
    });
}

template <typename Group, typename scalar_t>
void act4_backward_kernel(const scalar_t* grad, const scalar_t* X_ptr, const scalar_t* p_ptr, scalar_t* dX, scalar_t* dp, int batch_size) {
    // action on homogeneous point backward kernel

    using Tangent = Eigen::Matrix<scalar_t,Group::K,1>;
    using Grad = Eigen::Matrix<scalar_t,1,Group::K>;
    using Point = Eigen::Matrix<scalar_t,4,1>;
    using PointGrad = Eigen::Matrix<scalar_t,1,4>;
    using Transformation = Eigen::Matrix<scalar_t,4,4>;

    at::parallel_for(0, batch_size, 1, [&](int64_t start, int64_t end) {
        for (int64_t i=start; i<end; i++) {
            Group X(X_ptr + i*Group::N);
            Point p(p_ptr + i*4);
            PointGrad dq(grad + i*4);

            Eigen::Map<PointGrad>(dp + i*4) = dq * X.Matrix4x4();
            const Point q = X.act4(p);
            Eigen::Map<Grad>(dX + i*Group::N) = dq * Group::act4_jacobian(q);
        }
    });
}

template <typename Group, typename scalar_t>
void as_matrix_forward_kernel(const scalar_t* X_ptr, scalar_t* T_ptr, int batch_size) {
    // group inverse forward kernel
    using Tangent = Eigen::Matrix<scalar_t,Group::K,1>;
    using Data = Eigen::Matrix<scalar_t,Group::N,1>;
    using Matrix4 = Eigen::Matrix<scalar_t,4,4,Eigen::RowMajor>;

    at::parallel_for(0, batch_size, 1, [&](int64_t start, int64_t end) {
        for (int64_t i=start; i<end; i++) {
            Group X(X_ptr + i*Group::N);
            Eigen::Map<Matrix4>(T_ptr + i*16) = X.Matrix4x4();
        }
    });
}

template <typename Group, typename scalar_t>
void orthogonal_projector_kernel(const scalar_t* X_ptr, scalar_t* P_ptr, int batch_size) {
    // group inverse forward kernel
    using Proj = Eigen::Matrix<scalar_t,Group::N,Group::N,Eigen::RowMajor>;
    at::parallel_for(0, batch_size, 1, [&](int64_t start, int64_t end) {
        for (int64_t i=start; i<end; i++) {
            Group X(X_ptr + i*Group::N);
            Eigen::Map<Proj>(P_ptr + i*Group::N*Group::N) = X.orthogonal_projector();
        }
    });
}

template <typename Group, typename scalar_t>
void jleft_forward_kernel(const scalar_t* X_ptr, const scalar_t* a_ptr, scalar_t* b_ptr, int batch_size) {
    // left-jacobian inverse action
    using Tangent = Eigen::Matrix<scalar_t,Group::K,1>;
    using Data = Eigen::Matrix<scalar_t,Group::N,1>;

    at::parallel_for(0, batch_size, 1, [&](int64_t start, int64_t end) {
        for (int64_t i=start; i<end; i++) {
            Group X(X_ptr + i*Group::N);
            Tangent a(a_ptr + i*Group::K);
            Tangent b = Group::left_jacobian_inverse(X.Log()) * a;
            Eigen::Map<Tangent>(b_ptr + i*Group::K) = b;
        }
    });
}

// unary operations

torch::Tensor exp_forward_cpu(int group_id, torch::Tensor a) {
    int batch_size = a.size(0);
    torch::Tensor X;

    DISPATCH_GROUP_AND_FLOATING_TYPES(group_id, a.type(), "exp_forward_kernel", ([&] {
        X = torch::zeros({batch_size, group_t::N}, a.options());
        exp_forward_kernel<group_t, scalar_t>(
            a.data_ptr<scalar_t>(),
            X.data_ptr<scalar_t>(),
            batch_size);
    }));

    return X;
}

std::vector<torch::Tensor> exp_backward_cpu(int group_id, torch::Tensor grad, torch::Tensor a) {
    int batch_size = a.size(0);
    torch::Tensor da = torch::zeros(a.sizes(), grad.options());

    DISPATCH_GROUP_AND_FLOATING_TYPES(group_id, a.type(), "exp_backward_kernel", ([&] {
        exp_backward_kernel<group_t, scalar_t>(
            grad.data_ptr<scalar_t>(),
            a.data_ptr<scalar_t>(),
            da.data_ptr<scalar_t>(),
            batch_size);
    }));

    return {da};
}

torch::Tensor log_forward_cpu(int group_id, torch::Tensor X) {
    int batch_size = X.size(0);
    torch::Tensor a;

    DISPATCH_GROUP_AND_FLOATING_TYPES(group_id, X.type(), "log_forward_kernel", ([&] {
        a = torch::zeros({batch_size, group_t::K}, X.options());
        log_forward_kernel<group_t, scalar_t>(
            X.data_ptr<scalar_t>(),
            a.data_ptr<scalar_t>(),
            batch_size);
    }));

    return a;
}

std::vector<torch::Tensor> log_backward_cpu(int group_id, torch::Tensor grad, torch::Tensor X) {
    int batch_size = X.size(0);
    torch::Tensor dX = torch::zeros(X.sizes(), grad.options());

    DISPATCH_GROUP_AND_FLOATING_TYPES(group_id, X.type(), "log_backward_kernel", ([&] {
        log_backward_kernel<group_t, scalar_t>(
            grad.data_ptr<scalar_t>(),
            X.data_ptr<scalar_t>(),
            dX.data_ptr<scalar_t>(),
            batch_size);
    }));

    return {dX};
}

torch::Tensor inv_forward_cpu(int group_id, torch::Tensor X) {
    int batch_size = X.size(0);
    torch::Tensor Y = torch::zeros_like(X);

    DISPATCH_GROUP_AND_FLOATING_TYPES(group_id, X.type(), "inv_forward_kernel", ([&] {
        inv_forward_kernel<group_t, scalar_t>(
            X.data_ptr<scalar_t>(),
            Y.data_ptr<scalar_t>(),
            batch_size);
    }));

    return Y;
}

std::vector<torch::Tensor> inv_backward_cpu(int group_id, torch::Tensor grad, torch::Tensor X) {
    int batch_size = X.size(0);
    torch::Tensor dX = torch::zeros(X.sizes(), grad.options());

    DISPATCH_GROUP_AND_FLOATING_TYPES(group_id, X.type(), "inv_backward_kernel", ([&] {
        inv_backward_kernel<group_t, scalar_t>(
            grad.data_ptr<scalar_t>(),
            X.data_ptr<scalar_t>(),
            dX.data_ptr<scalar_t>(),
            batch_size);
    }));

    return {dX};
}

// binary operations
torch::Tensor mul_forward_cpu(int group_id, torch::Tensor X, torch::Tensor Y) {
    int batch_size = X.size(0);
    torch::Tensor Z = torch::zeros_like(X);

    DISPATCH_GROUP_AND_FLOATING_TYPES(group_id, X.type(), "mul_forward_kernel", ([&] {
        mul_forward_kernel<group_t, scalar_t>(
            X.data_ptr<scalar_t>(),
            Y.data_ptr<scalar_t>(),
            Z.data_ptr<scalar_t>(),
            batch_size);
    }));

    return Z;
}

std::vector<torch::Tensor> mul_backward_cpu(int group_id, torch::Tensor grad, torch::Tensor X, torch::Tensor Y) {
    int batch_size = X.size(0);
    torch::Tensor dX = torch::zeros(X.sizes(), grad.options());
    torch::Tensor dY = torch::zeros(Y.sizes(), grad.options());

    DISPATCH_GROUP_AND_FLOATING_TYPES(group_id, X.type(), "mul_backward_kernel", ([&] {
        mul_backward_kernel<group_t, scalar_t>(
            grad.data_ptr<scalar_t>(),
            X.data_ptr<scalar_t>(),
            Y.data_ptr<scalar_t>(),
            dX.data_ptr<scalar_t>(),
            dY.data_ptr<scalar_t>(),
            batch_size);
    }));

    return {dX, dY};
}

torch::Tensor adj_forward_cpu(int group_id, torch::Tensor X, torch::Tensor a) {
    int batch_size = X.size(0);
    torch::Tensor b = torch::zeros(a.sizes(), a.options());

    DISPATCH_GROUP_AND_FLOATING_TYPES(group_id, X.type(), "adj_forward_kernel", ([&] {
        adj_forward_kernel<group_t, scalar_t>(
            X.data_ptr<scalar_t>(),
            a.data_ptr<scalar_t>(),
            b.data_ptr<scalar_t>(),
            batch_size);
    }));

    return b;
}

std::vector<torch::Tensor> adj_backward_cpu(int group_id, torch::Tensor grad, torch::Tensor X, torch::Tensor a) {
    int batch_size = X.size(0);
    torch::Tensor dX = torch::zeros(X.sizes(), grad.options());
    torch::Tensor da = torch::zeros(a.sizes(), grad.options());

    DISPATCH_GROUP_AND_FLOATING_TYPES(group_id, X.type(), "adj_backward_kernel", ([&] {
        adj_backward_kernel<group_t, scalar_t>(
            grad.data_ptr<scalar_t>(),
            X.data_ptr<scalar_t>(),
            a.data_ptr<scalar_t>(),
            dX.data_ptr<scalar_t>(),
            da.data_ptr<scalar_t>(),
            batch_size);
    }));

    return {dX, da};
}

torch::Tensor adjT_forward_cpu(int group_id, torch::Tensor X, torch::Tensor a) {
    int batch_size = X.size(0);
    torch::Tensor b = torch::zeros(a.sizes(), a.options());

    DISPATCH_GROUP_AND_FLOATING_TYPES(group_id, X.type(), "adjT_forward_kernel", ([&] {
        adjT_forward_kernel<group_t, scalar_t>(
            X.data_ptr<scalar_t>(),
            a.data_ptr<scalar_t>(),
            b.data_ptr<scalar_t>(),
            batch_size);
    }));

    return b;
}

std::vector<torch::Tensor> adjT_backward_cpu(int group_id, torch::Tensor grad, torch::Tensor X, torch::Tensor a) {
    int batch_size = X.size(0);
    torch::Tensor dX = torch::zeros(X.sizes(), grad.options());
    torch::Tensor da = torch::zeros(a.sizes(), grad.options());

    DISPATCH_GROUP_AND_FLOATING_TYPES(group_id, X.type(), "adjT_backward_kernel", ([&] {
        adjT_backward_kernel<group_t, scalar_t>(
            grad.data_ptr<scalar_t>(),
            X.data_ptr<scalar_t>(),
            a.data_ptr<scalar_t>(),
            dX.data_ptr<scalar_t>(),
            da.data_ptr<scalar_t>(),
            batch_size);
    }));

    return {dX, da};
}

torch::Tensor act_forward_cpu(int group_id, torch::Tensor X, torch::Tensor p) {
    int batch_size = X.size(0);
    torch::Tensor q = torch::zeros(p.sizes(), p.options());

    DISPATCH_GROUP_AND_FLOATING_TYPES(group_id, X.type(), "act_forward_kernel", ([&] {
        act_forward_kernel<group_t, scalar_t>(
            X.data_ptr<scalar_t>(),
            p.data_ptr<scalar_t>(),
            q.data_ptr<scalar_t>(),
            batch_size);
    }));

    return q;
}

std::vector<torch::Tensor> act_backward_cpu(int group_id, torch::Tensor grad, torch::Tensor X, torch::Tensor p) {
    int batch_size = X.size(0);
    torch::Tensor dX = torch::zeros(X.sizes(), grad.options());
    torch::Tensor dp = torch::zeros(p.sizes(), grad.options());

    DISPATCH_GROUP_AND_FLOATING_TYPES(group_id, X.type(), "act_backward_kernel", ([&] {
        act_backward_kernel<group_t, scalar_t>(
            grad.data_ptr<scalar_t>(),
            X.data_ptr<scalar_t>(),
            p.data_ptr<scalar_t>(),
            dX.data_ptr<scalar_t>(),
            dp.data_ptr<scalar_t>(),
            batch_size);
    }));

    return {dX, dp};
}

torch::Tensor act4_forward_cpu(int group_id, torch::Tensor X, torch::Tensor p) {
    int batch_size = X.size(0);
    torch::Tensor q = torch::zeros(p.sizes(), p.options());

    DISPATCH_GROUP_AND_FLOATING_TYPES(group_id, X.type(), "act4_forward_kernel", ([&] {
        act4_forward_kernel<group_t, scalar_t>(
            X.data_ptr<scalar_t>(),
            p.data_ptr<scalar_t>(),
            q.data_ptr<scalar_t>(),
            batch_size);
    }));

    return q;
}

std::vector<torch::Tensor> act4_backward_cpu(int group_id, torch::Tensor grad, torch::Tensor X, torch::Tensor p) {
    int batch_size = X.size(0);
    torch::Tensor dX = torch::zeros(X.sizes(), grad.options());
    torch::Tensor dp = torch::zeros(p.sizes(), grad.options());

    DISPATCH_GROUP_AND_FLOATING_TYPES(group_id, X.type(), "act4_backward_kernel", ([&] {
        act4_backward_kernel<group_t, scalar_t>(
            grad.data_ptr<scalar_t>(),
            X.data_ptr<scalar_t>(),
            p.data_ptr<scalar_t>(),
            dX.data_ptr<scalar_t>(),
            dp.data_ptr<scalar_t>(),
            batch_size);
    }));

    return {dX, dp};
}

torch::Tensor as_matrix_forward_cpu(int group_id, torch::Tensor X) {
    int batch_size = X.size(0);
    torch::Tensor T4x4 = torch::zeros({X.size(0), 4, 4}, X.options());

    DISPATCH_GROUP_AND_FLOATING_TYPES(group_id, X.type(), "as_matrix_forward_kernel", ([&] {
        as_matrix_forward_kernel<group_t, scalar_t>(
            X.data_ptr<scalar_t>(),
            T4x4.data_ptr<scalar_t>(),
            batch_size);
    }));

    return T4x4;
}

torch::Tensor orthogonal_projector_cpu(int group_id, torch::Tensor X) {
    int batch_size = X.size(0);
    torch::Tensor P;

    DISPATCH_GROUP_AND_FLOATING_TYPES(group_id, X.type(), "orthogonal_projector_kernel", ([&] {
        P = torch::zeros({X.size(0), group_t::N, group_t::N}, X.options());
        orthogonal_projector_kernel<group_t, scalar_t>(X.data_ptr<scalar_t>(), P.data_ptr<scalar_t>(), batch_size);
    }));

    return P;
}

torch::Tensor jleft_forward_cpu(int group_id, torch::Tensor X, torch::Tensor a) {
    int batch_size = X.size(0);
    torch::Tensor b = torch::zeros(a.sizes(), a.options());

    DISPATCH_GROUP_AND_FLOATING_TYPES(group_id, X.type(), "jleft_forward_kernel", ([&] {
        jleft_forward_kernel<group_t, scalar_t>(
            X.data_ptr<scalar_t>(),
            a.data_ptr<scalar_t>(),
            b.data_ptr<scalar_t>(),
            batch_size);
    }));

    return b;
}


#include "urrp.h"
#include <boost/math/special_functions/digamma.hpp>
using namespace std;

inline double square(double x)
{
  return x * x;
}

inline double digamma(double x)
{
  if(x==0.0)
  {
    return 0.0;
  } else {
    return boost::math::digamma(x);
  }
}

void URRP::init_model()
{
  for(vector<rating*>::iterator it=trainratings.begin(); it != trainratings.end(); it++)
  {
    int user = (*it)->user;
    int item = (*it)->item;
    int rating = (*it)->value;
    vector< pair<int,int> >* words = &((*it)->words);

    //init attitude and topic
    int attitude = rand() % K;
    (*it)->attitude = attitude;
    (*muk)[user][attitude]++;
    (*ckvs)[attitude][item][rating]++;
    (*mu)[user]++;
    (*ckv)[attitude][item]++;

    int topic = rand() % K;
    for(vector< pair<int,int> >::iterator it2=words->begin(); it2 != words->end(); it2++)
    {
      (*it2).second = topic;
      (*nuk)[user][topic]++;
      (*nkw)[topic][(*it2).first]++;
      (*nu)[user]++;
      (*nk)[topic]++;
      topic = rand() % K;
    }
  }
}

void URRP::sample_attitudes(bool with_topic)
{
  double* p = new double[K];
  for(vector<rating*>::iterator it=trainratings.begin(); it != trainratings.end(); it++)
  {
    int user = (*it)->user;
    int item = (*it)->item;
    int rating = (*it)->value;

    //delete related attitude and topic for this rating;
    int current_attitude = (*it)->attitude;
    (*muk)[user][current_attitude]--;
    (*ckvs)[current_attitude][item][rating]--;
    (*mu)[user]--;
    (*ckv)[current_attitude][item]--;

    int k=0;
    //#pragma omp parallel for
    for(k=0; k<K; k++)
    {
      if(with_topic)
        p[k] = ((*nuk)[user][k] + (*muk)[user][k] + alpha[k]) / ((*nu)[user] + (*mu)[user] + sum_alpha) * (((*ckvs)[k][item][rating] + lambda[rating]) / ((*ckv)[k][item] + sum_lambda));
      else
        p[k] = ((*muk)[user][k] + alpha[k]) / ((*mu)[user] + sum_alpha) * (((*ckvs)[k][item][rating] + lambda[rating]) / ((*ckv)[k][item] + sum_lambda));
    }

    for(k=1; k<K; k++)
      p[k] += p[k-1];

    double randdouble = (rand() * 1.0 / RAND_MAX) * p[K-1];
    for(k=0; k<K; k++)
    {
      if(randdouble <= p[k])
      {
        break;
      }
    }
    (*it)->attitude = k;
    (*muk)[user][k]++;
    (*ckvs)[k][item][rating]++;
    (*mu)[user]++;
    (*ckv)[k][item]++;
  }
  delete[] p;
}

void URRP::sample_topics(bool with_attitude)
{
  double* p = new double[K];
  for(vector<rating*>::iterator it=trainratings.begin(); it != trainratings.end(); it++)
  {
    int user = (*it)->user;
    vector< pair<int,int> >* words = &((*it)->words);

    for(vector< pair<int,int> >::iterator it2=words->begin(); it2 != words->end(); it2++)
    {
      int w = (*it2).first;
      int current_topic = (*it2).second;
      (*nuk)[user][current_topic]--;
      (*nkw)[current_topic][w]--;
      (*nu)[user]--;
      (*nk)[current_topic]--;
      int k=0;
      for(k=0; k<K; k++)
      {
        if(with_attitude)
          p[k] = (((*nuk)[user][k] + (*muk)[user][k] + alpha[k]) / ((*nu)[user] + (*mu)[user] + sum_alpha) * (((*nkw)[k][w] + beta[w]) / ((*nk)[k] + sum_beta)));
        else
          p[k] = (((*nuk)[user][k] + alpha[k]) / ((*nu)[user] + sum_alpha) * (((*nkw)[k][w] + beta[w]) / ((*nk)[k] + sum_beta)));
      }
      for(k=1; k<K; k++)
        p[k] += p[k-1];

      double randdouble = (rand() * 1.0 / RAND_MAX) * p[K-1];
      for(k=0; k<K; k++)
      {
        if(randdouble <= p[k])
        {
          break;
        }
      }
      (*it2).second = k;
      (*nuk)[user][k]++;
      (*nkw)[k][w]++;
      (*nu)[user]++;
      (*nk)[k]++;
    }
  }
  delete[] p;
}

void URRP::readout_topic_theta(bool with_attitude)
{
  int u,k;
#pragma omp parallel for collapse(2)
  for(u=0; u<nUsers; u++)
    for(k=0; k<K; k++)
      if(with_attitude)
        (*theta)[u][k] = ((*nuk)[u][k] + (*muk)[u][k] + alpha[k]) / ((*nu)[u] + (*mu)[u] + sum_alpha);
      else
        (*theta)[u][k] = ((*nuk)[u][k] + alpha[k]) / ((*nu)[u] + sum_alpha);
}

void URRP::readout_attitude_theta(bool with_topic)
{
  int u,k;
#pragma omp parallel for collapse(2)
  for(u=0; u<nUsers; u++)
    for(k=0; k<K; k++)
      if(with_topic)
        (*theta)[u][k] = ((*nuk)[u][k] + (*muk)[u][k] + alpha[k]) / ((*nu)[u] + (*mu)[u] + sum_alpha);
      else
        (*theta)[u][k] = ((*muk)[u][k] + alpha[k]) / ((*mu)[u] + sum_alpha);
}

void URRP::readout_phi()
{
  int k,w;
#pragma omp parallel for collapse(2)
  for(k=0; k<K; k++)
    for(w=0; w<nWords; w++)
      (*phi)[k][w] = ((*nkw)[k][w] + beta[w]) / ((*nk)[k] + sum_beta);
}

void URRP::readout_xi()
{
  int v,k,s;
#pragma omp parallel for collapse(3)
  for(k=0; k<K; k++)
    for(v=0; v<nItems; v++)
      for(s=0; s<S; s++)
        (*xi)[k][v][s] = ((*ckvs)[k][v][s] + lambda[s]) / ((*ckv)[k][v] + sum_lambda);
}

void URRP::update_alpha_by_topic()
{
  int k,u;
  double ak, numerator, denominator;
  //update alpha
  for (k = 0; k < K; k++) {
    ak = alpha[k];
    numerator = 0, denominator = 0;
    #pragma omp parallel for reduction (+:numerator,denominator)
    for (u = 0; u < nUsers; u++) {
      numerator += digamma((*nuk)[u][k] + ak) - digamma(ak);
      denominator += digamma((*nu)[u] + sum_alpha) - digamma(sum_alpha);
      //numerator += digamma((*nuk)[u][k] + (*muk)[u][k] + ak) - digamma(ak);
      //denominator += digamma(get_nu(u) + get_mu(u) + sum_alpha) - digamma(sum_alpha);
    }
    if (numerator != 0)
      alpha[k] = ak * (numerator / denominator);
  }
  sum_alpha = 0;
  for (k = 0; k < K; k++)
  {
    sum_alpha += alpha[k];
  }
}

void URRP::update_alpha(bool with_topic)
{
  int k,u;
  double ak, numerator, denominator;
  //update alpha
  for (k = 0; k < K; k++) {
    ak = alpha[k];
    numerator = 0, denominator = 0;
    #pragma omp parallel for reduction (+:numerator,denominator)
    for (u = 0; u < nUsers; u++) {
      if (with_topic)
      {
        numerator += digamma((*nuk)[u][k] + (*muk)[u][k] + ak) - digamma(ak);
        denominator += digamma((*nu)[u] + (*mu)[u] + sum_alpha) - digamma(sum_alpha);
      }
      else
      {
        numerator += digamma((*muk)[u][k] + ak) - digamma(ak);
        denominator += digamma((*mu)[u] + sum_alpha) - digamma(sum_alpha);
      }
    }
    if (numerator != 0)
      alpha[k] = ak * (numerator / denominator);
  }
  sum_alpha = 0;
  for (k = 0; k < K; k++)
  {
    sum_alpha += alpha[k];
  }
}

void URRP::update_beta()
{
  int k,w;
  double betaw, numerator, denominator;
  //update beta
  for (w = 0; w < nWords; w++) {
    betaw = beta[w];
    numerator = 0, denominator = 0;
    //#pragma omp parallel for reduction (+:numerator,denominator)
    for (k = 0; k < K; k++) {
      numerator += digamma((*nkw)[k][w] + betaw) - digamma(betaw);
      denominator += digamma((*nk)[k] + sum_beta) - digamma(sum_beta);
    }
    if (numerator != 0)
      beta[w] = betaw * (numerator / denominator);
  }
  double tmp_sum_beta = 0;
  //#pragma omp parallel for reduction (+:tmp_sum_beta)
  for (w = 0; w < nWords; w++)
  {
    tmp_sum_beta += beta[w];
  }
  sum_beta = tmp_sum_beta;

}

void URRP::update_lambda()
{
  int v,k,s;
  double lambdas, numerator, denominator;
  //update lambda
  for (s = 0; s < S; s++) {
    lambdas = lambda[s];
    numerator = 0, denominator = 0;
    #pragma omp parallel for collapse(2) reduction (+:numerator,denominator)
    for(k=0; k<K; k++)
      for(v=0; v<nItems; v++)
      {
        numerator += digamma((*ckvs)[k][v][s] + lambdas) - digamma(lambdas);
        denominator += digamma((*ckv)[k][v] + sum_lambda) - digamma(sum_lambda);
      }
    if (numerator != 0)
      lambda[s] = lambdas * (numerator / denominator);
  }
  sum_lambda = 0;
  for (s = 0; s < S; s++)
  {
    sum_lambda += lambda[s];
  }
}

void URRP::evaluate(int iter)
{
  double train_err = 0.0;
  double validate_err = 0.0;
  double test_err = 0.0;
  double test_ste = 0.0;
  size_t i;

  #pragma omp parallel for reduction (+:validate_err)
  for(i=0; i<trainratings.size(); i++)
  {
    train_err += square(predict_with_expect(trainratings[i]) - trainratings[i]->value - 1);
  }

  #pragma omp parallel for reduction (+:validate_err)
  for(i=0; i<validratings.size(); i++)
  {
    validate_err += square(predict_with_expect(validratings[i]) - validratings[i]->value - 1);
  }

  #pragma omp parallel for reduction (+:test_err,test_ste)
  for(i=0; i<testratings.size(); i++)
  {
    double err = square(predict_with_expect(testratings[i]) - testratings[i]->value - 1);
    test_err += err;
    test_ste += err*err;
  }
  train_err /= trainratings.size();
  validate_err /= validratings.size();
  test_err /= testratings.size();
  test_ste /= testratings.size();
  test_ste = sqrt((test_ste-test_err*test_err)/testratings.size());

  if(test_err < current_best)
  {
    current_best = test_err;
    current_best_ste = test_ste;
  }

  double delta = validate_err - prev_mse;
  printf("\nRecommend: Iter %d, Train MSE: %.4lf, Validation MSE: %.4lf, Test MSE: %.4lf (%.2lf), validate_mse_delta: %.4lf\n", iter, train_err, validate_err, test_err, test_ste, delta);
  printf("\nCurrent best MSE:\t%.4lf (%.2lf)\n", current_best, current_best_ste);
  prev_mse = validate_err;
  fflush(stdout);
}

/// Train a model
void URRP::train()
{
  init_model();
  readout_attitude_theta(false);
  readout_xi();
  evaluate(0);
  // learn topic distribution by lda
  for(int i=1; i<=burn_in; i++)
  {
    sample_topics(false);
    if (i % 10 == 0)
    {
      printf("LDA stage: iter %d\n", i);
      fflush(stdout);
    }
  }
  readout_topic_theta(false);
  readout_phi();
  topic_words();
  for (int iter = 0; iter <= max_iter; iter++) {
    // sample topic and attitude for all words and ratings
    sample_attitudes(false);
    sample_topics(false);
    // update hyper-parameters
    update_alpha(true);
    update_beta();
    update_lambda();
    // get statistics after burn-in
    if (iter % sample_lag == 0)
    {
      readout_attitude_theta(false);
      readout_phi();
      readout_xi();
      evaluate(iter);
    }
  }
  topic_words();
}

bool word_prob_com(pair<int, double> p1, pair<int, double> p2)
{
  return p1.second > p2.second;
}

void URRP::topic_words()
{
  map<int, string> * id2word = &(corp->id2word);
  for(int k=0; k<K; k++)
  {
    vector< pair<int, double> > topic_words;
    for(int w=0; w<nWords; w++)
    {
      topic_words.push_back(make_pair(w, (*phi)[k][w]));
    }
    sort(topic_words.begin(), topic_words.end(), word_prob_com);
    printf("\nTopic %d:", k+1);
    for(int i=0; i<10; i++)
    {
      printf(" %s(%.4lf)", (*id2word)[topic_words[i].first].c_str(), (*phi)[k][topic_words[i].first]);
    }
  }
  printf("\n");
  fflush(stdout);
}

// Predict a particular rating given the current parameter values
double URRP::predict_with_expect(rating* vi)
{
  int user = vi->user;
  int item = vi->item;
  double pred = 0.0;
  double ps = 0.0;
  for(int s=0; s<S; s++)
  {
    ps = 0.0;
    for(int k=0; k<K; k++)
    {
      ps += (*theta)[user][k] * (*xi)[k][item][s];
    }
    pred += ps * (s+1);
  }
  return pred;
}

// Predict a particular rating given the current parameter values
double URRP::predict_with_most_prob(rating* vi)
{
  int user = vi->user;
  int item = vi->item;
  double pred = 0.0;
  double most_prob = 0.0;
  double ps = 0.0;
  for(int s=0; s<S; s++)
  {
    ps = 0.0;
    for(int k=0; k<K; k++)
    {
      ps += (*theta)[user][k] * (*xi)[k][item][s];
    }
    if(ps > most_prob)
    {
      pred = s+1;
      most_prob = ps;
    }
  }
  return pred;
}


#include "policy.hpp"
#include "helpers.hpp"
#include <ctime>
#include <sys/time.h>
#include <boost/math/distributions.hpp>
#include <stdlib.h>

template <class Policy>
double timePolicy(Policy& pol, int arms, int horizon, bool gaussian, RandUnif& urandom, RandNorm& nrandom)
{

  int trials = 1;
  static double mean_rewards[] = {0.6, 0.1, 0.7, 0.45, 0.43, 0.32, 0.8, 0.93, 0.4, 0.5};
  clock_t timing = clock();
  for (int tr= 1; tr <= trials; ++tr)
  {
    Table rewards(arms, vector<double>(horizon, 0));
    for (int a = 0; a < arms; ++a)
    {
      double mu = mean_rewards[a];
      for (int t = 0; t < horizon; ++t)
      {
        rewards[a][t] = gaussian ? (mu + nrandom()) : (urandom() < mu ? 1 : 0);
      }
    }
    for (int t = 1; t <= horizon; ++t)
    {
      int arm = pol.choosearm(t);

      pol.totalsampled[arm] += 1;
      double outcome = rewards[arm][pol.sample_index[arm]];
      pol.totalrewards[arm] += outcome;
      pol.sample_index[arm]++;
      pol.totalscore += outcome;
    }
  }
  timing = clock() - timing;
  return ((double)timing)/CLOCKS_PER_SEC;
}

int main(int argc, char** argv)
{
  timeval tv;
  gettimeofday(&tv,NULL);

  boost::mt19937 seed(tv.tv_sec);
  boost::uniform_real<> dist(0.0,1.0);
  boost::normal_distribution<> normdist(0.0,1.0);
  boost::uniform_int<> idist(1,1000000);
  RandomSampleHelper bh(tv.tv_sec);
  RandUnif urandom(seed,dist);
  RandNorm nrandom(seed,normdist);
  int horizon = 1000;
  int arms = 10;

  vector<double> uniforms(horizon, 0);
  IDSPolicy idsBer(horizon, uniforms,  arms, false, 1000);
  IDSPolicy idsGauss(horizon, uniforms, arms, true, 1000);
  InterestingPolicy nitpBer(horizon, arms, 1, 0, false);
  InterestingPolicy nitpBer3(horizon, arms, 1, 0, false, 3);
  InterestingPolicy nitpGauss(horizon, arms, 1, 0, true);

  ThompsonPolicy tompBer(bh, horizon, arms, false);
  ThompsonPolicy tompGauss(bh, horizon, arms, true);
  BayesUCBPolicy bucbBer(horizon, arms, false);
  BayesUCBPolicy bucbGauss(horizon, arms, true);

  cout << "TOM:" << timePolicy(tompBer, arms, horizon, false, urandom, nrandom) << endl;
  cout << "BUCB:" << timePolicy(bucbBer, arms, horizon, false, urandom, nrandom) << endl;
  cout << "IDS:" << timePolicy(idsBer, arms, horizon, false, urandom, nrandom) << endl;
  cout << "OURS-1:" << timePolicy(nitpBer, arms, horizon, false, urandom, nrandom) << endl;
  cout << "OURS-3:" <<  timePolicy(nitpBer3, arms, horizon, false, urandom, nrandom) << endl;

  cout << endl << "Gaussian!" << endl;
  cout << "TOM:" << timePolicy(tompGauss, arms, horizon, true, urandom, nrandom) << endl;
  cout << "BUCB:" << timePolicy(bucbGauss, arms, horizon, true, urandom, nrandom) << endl;
  cout << "IDS:" << timePolicy(idsGauss, arms, horizon, true, urandom, nrandom) << endl;
  cout << "OURS-1:" << timePolicy(nitpGauss, arms, horizon, true, urandom, nrandom) << endl;
}


#include <iostream>
#include <fstream>
#include <Eigen/Eigen>
//include the bie header files
#include "material.hh"
#include "precomputed_kernel.hh"
#include "bimat_interface.hh"
#include "infinite_boundary.hh"
//include the fem header files
#include "mesh_Generated_multi_faults.hpp"
#include "bcdof.hpp"
#include "bcdof_ptr.hpp"
#include "cal_ke.hpp"
#include "cal_fe_global_const_ke.hpp"
#include "mapglobal.hpp"
#include "Slip_Weakening_reg.hpp"
#include "Slip_Weakening.hpp"
#include "cal_slip_sliprate.hpp"
#include "time_advance.hpp"
#include "BIE_correct.hpp"
#include "igl/list_to_matrix.h"
#include <omp.h>
using namespace Eigen;
using namespace std;

typedef Eigen::Matrix<int, -1, -1,RowMajor> MatrixXi_rm;

void read_matrix(std::string fileName, Eigen::MatrixXd &outputMat);

double time_fem=0;
double time_bie;
int main() {
     time_fem = 0.0;
    // Domain Size
    double x_min = -30e3;
    double x_max = 30e3;
    double y_min = -0.5e3;
    double y_max = 0.5e3;
    int dim = 2.0;
    double dx = 3.125;
    double dy = 3.125;
    int nx_el = (x_max-x_min)/dx;
    int ny = (y_max-y_min)/dy;
    MatrixXd Node = MatrixXd::Zero((nx_el+1)*(ny+1),2);
    MatrixXi_rm Element(nx_el*ny,4); Element.setZero();
    ArrayXi BIE_top_surf_nodes = ArrayXi::Zero((nx_el+1),1);
    ArrayXi BIE_bot_surf_nodes = ArrayXi::Zero((nx_el+1),1);
    int num_faults = 5;
    // Mesh
    // Position of the fault : y_pos, x_left, x_right
    MatrixXd fault_pos(num_faults,3);
    fault_pos << 0.0,  -30.0e3 ,30.0e3,
                 0.3e3, 3.0e3, 30.0e3,
                 0.3e3, -30.0e3, -3.0e3,
                 -0.3e3, 3.0e3, 30.0e3,
                 -0.3e3, -30.0e3, -3.0e3;

//    MatrixXd fault_pos;
//    read_matrix("fault_pos.txt", fault_pos);
//    int num_faults = fault_pos.rows();
    cout<<"test_mat="<<"\n"<<fault_pos<<"\n"<<"rows="<<fault_pos.rows()<<"\n";
    std::vector<std::vector<int>> fault_nodes(num_faults*2);
    mesh_Generated_multi_faults(x_min,x_max,y_min,y_max,dx,dy,nx_el,ny,Node, Element,BIE_top_surf_nodes, BIE_bot_surf_nodes,fault_pos, fault_nodes);

    int n_nodes = Node.rows();
    int n_el = Element.rows();
    int Ndofn = 2;
    int Nnel = Element.cols();
    // Material
    double density = 2670.0;
    double v_s =3.464e3;
    double v_p = 6.0e3;
    double G= pow(v_s,2)*density;
    double Lambda = pow(v_p,2)*density-2.0*G;
    double E  = G*(3.0*Lambda+2.0*G)/(Lambda+G);
    double nu = Lambda/(2.0*(Lambda+G));
    // Time
    double alpha = 0.4;
    double dt = alpha*dx/v_p;
    dt = 0.0002;
    // Reyleigh Damping
    double beta =0.1;
    double q = beta*dt;
    double time_run = 6.0;
    int numt = time_run/dt;
    //numt = 1;
    //numt = 200;
    //numt = 3;
    //numt =4;
    VectorXd time = dt*VectorXd::LinSpaced(numt,1,numt);
    // Slip weakening friction parameters
    double Dc = 0.2;
    double mu_d= 0.2;
    double mu_s = 0.6;
    // Intialization
    // disp velocity current and next time step (new)
    VectorXd u_n = VectorXd::Zero(n_nodes*Ndofn,1);
    VectorXd v_n = VectorXd::Zero(n_nodes*Ndofn,1);
    VectorXd u_new = VectorXd::Zero(n_nodes*Ndofn,1);
    VectorXd v_new = VectorXd::Zero(n_nodes*Ndofn,1);
    VectorXd a_n = VectorXd::Zero(n_nodes*Ndofn,1);
    // slip and slip-rate
    std::vector<Eigen::Array<double, -1, 1> > delt_u_n(num_faults);
    std::vector<Eigen::Array<double, -1, 1> > delt_v_n(num_faults);
    std::vector<Eigen::Array<double, -1, 1> > T_c(num_faults);
    std::vector<Eigen::Array<double, -1, 1> > T_0(num_faults);
    std::vector<Eigen::Array<double, -1, 1> > tau_s(num_faults);

    // Vector containing number of elements on each faults (nx)
    std::vector<int> nx_faults(num_faults);
    for (int i=0;i<num_faults;i++)
    {
        nx_faults[i] =fault_nodes[2*i].size();
    }

    for (int i=0;i<num_faults;i++)
    {
        delt_u_n[i] = ArrayXd::Zero(Ndofn*(nx_faults[i]),1);
        delt_v_n[i] = ArrayXd::Zero(Ndofn*(nx_faults[i]),1);
        T_c[i] = ArrayXd::Zero(Ndofn*(nx_faults[i]),1);
        T_0[i] = ArrayXd::Zero(Ndofn*(nx_faults[i]),1);
       // tau_s[i] = ArrayXd::Zero(Ndofn*(nx_faults[i]),1);
        tau_s[i] = 0.6*ArrayXd::Ones((nx_faults[i]),1)*50.0e6;
    }
    VectorXd F_ext_global = VectorXd::Zero(n_nodes*Ndofn,1);
    // Setting intial stress on the fault
    for (int j=0; j<num_faults;j++)
    {
        for (int i=0; i<nx_faults[j]; i++)
        {
            T_0[j](2*i+1) = -50.0e6;
            T_c[j](2*i+1) = T_0[j](2*i+1);
        }
        if (j==0)
        {
            VectorXd x = VectorXd::LinSpaced(nx_faults[j], fault_pos(j,1), fault_pos(j,2));
            for (int i=0 ; i<nx_faults[j]; i++)
            {
                if ((x(i)<=(fault_pos(j,1)+fault_pos(j,2))/2+0.2e3)&&(x(i)>=(fault_pos(j,1)+fault_pos(j,2))/2-0.2e3))
                {
                    T_0[j](2*i) = 31.0e6;
                    T_c[j](2*i) = T_0[j](2*i);
                }
                else
                {
                    T_0[j](2*i) = 20.0e6;
                    T_c[j](2*i) = T_0[j](2*i);
                }
            }
        }
        else
        {
            for (int i=0 ; i<nx_faults[j]; i++)
            {
                T_0[j](2*i) = 20.0e6;
            }
        }
    }
    // Get the index degree of freedome for each element
    VectorXi index_el = VectorXi::Zero(Ndofn*Nnel,1);
    MatrixXi index_store = MatrixXi::Zero(Nnel*Ndofn,n_el);
    for (int i=0;i<n_el;i++)
    {
        bcdof(Element.row(i),dim,index_el);
        index_store.col(i) = index_el;
    }
    VectorXi BIE_top_surf_index = VectorXi::Zero(Ndofn*(BIE_top_surf_nodes.size()),1);
    VectorXi BIE_bot_surf_index = VectorXi::Zero(Ndofn*(BIE_bot_surf_nodes.size()),1);

    std::vector<Eigen::ArrayXi> Fault_surf_index(num_faults*2);
    for (int i=0; i <num_faults*2; i++)
    {
        Fault_surf_index[i] = ArrayXi::Zero(Ndofn*(fault_nodes[i].size()),1);
    }

    // Getting the faults DOF index
    for (int i=0; i<num_faults*2;i++)
    {
        bcdof_ptr(fault_nodes[i], dim, Fault_surf_index[i].data());
    }
    bcdof(BIE_top_surf_nodes,dim,BIE_top_surf_index);
    bcdof(BIE_bot_surf_nodes,dim,BIE_bot_surf_index);
    // Calculating the Global Mass Vector (lumped mass)
    // Element mass
    double M=density*dx*dy*1.0;
    VectorXd M_el_vec = M/4*VectorXd::Ones(Nnel*Ndofn,1);
    VectorXd M_global_vec=VectorXd::Zero(n_nodes*Ndofn,1);
    for (int i=0 ; i<n_el;i++)
    {
        index_el = index_store.col(i);
        mapglobal(index_el,M_global_vec,M_el_vec);
    }
    // Element matrix
    MatrixXd ke = MatrixXd::Zero(8,8);
    MatrixXd coord = MatrixXd::Zero(4,2);
    VectorXi Element_0= Element.row(0);
    coord.row(0) = Node.row(Element_0(0));
    coord.row(1) = Node.row(Element_0(1));
    coord.row(2) = Node.row(Element_0(2));
    coord.row(3) = Node.row(Element_0(3));
    cal_ke (coord,E,nu,ke);
    // BIE part initiation
    // Setting up the material property for the BIE code
    Material BIE_top_mat = Material(E,nu,density);
    Material BIE_bot_mat = Material(E,nu,density);
    double length = x_max-x_min;
    // infinte bc BIE call infinite_boundary.cc
    PrecomputedKernel h11("kernels/nu_.25_h11.dat");
    PrecomputedKernel h12("kernels/nu_.25_k12.dat");
    PrecomputedKernel h22("kernels/nu_.25_h22.dat");
    InfiniteBoundary BIE_inf_top(length,nx_el+1,1.0,&BIE_top_mat,&h11,&h12,&h22);
    InfiniteBoundary BIE_inf_bot(length,nx_el+1,-1.0,&BIE_bot_mat,&h11,&h12,&h22);
    // BIE setting time step
    BIE_inf_top.setTimeStep(dt);
    BIE_inf_bot.setTimeStep(dt);
    // BIE initialization
    BIE_inf_top.init();
    BIE_inf_bot.init();
    printf("ready to start\n");
    // Output
    ofstream file;
    file.open("results/num_nodes_fault.bin",ios::binary);
    file.write((char*)(nx_faults.data()),nx_faults.size()*sizeof(int));
    file.close();

    for (int i=0;i<num_faults;i++)
    {
        std::string slip = "results/slip_"+std::to_string(i)+".bin";
        file.open(slip);
        file.close();
        std::string slip_rate = "results/slip_rate_"+std::to_string(i)+".bin";
        file.open(slip_rate);
        file.close();
        std::string shear = "results/shear_"+std::to_string(i)+".bin";
        file.open(shear);
        file.close();
    }
    double start = omp_get_wtime();

    // Main time loop
    for (int j=0;j<numt;j++)
    {
        // Compute the global internal force
        VectorXd fe_global= VectorXd::Zero(n_nodes*Ndofn,1);
        cal_fe_global_const_ke(n_nodes, n_el, index_store, q, u_n, v_n, Ndofn, ke, fe_global);
        // Friction subroutine
        VectorXd F_total = F_ext_global-fe_global;
        for (int i=0 ; i<num_faults; i++)
        {
            VectorXd F_fault = VectorXd::Zero(Ndofn*(nx_faults[i]),1);
            Slip_Weakening(M_global_vec, Fault_surf_index[2*i], Fault_surf_index[2*i+1], fe_global, dt, dx, dy, nx_faults[i]-1, delt_v_n[i], delt_u_n[i], T_0[i], tau_s[i], mu_s, mu_d, Dc, Ndofn, M, F_fault, T_c[i]);
            mapglobal(Fault_surf_index[2*i],F_total,-F_fault);
            mapglobal(Fault_surf_index[2*i+1],F_total,F_fault);
        }

        // Central Difference Time integration
        time_advance(u_n, v_n, F_total, M_global_vec, dt);
        // Get the slip and slip rate

        for (int i=0; i<num_faults; i++)
        {
            cal_slip_slip_rate(u_n, v_n,  Fault_surf_index[2*i], Fault_surf_index[2*i+1] , Ndofn, nx_faults[i]-1, delt_u_n[i], delt_v_n[i]);
        }
        // Correct the BIE surf nodes solutions from FEM with the BIE solution
        BIE_correct(BIE_top_surf_index, BIE_bot_surf_index, fe_global, Ndofn, nx_el, dx, BIE_inf_top, BIE_inf_bot, u_n, v_n);

//        VectorXd delt_u_n_out = delt_u_n[0];
//        VectorXd delt_v_n_out = delt_v_n[0];
//        VectorXd T_c_out = T_c[0];
//        file.open("results/slip_main.bin",ios::binary | ios::app);
//        file.write((char*)(delt_u_n[0].data()),delt_u_n[0].size()*sizeof(double));
//        file.close();
//        file.open("results/slip_rate_main.bin",ios::binary | ios::app);
//        file.write((char*)(delt_v_n[0].data()),delt_v_n[0].size()*sizeof(double));
//        file.close();
//        file.open("results/shear_main.bin",ios::binary | ios::app);
//        file.write((char*)(T_c[0].data()),T_c[0].size()*sizeof(double));
//        file.close();
        //
        if (j%4==3)
        {
        for (int i=0;i<num_faults;i++)
        {
            std::string slip = "results/slip_"+std::to_string(i)+".bin";
            file.open(slip,ios::binary | ios::app);
            file.write((char*)(delt_u_n[i].data()),delt_u_n[i].size()*sizeof(double));
            file.close();

            std::string slip_rate = "results/slip_rate_"+std::to_string(i)+".bin";
            file.open(slip_rate,ios::binary | ios::app);
            file.write((char*)(delt_v_n[i].data()),delt_v_n[i].size()*sizeof(double));
            file.close();

            std::string shear = "results/shear_"+std::to_string(i)+".bin";
            file.open(shear,ios::binary | ios::app);
            file.write((char*)(T_c[i].data()),T_c[i].size()*sizeof(double));
            file.close();
        }
        }
        printf("Simulation time = %f\n",time(j));
      //  double end_t = omp_get_wtime();
       // std::cout<<"time_cpu_t="<<end_t-start<<std::endl;
    //   std::cout<<"time_fem="<< time_fem<<std::endl;
    //   std::cout<<"time_bie="<< time_bie<<std::endl;
    }
    double end = omp_get_wtime();
    std::cout<<"time_cpu="<<end-start<<std::endl;
    return 0;
}

void read_matrix(std::string fileName, Eigen::MatrixXd &outputMat) {
    fstream cin;
    cin.open(fileName.c_str());
    if (cin.fail())
    {
        std::cerr << "Failed to open file: " << fileName << std::endl;
        std::cin.get(); }
    string s;
    vector <vector <double> > matrix;
    while (getline(cin, s)) {
        stringstream input(s);
        double temp;
        vector <double> currentLine;
        while (input >> temp)
            currentLine.push_back(temp);
        matrix.push_back(currentLine);
    }
    if (!igl::list_to_matrix(matrix, outputMat))
    { std::cerr << "list tom matrix error" << std::endl; std::cin.get();
        //return false;
    }
}


//  Copyright John Maddock 2006, 2007.
//  Copyright Paul A. Bristow 2006, 2007.

//  Use, modification and distribution are subject to the
//  Boost Software License, Version 1.0. (See accompanying file
//  LICENSE_1_0.txt or copy at http://www.boost.org/LICENSE_1_0.txt)

#ifndef BOOST_STATS_NORMAL_HPP
#define BOOST_STATS_NORMAL_HPP

// http://en.wikipedia.org/wiki/Normal_distribution
// http://www.itl.nist.gov/div898/handbook/eda/section3/eda3661.htm
// Also:
// Weisstein, Eric W. "Normal Distribution."
// From MathWorld--A Wolfram Web Resource.
// http://mathworld.wolfram.com/NormalDistribution.html

#include <boost/math/distributions/fwd.hpp>
#include <boost/math/special_functions/erf.hpp> // for erf/erfc.
#include <boost/math/distributions/complement.hpp>
#include <boost/math/distributions/detail/common_error_handling.hpp>

#include <utility>

namespace boost{ namespace math{

template <class RealType = double, class Policy = policies::policy<> >
class normal_distribution
{
public:
   typedef RealType value_type;
   typedef Policy policy_type;

   normal_distribution(RealType l_mean = 0, RealType sd = 1)
      : m_mean(l_mean), m_sd(sd)
   { // Default is a 'standard' normal distribution N01.
     static const char* function = "boost::math::normal_distribution<%1%>::normal_distribution";

     RealType result;
     detail::check_scale(function, sd, &result, Policy());
     detail::check_location(function, l_mean, &result, Policy());
   }

   RealType mean()const
   { // alias for location.
      return m_mean;
   }

   RealType standard_deviation()const
   { // alias for scale.
      return m_sd;
   }

   // Synonyms, provided to allow generic use of find_location and find_scale.
   RealType location()const
   { // location.
      return m_mean;
   }
   RealType scale()const
   { // scale.
      return m_sd;
   }

private:
   //
   // Data members:
   //
   RealType m_mean;  // distribution mean or location.
   RealType m_sd;    // distribution standard deviation or scale.
}; // class normal_distribution

typedef normal_distribution<double> normal;

#ifdef BOOST_MSVC
#pragma warning(push)
#pragma warning(disable:4127)
#endif

template <class RealType, class Policy>
inline const std::pair<RealType, RealType> range(const normal_distribution<RealType, Policy>& /*dist*/)
{ // Range of permissible values for random variable x.
  if (std::numeric_limits<RealType>::has_infinity)
  {
     return std::pair<RealType, RealType>(-std::numeric_limits<RealType>::infinity(), std::numeric_limits<RealType>::infinity()); // - to + infinity.
  }
  else
  { // Can only use max_value.
    using boost::math::tools::max_value;
    return std::pair<RealType, RealType>(-max_value<RealType>(), max_value<RealType>()); // - to + max value.
  }
}

template <class RealType, class Policy>
inline const std::pair<RealType, RealType> support(const normal_distribution<RealType, Policy>& /*dist*/)
{ // This is range values for random variable x where cdf rises from 0 to 1, and outside it, the pdf is zero.
  if (std::numeric_limits<RealType>::has_infinity)
  {
     return std::pair<RealType, RealType>(-std::numeric_limits<RealType>::infinity(), std::numeric_limits<RealType>::infinity()); // - to + infinity.
  }
  else
  { // Can only use max_value.
   using boost::math::tools::max_value;
   return std::pair<RealType, RealType>(-max_value<RealType>(),  max_value<RealType>()); // - to + max value.
  }
}

#ifdef BOOST_MSVC
#pragma warning(pop)
#endif

template <class RealType, class Policy>
inline RealType pdf(const normal_distribution<RealType, Policy>& dist, const RealType& x)
{
   BOOST_MATH_STD_USING  // for ADL of std functions

   RealType sd = dist.standard_deviation();
   RealType mean = dist.mean();

   static const char* function = "boost::math::pdf(const normal_distribution<%1%>&, %1%)";

   RealType result = 0;
   if(false == detail::check_scale(function, sd, &result, Policy()))
   {
      return result;
   }
   if(false == detail::check_location(function, mean, &result, Policy()))
   {
      return result;
   }
   if((boost::math::isinf)(x))
   {
     return 0; // pdf + and - infinity is zero.
   }
   // Below produces MSVC 4127 warnings, so the above used instead.
   //if(std::numeric_limits<RealType>::has_infinity && abs(x) == std::numeric_limits<RealType>::infinity())
   //{ // pdf + and - infinity is zero.
   //  return 0;
   //}
   if(false == detail::check_x(function, x, &result, Policy()))
   {
      return result;
   }

   RealType exponent = x - mean;
   exponent *= -exponent;
   exponent /= 2 * sd * sd;

   result = exp(exponent);
   result /= sd * sqrt(2 * constants::pi<RealType>());

   return result;
} // pdf

template <class RealType, class Policy>
inline RealType cdf(const normal_distribution<RealType, Policy>& dist, const RealType& x)
{
   BOOST_MATH_STD_USING  // for ADL of std functions

   RealType sd = dist.standard_deviation();
   RealType mean = dist.mean();
   static const char* function = "boost::math::cdf(const normal_distribution<%1%>&, %1%)";
   RealType result = 0;
   if(false == detail::check_scale(function, sd, &result, Policy()))
   {
      return result;
   }
   if(false == detail::check_location(function, mean, &result, Policy()))
   {
      return result;
   }
   if((boost::math::isinf)(x))
   {
     if(x < 0) return 0; // -infinity
     return 1; // + infinity
   }
   // These produce MSVC 4127 warnings, so the above used instead.
   //if(std::numeric_limits<RealType>::has_infinity && x == std::numeric_limits<RealType>::infinity())
   //{ // cdf +infinity is unity.
   //  return 1;
   //}
   //if(std::numeric_limits<RealType>::has_infinity && x == -std::numeric_limits<RealType>::infinity())
   //{ // cdf -infinity is zero.
   //  return 0;
   //}
   if(false == detail::check_x(function, x, &result, Policy()))
   {
     return result;
   }
   RealType diff = (x - mean) / (sd * constants::root_two<RealType>());
   result = boost::math::erfc(-diff, Policy()) / 2;
   return result;
} // cdf

template <class RealType, class Policy>
inline RealType quantile(const normal_distribution<RealType, Policy>& dist, const RealType& p)
{
   BOOST_MATH_STD_USING  // for ADL of std functions

   RealType sd = dist.standard_deviation();
   RealType mean = dist.mean();
   static const char* function = "boost::math::quantile(const normal_distribution<%1%>&, %1%)";

   RealType result = 0;
   if(false == detail::check_scale(function, sd, &result, Policy()))
      return result;
   if(false == detail::check_location(function, mean, &result, Policy()))
      return result;
   if(false == detail::check_probability(function, p, &result, Policy()))
      return result;

   result= boost::math::erfc_inv(2 * p, Policy());
   result = -result;
   result *= sd * constants::root_two<RealType>();
   result += mean;
   return result;
} // quantile

template <class RealType, class Policy>
inline RealType cdf(const complemented2_type<normal_distribution<RealType, Policy>, RealType>& c)
{
   BOOST_MATH_STD_USING  // for ADL of std functions

   RealType sd = c.dist.standard_deviation();
   RealType mean = c.dist.mean();
   RealType x = c.param;
   static const char* function = "boost::math::cdf(const complement(normal_distribution<%1%>&), %1%)";

   RealType result = 0;
   if(false == detail::check_scale(function, sd, &result, Policy()))
      return result;
   if(false == detail::check_location(function, mean, &result, Policy()))
      return result;
   if((boost::math::isinf)(x))
   {
     if(x < 0) return 1; // cdf complement -infinity is unity.
     return 0; // cdf complement +infinity is zero
   }
   // These produce MSVC 4127 warnings, so the above used instead.
   //if(std::numeric_limits<RealType>::has_infinity && x == std::numeric_limits<RealType>::infinity())
   //{ // cdf complement +infinity is zero.
   //  return 0;
   //}
   //if(std::numeric_limits<RealType>::has_infinity && x == -std::numeric_limits<RealType>::infinity())
   //{ // cdf complement -infinity is unity.
   //  return 1;
   //}
   if(false == detail::check_x(function, x, &result, Policy()))
      return result;

   RealType diff = (x - mean) / (sd * constants::root_two<RealType>());
   result = boost::math::erfc(diff, Policy()) / 2;
   return result;
} // cdf complement

template <class RealType, class Policy>
inline RealType quantile(const complemented2_type<normal_distribution<RealType, Policy>, RealType>& c)
{
   BOOST_MATH_STD_USING  // for ADL of std functions

   RealType sd = c.dist.standard_deviation();
   RealType mean = c.dist.mean();
   static const char* function = "boost::math::quantile(const complement(normal_distribution<%1%>&), %1%)";
   RealType result = 0;
   if(false == detail::check_scale(function, sd, &result, Policy()))
      return result;
   if(false == detail::check_location(function, mean, &result, Policy()))
      return result;
   RealType q = c.param;
   if(false == detail::check_probability(function, q, &result, Policy()))
      return result;
   result = boost::math::erfc_inv(2 * q, Policy());
   result *= sd * constants::root_two<RealType>();
   result += mean;
   return result;
} // quantile

template <class RealType, class Policy>
inline RealType mean(const normal_distribution<RealType, Policy>& dist)
{
   return dist.mean();
}

template <class RealType, class Policy>
inline RealType standard_deviation(const normal_distribution<RealType, Policy>& dist)
{
   return dist.standard_deviation();
}

template <class RealType, class Policy>
inline RealType mode(const normal_distribution<RealType, Policy>& dist)
{
   return dist.mean();
}

template <class RealType, class Policy>
inline RealType median(const normal_distribution<RealType, Policy>& dist)
{
   return dist.mean();
}

template <class RealType, class Policy>
inline RealType skewness(const normal_distribution<RealType, Policy>& /*dist*/)
{
   return 0;
}

template <class RealType, class Policy>
inline RealType kurtosis(const normal_distribution<RealType, Policy>& /*dist*/)
{
   return 3;
}

template <class RealType, class Policy>
inline RealType kurtosis_excess(const normal_distribution<RealType, Policy>& /*dist*/)
{
   return 0;
}

template <class RealType, class Policy>
inline RealType entropy(const normal_distribution<RealType, Policy> & dist)
{
   using std::log;
   RealType arg = constants::two_pi<RealType>()*constants::e<RealType>()*dist.standard_deviation()*dist.standard_deviation();
   return log(arg)/2;
}

} // namespace math
} // namespace boost

// This include must be at the end, *after* the accessors
// for this distribution have been defined, in order to
// keep compilers that support two-phase lookup happy.
#include <boost/math/distributions/detail/derived_accessors.hpp>

#endif // BOOST_STATS_NORMAL_HPP


#include "rememberBoostHistogram.h"
#include <iostream>
#include <fstream>
#include <string>
#include <boost/lexical_cast.hpp>
#include "sgcu/sgcu.hpp"
#include "utils.hpp"

using std::cout;
using std::cerr;
using std::endl;
using std::string;

namespace SG{

//using namespace boost::accumulators;

// build a histogram for a data vector
template <typename T>
void getHistogram(std::vector<T> const & data, histogram_type & hist)
{
	// create a boost accumulator
	acc myAccumulator( boost::accumulators::tag::density::num_bins = data.size()/10.0, boost::accumulators::tag::density::cache_size = data.size()-2);
	// fill accumulator
	for (unsigned int j = 0; j < data.size(); ++j){
		myAccumulator(data[j]);
	}
	// do the boost histogram
	hist = boost::accumulators::density(myAccumulator);
}

unsigned int rememberBoostHistogram::size() const
{
	unsigned int thisSize = values.size();
	if (indices.size() < thisSize)
		thisSize = indices.size();
	return thisSize;
}

void rememberBoostHistogram::clear()
{
	values.clear();
	indices.clear();
}

void rememberBoostHistogram::rememberThisHistogram(histogram_type& hist)
{
	clear();
	for (unsigned int i = 0; i < (unsigned int) hist.size(); ++i)
	{
		indices.push_back(hist[i].first);
		values.push_back(hist[i].second);
	}
}

void rememberBoostHistogram::operator()(histogram_type& hist)
{
	rememberThisHistogram(hist);
}

rememberBoostHistogram::rememberBoostHistogram()
{
}

rememberBoostHistogram::rememberBoostHistogram(histogram_type& hist)
{
	rememberThisHistogram(hist);
}

double computeArea(struct rememberBoostHistogram const & data)
{
	double area = 0.0;
	for (unsigned int i = 0; i < data.size(); ++i)
	{
		area += data.values[i];
	}
	return area;
}

void getFitHistogramOwn(struct rememberBoostHistogram const & dataOrg, struct rememberBoostHistogram& fit)
{
	// NOTE: this implicitly assumes all histogram values to be positive. true for a histograms, not necessarily true for all inputs!
	if ( computeArea(dataOrg) == 0.0 )
	{
		fit = dataOrg;
		return;
	}

	const bool doDebugOutput = false;

	struct rememberBoostHistogram data = dataOrg;
	if (doDebugOutput) saveHistogramToFile(data, "comparisonValsNormalization3.txt");

	// make sure the histogram is normalized to 1.0
	normalizeHistogramAreaTo(data, 1.0);

	if (doDebugOutput) saveHistogramToFile(data, "comparisonValsNormalization3Renorm.txt");

	unsigned int lastIndex = computeLastIndex(data);

	//	cerr << "new: setting last entry of histogram to zero" << endl;
	data.values[data.size()-1] = 0.0;

	struct rememberBoostHistogram dataBkp = data;
	struct rememberBoostHistogram dataRenormalizedBkp = data;

	smoothHistogram(data);

	if (doDebugOutput) saveHistogramToFile(data, "comparisonValsNormalization3Smoothed.txt");

	unsigned int firstIndex = getFirstIndex(data, lastIndex);

	// do not consider values from before the tail to fit the tail!
	eraseEverythingButTail(data, firstIndex);
	eraseEverythingButTail(dataBkp, firstIndex);

	double areaBefore = computeArea(dataBkp);
	normalizeHistogramAreaTo(data, areaBefore);

	if (doDebugOutput) saveHistogramToFile(data, "comparisonValsNormalization3SmoothedRenorm.txt");

	double scalingFactor = 0.0;
	double exponent = 0.0;
	obtainFitParameters(data, firstIndex, lastIndex, dataRenormalizedBkp, scalingFactor, exponent);

	generateActualFit(data, scalingFactor, exponent, fit);

	mergeOriginalAndFit(dataRenormalizedBkp, firstIndex, lastIndex, fit);

	if (doDebugOutput) saveHistogramToFile(fit, "comparisonValsNormalization3Fitted.txt");
}

void getFitHistogramOwn(const std::vector<double>& data, struct rememberBoostHistogram& fit)
{
	// do a histogram
	acc myAccumulator( boost::accumulators::tag::density::num_bins = 100, boost::accumulators::tag::density::cache_size = data.size());
	for (unsigned int i = 0; i < data.size(); ++i)
	{
		myAccumulator(data[i]);
	}
	histogram_type hist = boost::accumulators::density(myAccumulator);
	struct rememberBoostHistogram histNew(hist);

	getFitHistogramOwn(histNew, fit);
}

void outputHistogram(histogram_type const & hist, std::string fileName)
{
	std::ofstream raus;
	raus.open(fileName.c_str());
	if ( (! raus.is_open()) || (! raus.good()) )
	{
		std::cerr << "Failed to open file " << fileName << " for histogram output." << endl;
		throw -1;
	}

	for (unsigned int i = 0; i < (unsigned int) hist.size(); ++i)
	{
		raus << hist[i].first << " " << hist[i].second << endl;
	}

	if (! raus.good() )
	{
		std::cerr << "Failed to write histogram to file " << fileName << endl;
		throw -1;
	}
	raus.close();

}

void outputRememberHistogram(rememberBoostHistogram const & fit, std::string fileName)
{
	std::ofstream raus;
	raus.open(fileName.c_str());
	if ( (! raus.is_open()) || (! raus.good()) )
	{
		std::cerr << "Failed to open file " << fileName << " for rememberBoostHistogram output." << endl;
		throw -1;
	}

	for (unsigned int i = 0; i < fit.size(); ++i)
	{
		raus << fit.indices[i] << " " << fit.values[i] << endl;
	}
	raus.close();

	if (! raus.good() )
	{
		std::cerr << "Failed to write rememberBoostHistogram to file " << fileName << endl;
		throw -1;
	}
	raus.close();
}

// TODO: check if both outputs are actually used!
// for a fixed segmentation threshold: compute histograms for each metric from the clusters found in all pseudo experiments
void computeAndOutputHistogramsFromClusters(std::string folderName, unsigned int n, std::vector<std::vector<double> > & relevanceComparisonVals,
		std::string searchStringSegmentationThreshold,
		std::vector<std::vector<struct rememberBoostHistogram > >& histsAllSegmentationsAllRelevances)
{
		histsAllSegmentationsAllRelevances[n].resize(relevanceComparisonVals.size());
		for (unsigned int j = 0; j < relevanceComparisonVals.size(); ++j)	// for all metrics
		{
			// create a histogram from all clusters that match this segmentation and metric
			acc myAcc(boost::accumulators::tag::density::num_bins = 100, boost::accumulators::tag::density::cache_size = relevanceComparisonVals[j].size());
			for (unsigned int i = 0; i < relevanceComparisonVals[j].size(); ++i)
			{
				myAcc(relevanceComparisonVals[j][i]);
			}
			histogram_type hist = boost::accumulators::density(myAcc);
			outputHistogram(hist, folderName + "/hist_Size" + searchStringSegmentationThreshold + "_Metric" + boost::lexical_cast<std::string>(j) + ".txt");

			struct rememberBoostHistogram data(hist);
			struct rememberBoostHistogram fit;
			// compute a fit that approximates the tail of the observed data
			getFitHistogramOwn(data, fit);
			outputRememberHistogram(fit, folderName + "/hist_Size" + searchStringSegmentationThreshold + "_Metric" + boost::lexical_cast<std::string>(j) + "fit.txt");

			histsAllSegmentationsAllRelevances[n][j] = fit;
		}
}

void saveHistogramToFile(struct rememberBoostHistogram const & data, std::string const & fileName)
{
	std::ofstream raus(fileName.c_str());
	if (! raus.good())
	{
		std::cerr << "Failed to open file " << fileName << " for output of histogram data" << std::endl;
		throw (-1);
	}
	for (unsigned int i = 0; i < data.size(); ++i)
	{
		raus << data.indices[i] << " " << data.values[i] << endl;
		if (! raus.good())
		{
			std::cerr << "Failed to write histogram data entry " << i << " of " << data.size() << " to file " << fileName << std::endl;
			throw (-2);
		}
	}
	if (! raus.good())
	{
		std::cerr << "Failed to write last of " << data.size() << " histogram data entries to file " << fileName << std::endl;
		throw (-2);
	}
	raus.close();
}

void normalizeHistogramAreaTo(struct rememberBoostHistogram & data, double targetValue)
{
	double area = computeArea(data);
	double rescaleFactor = 0.0;
	if (area != 0.0)
	{
		rescaleFactor = targetValue/area;
		for (unsigned int i = 0; i < data.size(); ++i)
		{
			data.values[i] *= rescaleFactor;
		}
	}
}

unsigned int computeLastIndex(struct rememberBoostHistogram const & data)
{
	unsigned int lastIndex = 0;
	unsigned int notZeroIndex = 0;
	unsigned int notZeroIndexBefore = 0;
	for (unsigned int i = 0; i < data.size(); ++i)
	{
		if (data.values[i] > 0.0)
		{
			lastIndex = notZeroIndexBefore;
			notZeroIndexBefore = notZeroIndex;
			notZeroIndex = i;
		}
	}
	if (lastIndex == 0)
	{
		lastIndex = notZeroIndexBefore;
	}
	if (lastIndex == 0)
	{
		lastIndex = notZeroIndex;
	}
	if (lastIndex == 0)
	{
		lastIndex = data.size();
	}
	return lastIndex;
}

void smoothHistogram(struct rememberBoostHistogram & data)
{
	// NOTE: I just noted that the number of smoothings correlates with the size of the vector.
	// This has been intentional, but it hasn't been checked if this scales to less or more entries.
	for (unsigned int i = 0; i < 0.75*data.size(); ++i)
		smoothNoEnds(data.values);
	for (unsigned int i = 0; i < 0.25*data.size(); ++i)
		sgcu::smooth(data.values);
}

double evaluateFitL2(const std::vector<double>& data, const std::vector<double>& fit){
	if (data.size() != fit.size()){
		cerr << "evaluateFitL2 cannot compare data and fit of different size, data: " << data.size() << " and fit: " << fit.size() << endl;
		return std::numeric_limits<double>::max();
	}

	double diffSum = 0.0;
	for (unsigned int i = 0; i < data.size(); ++i){
		double diff = data[i] - fit[i];
		diffSum += (diff*diff);
	}
	return diffSum;
}

void obtainFitParameters(struct rememberBoostHistogram & data, unsigned int firstIndex, unsigned int lastIndex,
		struct rememberBoostHistogram const & dataRenormalizedBkp,
		double & scalingFactorFitFinal, double & exponent)
{
	// evaluate when each value of the smoothed right flanc is below half of its value - store the distance in x
	//cout << "maxIndex = " << maxIndex << endl;
	//cout << "firstIndex = " << firstIndex << endl;
	//cout << "lastIndex = " << lastIndex << endl;
	unsigned int num = 0;
	double howLongUntilBelowHalf = 0.0;
	for (unsigned int i = firstIndex; i < lastIndex; ++i){
		double currentY100 = data.values[i];
		if (currentY100 <= 0.0)
			break;
		double currentX = data.indices[i];
		for (unsigned int j = i+1; j < data.size(); ++j){
	//		for (unsigned int j = i+1; j < lastIndex; ++j){	// this was the default
			double currentY = data.values[j];
			if (currentY <= 0.5*currentY100){
				if (currentY > 0.0){
					double distance = data.indices[j] - currentX;
					//cout << i << " " << j << " distance: " << distance << "    histjfirst = " << data.indices[j] << " currentX = " << currentX << " currentY100 = " << currentY100 << " currentY = " << currentY << endl;
					howLongUntilBelowHalf += distance;
					++num;
				}
				break;
			}
		}
	}
	//cout << "nowLongUntil = " << howLongUntilBelowHalf << endl;
	//cout << "num = " << num << endl;
	if (num != 0){
		howLongUntilBelowHalf /= static_cast<double>(num);
	}

	double stepSize = data.indices[1] - data.indices[0];

	//cout << "Mean x distance until y value is below half of its starting value: " << howLongUntilBelowHalf << endl;
	howLongUntilBelowHalf  -= 0.5*stepSize;
	//cout << "x distance corrected for stepsize: " << howLongUntilBelowHalf << endl;
	//double
	exponent = -1.0/howLongUntilBelowHalf;
	//cout << "exponent of fit: " << exponent << endl;

	double minDistance = std::numeric_limits<double>::max();
	//double
	scalingFactorFitFinal = 0.0;

	std::vector<double> histValsFitComparison = dataRenormalizedBkp.values;
	for (unsigned int j = 0; j < firstIndex; ++j){
		histValsFitComparison[j] = 0.0;
	}
	for (unsigned int j = lastIndex; j < histValsFitComparison.size(); ++j){
		histValsFitComparison[j] = 0.0;
	}

	std::vector<double> fitVals = histValsFitComparison;
	for (unsigned int i = firstIndex; i < lastIndex; ++i){
		double startingXValue = data.indices[i];
		double expectedStartingValue = data.values[i];
		double plainStartingValue = pow(2.0, exponent*startingXValue);
		double scalingFactorFit = expectedStartingValue/plainStartingValue;

		for (unsigned int j = firstIndex; j < lastIndex; ++j){
			double currentX = data.indices[j];
			fitVals[j] = scalingFactorFit*pow(2.0, exponent*currentX);
		}

		double distance = evaluateFitL2(histValsFitComparison, fitVals);
		//cout << "index " << i << " distance " << distance << endl;
		if (distance < minDistance){
			minDistance = distance;
			scalingFactorFitFinal = scalingFactorFit;
		}
	}

	//cout << "scalingfactor: " << scalingFactorFitFinal << endl;
}

void generateActualFit(struct rememberBoostHistogram const & data, double scalingFactor, double exponent,
		struct rememberBoostHistogram& fit, double longerFactor, double finerFactor)
{
	double stepSize = data.indices[1] - data.indices[0];
	unsigned int targetSize = data.size()*longerFactor*finerFactor;
	fit.clear();
	fit.indices.resize(targetSize);
	fit.values.resize(targetSize);
	double startIndex = data.indices[0];
	for (unsigned int i = 0; i < targetSize; ++i){
		fit.indices[i] = startIndex+stepSize*(static_cast<double>(i)/finerFactor);
	}
	for (unsigned int i = 0; i < targetSize; ++i){
		fit.values[i] = scalingFactor*pow(2.0, exponent*fit.indices[i]);
	}
}

void mergeOriginalAndFit(struct rememberBoostHistogram const & dataRenormalizedBkp, unsigned int firstIndex,
		unsigned int lastIndex, struct rememberBoostHistogram& fit, double finerFactor)
{
	unsigned int firstIndexFiner = static_cast<unsigned int>(static_cast<double>(firstIndex)*finerFactor+0.5);
	unsigned int lastIndexFiner = static_cast<unsigned int>(static_cast<double>(lastIndex)*finerFactor+0.5);

	// only start using the fit when you don't trust the actual distribution anymore:
//	double shiftBack = 0.8;	// default
	double shiftBack = 0.6;

	unsigned int startReallyUsingIt = ((1.0-shiftBack)*firstIndex + shiftBack*lastIndex);
	unsigned int startReallyUsingItFiner = ((1.0-shiftBack)*firstIndexFiner + shiftBack*lastIndexFiner)-0.55*finerFactor;

	double area2 = 0.0;
	for (unsigned int i = startReallyUsingIt; i < dataRenormalizedBkp.size(); ++i){
		area2 += dataRenormalizedBkp.values[i];
	}

	double areaFit = 0.0;
	for (unsigned int i = startReallyUsingItFiner; i < fit.values.size(); ++i){
		areaFit += fit.values[i];
	}
	double rescalingFactor2 = 1.0/finerFactor;
	if (areaFit != 0.0)
		rescalingFactor2 = area2/areaFit;

	for (unsigned int i = 0; i < startReallyUsingItFiner; ++i){
		fit.values[i] = 0.0;
	}
	for (unsigned int i = startReallyUsingItFiner; i < fit.size(); ++i){
		fit.values[i] *= rescalingFactor2;
	}
	for (unsigned int i = 0; i < startReallyUsingIt; ++i){
		fit.values[static_cast<unsigned int>(static_cast<double>(i)*finerFactor+0.5)] = dataRenormalizedBkp.values[i];
	}

	// make sure the histogram is normalized to 1.0
	normalizeHistogramAreaTo(fit, 1.0);
}

unsigned int getFirstIndex(struct rememberBoostHistogram const & data, unsigned int lastIndex)
{
	unsigned int maxIndex = (std::max_element(data.values.begin(), data.values.end())-data.values.begin());
	return static_cast<unsigned int>((2.0*maxIndex+1.0*lastIndex)/3.0);
}

void eraseEverythingButTail(struct rememberBoostHistogram & data, unsigned int firstIndex)
{
	for (unsigned int i = 0; i < firstIndex; ++i)
	{
		data.values[i] = 0.0;
	}
}

} // end of namespace SG


//  Copyright John Maddock 2007.
//  Use, modification and distribution are subject to the
//  Boost Software License, Version 1.0. (See accompanying file
//  LICENSE_1_0.txt or copy at http://www.boost.org/LICENSE_1_0.txt)

// Note that this file contains quickbook mark-up as well as code
// and comments, don't change any of the special comment mark-ups!

#include <iostream>

//[policy_eg_7

#include <boost/math/distributions.hpp>

namespace {

using namespace boost::math::policies;

typedef policy<
   // return infinity and set errno rather than throw:
   overflow_error<errno_on_error>,
   // Don't promote double -> long double internally:
   promote_double<false>,
   // Return the closest integer result for discrete quantiles:
   discrete_quantile<integer_round_nearest>
> my_policy;

BOOST_MATH_DECLARE_DISTRIBUTIONS(double, my_policy)

} // close namespace my_namespace

int main()
{
   //
   // Start with something we know will overflow:
   //
   normal norm(10, 2);
   errno = 0;
   std::cout << "Result of quantile(norm, 0) is: "
      << quantile(norm, 0) << std::endl;
   std::cout << "errno = " << errno << std::endl;
   errno = 0;
   std::cout << "Result of quantile(norm, 1) is: "
      << quantile(norm, 1) << std::endl;
   std::cout << "errno = " << errno << std::endl;
   //
   // Now try a discrete distribution:
   //
   binomial binom(20, 0.25);
   std::cout << "Result of quantile(binom, 0.05) is: "
      << quantile(binom, 0.05) << std::endl;
   std::cout << "Result of quantile(complement(binom, 0.05)) is: "
      << quantile(complement(binom, 0.05)) << std::endl;
}

//] ends quickbook imported section


#include <armadillo>
#include <iostream>
#include <cmath>
#include "LowPassFilter.hpp"
#include "LowPassFilter.cpp"
#include <ros/ros.h>
#include "std_msgs/String.h"
#include <sigproc_lib/sigproc_lib.h>
#include "sigproc_lib/sigproc.h"

using namespace arma;
using namespace as64_;

class GDMP
{
public:

  int BFs, NSamples, extra_train, Ns, id, window;
  double tNyq, dt, Tend;
  double y0, dy0, g, y0d_hat, gd_hat, Ry0d_hat, Rgd_hat;
  double az, bz, K, D, k, ks, kt;

  mat c, h, psi, w;
  mat y, dy, ddy, y_smoothed, Ry_smoothed;
  mat Ry, Rdy, Rddy;

  GDMP();
  ~GDMP();
  void training(mat y_desired,mat dy_desired,mat ddy_desired, vec timed);
  mat find_reference_desired(mat y_desired,mat dy_desired,mat ddy_desired, vec timed);
  void calculate_centers(vec timed);
  void generate_weights(mat fp_des);
  void run_solution(double goal, CanonicalStructure cs, double ts, double main_time, double extra_time, int extra_samples);
  void run_reverse_solution(double goal, CanonicalStructure cs, double ts, double main_time, double extra_time, int extra_samples);
  double fNyquistFunc(double *original, int size, double fs, double T);
  vec run_solution_dt(double t_now, double goal, double ts, int i, double f_prev, double df_prev, double ddf_prev );
  vec init_solution_dt(double goal, CanonicalStructure cs, double main_time, int extra_samples, double t_now, int i );
  vec init_Rsolution_dt(double goal, CanonicalStructure cs, double main_time, int extra_samples, double t_now, int i);
  vec run_Rsolution_dt(double t_now, double goal, double ts, int i, double f_prev, double df_prev, double ddf_prev );

};

GDMP::GDMP()
{
  std::cout << "/* Constructor of GDMP */" << '\n';
}

GDMP::~GDMP()
{
  std::cout << "/* Deconstructor of GDMP */" << '\n';
}

void GDMP::training(mat y_desired,mat dy_desired,mat ddy_desired, vec timed)
{
  std::cout << "/* Start of Training */" << '\n';

  NSamples = y_desired.n_cols;
  dt = timed(3)-timed(2);
  Tend = timed(NSamples-1);
  extra_train = ceil(NSamples/2);

  /* Reference Desired*/
  mat mat_refD(size(y_desired));
  mat_refD = find_reference_desired(y_desired, dy_desired, ddy_desired, timed);

  /* Calculate centers h psi */
  calculate_centers(timed);

  /* Generate weights */
  generate_weights(mat_refD.row(0));

  std::cout << "/* End of Training */" << '\n';
}

mat GDMP::find_reference_desired(mat y_desired,mat dy_desired,mat ddy_desired, vec timed)
{
  cout<<"Reference Desired Checked"<<endl;

  ks = 1;
  kt = 1;
  y0 = y_desired(0);
  dy0 = dy_desired(0);
  g =  y_desired(NSamples-1);

  /* sig y_desired to 0*/
  mat tsig = linspace(-Tend-extra_train*dt,extra_train*dt,NSamples+extra_train);
  mat sig = 1-1/(1+exp(-0.08/dt*tsig));

  double y_scaled_ext[NSamples+extra_train], dy_ext[NSamples+extra_train], ddy_ext[NSamples+extra_train], timed_ext[NSamples+extra_train];

  for (int i = 0; i < NSamples; i++)
  {
    y_scaled_ext[i] = sig(i)*(y_desired[i] - y0);
    dy_ext[i] = sig(i)*dy_desired[i];
    ddy_ext[i] = sig(i)*ddy_desired[i];
    timed_ext[i] = timed(i);
  }

  for (int i = NSamples;i<NSamples+extra_train;i++)
  {
    y_scaled_ext[i] = sig(i)*(y_desired[NSamples-1] - y0);
    dy_ext[i] = sig(i)*dy_desired(NSamples-1);
    ddy_ext[i] = sig(i)*ddy_desired(NSamples-1);
    timed_ext[i] = i*dt ;
  }

  /* original signal */
  double fd_original_scaled[NSamples+extra_train];
  double fd_filtered[NSamples];

  vec fd_original_ext(3*NSamples+extra_train);
  vec fd_filtered_ext(3*NSamples+extra_train);

  /* find original scaled */
  for (int i = 0; i < NSamples+extra_train; i++)
  {
    fd_original_scaled[i] = ks*(y_scaled_ext[i]-y0) + y0; //basically is ext too
  }

  /* extent before filter */
  for (int i = 0; i < NSamples; i++)
  {
    fd_original_ext(i) = fd_original_scaled[0];
  }
  for (int i = 0; i < NSamples+extra_train; i++)
  {
    fd_original_ext( i + NSamples ) = fd_original_scaled[i] ;
  }
  for (int i = 2*NSamples+extra_train; i < 3*NSamples+extra_train; i++)
  {
    fd_original_ext(i) = fd_original_scaled[NSamples+extra_train-1];
  }

  double  tau = Tend, a = 10, b = a/4;

  /* Nyquist - fc */
  double fNyquist;
  fNyquist = fNyquistFunc(fd_original_scaled, NSamples+extra_train, 1/dt, Tend + extra_train*dt );

  BFs = ceil(Tend*fNyquist + extra_train*dt*fNyquist);

  std::cout << fNyquist << '\n';
  cout<<BFs<<endl;
  std::cout << Tend*fNyquist << '\n';
  std::cout << extra_train*dt*fNyquist << '\n';

  tNyq = 1/fNyquist;
//  Ns = ceil(NSamples/BFs);

  /* Lowpass Filter */
  // LowPassFilter lpf((1/(2*tNyq)), dt);
  // for(int i = 0; i < 3*NSamples+extra_train; i++)
  // {
  // 		fd_filtered_ext[i] = lpf.update(fd_original_ext[i]) ; //Update with 1.0 as input value.
  // }

  std::cout << "Filtering...";

  arma::vec filter_coeff = spl_::fir1(100 /*filter order*/, fNyquist*dt /*normalized cutoff Frequency*/);
  fd_filtered_ext = arma::conv(fd_original_ext, filter_coeff.t(), "same");
  std::cout << "[DONE]!\n";

  /* keep specific fd_filtered size of NSamples */
  for (int i = 0; i < NSamples+extra_train; i++)
  {
    fd_filtered[i] = fd_filtered_ext(i + NSamples) + y0;
  }

  /* A to return in once Reference y,dy,ddy */
  mat A(3,NSamples+extra_train); //3 y,dy,ddy not because of DIM
  for (int i = 0; i < NSamples+extra_train; i++)
  {
    A(0,i) = fd_filtered[i];
    A(1,i) = ks*kt*dy_ext[i];
    A(2,i) = ks*pow(kt,2)*ddy_ext[i];
  }

  /* Printed in files for CHECK*/
  ofstream myFile;

  std::ostringstream oss,oss2;

  oss << "CHECK/fd_ori" << id <<".log";
  myFile.open((oss.str()).c_str());
  for (int i = 0; i < NSamples+extra_train; i++)
    myFile<<fd_original_scaled[i]<<endl;
  myFile.close();

  oss2 << "CHECK/fd_filtered" << id <<".log";
  myFile.open((oss2.str()).c_str());
  for (int i = 0; i < NSamples+extra_train; i++)
    myFile<<fd_filtered[i]<<endl;
  myFile.close();

  // oss << "CHECK/drefD" << id <<".log";
  // myFile.open((oss.str()).c_str());
  // myFile<<A.row(1)<<endl;
  // myFile.close();
  //
  // oss << "CHECK/ddrefD" << id <<".log";
  // myFile.open((oss.str()).c_str());
  // myFile<<A.row(2)<<endl;
  // myFile.close();

  return A;
}

/* Calculate centers of kernels */
void GDMP::calculate_centers(vec timed)
{
  c.set_size(1,BFs);

  for (int idx = 0; idx<BFs; idx++)
  {
    c(idx) = (idx)*tNyq;
  }

  /* ext time for extra train*/
  vec timed_ext_train(NSamples+extra_train);
  for (int i = 0; i <NSamples; i++ )
    timed_ext_train(i) = timed(i);
  for (int i = 0; i < extra_train; i++ )
    timed_ext_train(i+NSamples) =  (i+1)*dt + timed(NSamples-1) ;

  /* Sincs == psi */
  psi.set_size(BFs , NSamples + extra_train);
  double x;
  for (int b = 0; b < BFs ; b++)
  {
    for (int i = 0; i <NSamples+ extra_train; i++ )
    {
      x = ((timed_ext_train(i) - c(b))/tNyq);
      if (x != 0)
        psi(b,i) = sin(M_PI*x)/(M_PI*x);
      else
        psi(b,i) = 1;
    }
  }

}

/* Weights */
void GDMP::generate_weights(mat fp_des)
{
  std::cout << "/* start generate weights */" << '\n';

  mat scaled_fp_d = fp_des.t() - y0*ones(NSamples+extra_train);

  w.set_size(BFs ,1);
  /* weights for BFs*/
  // !! DERIVATE WITH 0 ?
  if ( g == y0 )
  {
    for (int idx = 0; idx<BFs; idx++)
    {
        w(idx) = 1;
    }
  }
  else
  {
    for (int idx = 0; idx<BFs; idx++)
    {
        w(idx) = scaled_fp_d(round(c(idx)/dt))/(g - y0);
    }
  }

}

vec GDMP::init_solution_dt(double goal, CanonicalStructure cs, double main_time, int extra_samples, double t_now, int i)
{
  /* Set size = NSamples + extra samples
      extra samples oq from extra time to run
  */
  y.set_size(1,NSamples+extra_samples);
  dy.set_size(1,NSamples+extra_samples);
  ddy.set_size(1,NSamples+extra_samples);

  az = 10;
  bz = az/4;

  K = az*bz/pow(cs.taf,2);
  D = (az + cs.dtaf)/cs.taf;

  k = goal - y0;

  /* Init fp*/
  double s = 0.0 ;
  for (int l = 0; l < BFs ; l++)
  {
    s = s + w(l)*psi(l,i);
  }

  /* Declare double mat vec fp d dd */
  double fp =  k*s + y0;
  double dfp = 0;
  double ddfp = 0;

  y0d_hat = fp;

  /* needed to get gd_hat */
  s = 0.0 ;
  for (int l = 0; l < BFs ; l++)
  {
    s = s + w(l)*psi(l,NSamples-1);
  }

  gd_hat = k*s + y0;

  // // !! DERIVATE WITH 0 ?
  // if (gd_hat == y0d_hat)
  // {
  // }
  // else
  // {
  //   ks = (goal - y0) / (gd_hat - y0d_hat);
  // }
  ks = 1;

  /* Ref compute*/
  double y_ref  = ks*(fp - y0d_hat) + y0;
  double dy_ref  = ks*dfp ;
  double ddy_ref  = ks*ddfp ;

  /* first time loop i = 0*/

  // y dy ddy: only them are stored
  y(i) = y0;
  dy(i) = 0; //dy0;
  double fs = ddy_ref + D*dy_ref + K*(y_ref-goal);
  ddy(i) = K*(goal-y(i)) - D*dy(i) + fs;

  vec info_back(3);
  info_back(0)  = fp;
  info_back(1)  = dfp;
  info_back(2)  = ddfp;

  return info_back;
}

 vec GDMP::run_solution_dt(double t_now, double goal, double ts, int i, double f_prev, double df_prev, double ddf_prev )
 {
     double fp  ;
     double dfp  ;
     double ddfp ;

     /* Compute fp */
     if (i > (NSamples-1))
     {
         fp = f_prev  ; //same as prev or (NSamples-1)
         dfp = df_prev ;
         ddfp = df_prev;
     }
     else
     {
         double s = 0.0 ;
         for (int l = 0; l < BFs ; l++)
         {
           s = s + w(l)*psi(l,i);
         }
         fp = k*s + y0;
         dfp =  (fp - f_prev)/ts;
         ddfp = (dfp - df_prev)/ts;
     }

     /* Compute refs */
     double y_ref  = ks*(fp - y0d_hat) + y0;
     double dy_ref  = ks*dfp ;
     double ddy_ref  = ks*ddfp ;

     /* Euler Solution dt*/
     dy(i) = ddy(i-1)*ts + dy(i-1);
     y(i) = dy(i)*ts + y(i-1);

     /* dynamical system */
     double fs = ddy_ref + D*dy_ref + K*(y_ref-goal);
     ddy(i) = K*(goal-y(i)) - D*dy(i) + fs;

     /* keep it for next loop dt*/
     vec info_back(3);
     info_back(0)  = fp;
     info_back(1)  = dfp;
     info_back(2)  = ddfp;

     return info_back;
 }

vec GDMP::init_Rsolution_dt(double goal, CanonicalStructure cs, double main_time, int extra_samples, double t_now, int i)
{
  /* Set size = NSamples + extra samples
      extra samples oq from extra time to run
  */
  Ry.set_size(1,NSamples+extra_samples);
  Rdy.set_size(1,NSamples+extra_samples);
  Rddy.set_size(1,NSamples+extra_samples);

  /* Swap gd,y0 */
  double Ry0 = goal;
  double Rgoal = y0;

  /* Swap gd,y0 hats */
  Ry0d_hat = Ry0; //gd_hat;
  Rgd_hat = Rgoal; //y0d_hat;

  /* init Rfp */
  double Rfp = gd_hat;
  double Rdfp = 0;
  double Rddfp = 0;

  /* Ref compute */
  double Ry_ref  = ks*(Rfp - Ry0d_hat) + Ry0;
  double Rdy_ref  = ks*Rdfp ;
  double Rddy_ref  = ks*Rddfp ;

  /* Init R solutions dt */
  Ry(i) = goal;
  Rdy(i) = 0; //Ry0;
  double Rfs = Rddy_ref + D*Rdy_ref + K*(Ry_ref-Rgoal);
  Rddy(i) = K*(Rgoal-Ry(i)) - D*Rdy(i) + Rfs;

/* keep it for the nect loop dt */
 vec info_back(3);
 info_back(0)  = Rfp;
 info_back(1)  = Rdfp;
 info_back(2)  = Rddfp;

 return info_back;
}

vec GDMP::run_Rsolution_dt(double t_now, double goal, double ts, int i, double f_prev, double df_prev, double ddf_prev )
{
    double Rfp  ;
    double Rdfp  ;
    double Rddfp ;
    double Ry0 = goal;
    double Rgoal = y0;

    /* Compute Rfp */
    if (i > (NSamples-1))
    {
        Rfp = f_prev  ; //same as prev or (NSamples-1)
        Rdfp = df_prev ;
        Rddfp = df_prev;
    }
    else
    {
        double s = 0.0 ;
        for (int l = 0; l < BFs ; l++)
        {
          s = s + w(l)*psi(l,NSamples - i );
        }

        Rfp = k*s + y0;
        Rdfp =  (Rfp - f_prev)/ts;
        Rddfp = (Rdfp - df_prev)/ts;
    }

    /* Compute R refs */
    double Ry_ref  = ks*(Rfp - Ry0d_hat) + Ry0;
    double Rdy_ref  = ks*Rdfp ;
    double Rddy_ref  = ks*Rddfp ;

    /* Euler Solution dt*/
    Rdy(i) = Rddy(i-1)*ts + Rdy(i-1);
    Ry(i) = Rdy(i)*ts + Ry(i-1);

    /* Dynamical system dt*/
    double Rfs = Rddy_ref + D*Rdy_ref + K*(Ry_ref-Rgoal);
    Rddy(i) = K*(Rgoal-Ry(i)) - D*Rdy(i) + Rfs;

    /* keep it for the nect loop dt */
    vec info_back(3);
    info_back(0)  = Rfp;
    info_back(1)  = Rdfp;
    info_back(2)  = Rddfp;

    return info_back;
}

/* fNyquist Function*/
double GDMP::fNyquistFunc(double *original, int size, double fs, double T)
{
  int L;
  L = (size % 2) == 0 ? size : size+1;

  vec signal_in(size);
  // cx_vec signal_xf(size);
  //
  // cx_mat XF_in(1,int(l1/2+1));
  // vec A1(size);
  //
  double f[int(L/2)];
  for ( int i = 0; i < (L/2); i++ )
      f[i] = fs*i/L;
  // complex<double> epi;
  // double EdXF = 0;
  // double  img_EdXF = 0;
  for (int i = 0; i<size ; i++)
    signal_in(i) = original[i];

  cx_vec P2(size);
  P2 = fft(signal_in) ;

  cx_vec Fd_original(int(L/2));
  for ( int i = 0; i < (L/2); i++ )
      Fd_original(i) = P2(i);

  double EdD = 0;
  for ( int i = 0; i < (L/2); i++)
  {
    if (f[i] < 40)
      EdD += abs(pow(Fd_original(i),2));
  }

  double h = 0.001;
  double EdDesired = (1-h)*EdD;
  double energy_inter = 0;

  int i = -1;
  while( energy_inter < EdDesired)
  {
    i++;
    energy_inter = energy_inter + abs(pow(Fd_original(i),2));
  }
  // cout<<i<<endl;
  double fNyq ;

  if ( i < 0)
  {
      fNyq = 2*f[1];
  }
  else
  {
      fNyq = 2*f[i];
  }
  return fNyq;
}