#include <Python.h>
#include <windows.h>
#ifdef min
#undef min
#undef max
#endif

#include <set>
#include <vector>
#include <numpy/arrayobject.h>
#include <maya/MGlobal.h>
#include <maya/MPoint.h>
#include <maya/MMatrix.h>
#include <maya/MQuaternion.h>
#include <maya/MTransformationMatrix.h>
#include <Eigen/Core>
#include <Eigen/SparseCore>
#include <Eigen/Eigen>
#include <tbb/blocked_range.h>
#include <tbb/parallel_for.h>
#include <osqp/osqp.h>
#include <osqp/auxil.h>

#pragma region Macros
inline void SET_LONG2(PyArrayObject* m, long i0, long i1, long v)
{
    *reinterpret_cast<npy_longlong*>(PyArray_GETPTR2(m, i0, i1)) = static_cast<npy_longlong>(v);
}
inline void SET_DOUBLE2(PyArrayObject* m, long i0, long i1, double v)
{
    *reinterpret_cast<npy_float64*>(PyArray_GETPTR2(m, i0, i1)) = static_cast<npy_float64>(v);
}
inline void SET_DOUBLE4(PyArrayObject* m, long i0, long i1, long i2, long i3, double v)
{
    *reinterpret_cast<npy_float64*>(PyArray_GETPTR4(m, i0, i1, i2, i3)) = static_cast<npy_float64>(v);
}
inline long GET_LONG2(PyArrayObject* m, long i0, long i1)
{
    return static_cast<long>(*reinterpret_cast<npy_longlong*>(PyArray_GETPTR2(m, i0, i1)));
}
inline double GET_DOUBLE2(PyArrayObject* m, long i0, long i1)
{
    return static_cast<double>(*reinterpret_cast<npy_float64*>(PyArray_GETPTR2(m, i0, i1)));
}
inline double GET_DOUBLE3(PyArrayObject* m, long i0, long i1, long i2)
{
    return static_cast<double>(*reinterpret_cast<npy_float64*>(PyArray_GETPTR3(m, i0, i1, i2)));
}
inline double GET_DOUBLE4(PyArrayObject* m, long i0, long i1, long i2, long i3)
{
    return static_cast<double>(*reinterpret_cast<npy_float64*>(PyArray_GETPTR4(m, i0, i1, i2, i3)));
}
#pragma endregion

struct Input
{
    //! number of model vertices
    long numVertices;
    //! number of sample frames
    long numSamples;
    //! desirable number of joints
    long numJoints;
    //! local neighbors (0 / 1 / 2)
    long numRings;
    //! rest/bind shape (numVertices)
    std::vector<MPoint> restShape;
    //! samples (numSamples x numVertices)
    std::vector<MPoint> sample;
    //! neighbor information
    std::vector<std::vector<std::pair<int, double>>> neighbor;
    //! model normalization matrix
    MMatrix normalizer;

    Input() :
        numVertices(0), numSamples(0), numJoints(0), normalizer(MMatrix::identity)
	{
	}
	~Input()
	{
		restShape.clear();
		sample.clear();
        neighbor.clear();
	}
};

struct Output
{
    //! number of inserted joints
    long numJoints;
    //! number of skinning joint index (max influence)
    long numIndices;
    //! skinning weight (numVertices x numIndices)
    std::vector<double> skinWeight;
    //! skinning joint index (numVertices x numIndices)
    std::vector<long> skinIndex;
    //! skinning matrix (numJoints x numSamples)
    std::vector<MTransformationMatrix> skinMatrix;
    //! rest joint position (numJoints)
    std::vector<MPoint> restJointPos;
    //! clusters each vertex belongs
    std::vector<std::vector<int>> vertCluster;

	Output() :
        numJoints(0), numIndices(0)
	{
	}
	~Output()
	{
		skinWeight.clear();
		skinIndex.clear();
		skinMatrix.clear();
        restJointPos.clear();
	}
};

#pragma region Subroutines for joint transformation
MTransformationMatrix calcRegistrationT(long numPoints, std::vector<MPoint>::const_iterator ps, std::vector<MPoint>::const_iterator pd)
{
    MTransformationMatrix transform;
    MPoint cs = MVector::zero, cd = MVector::zero;
    std::vector<MPoint>::const_iterator sit = ps;
    std::vector<MPoint>::const_iterator dit = pd;
    for (long i = 0; i < numPoints; ++i, ++sit, ++dit)
    {
        cs += *sit;
        cd += *dit;
    }
    transform.setTranslation((cd - cs) / static_cast<double>(numPoints), MSpace::kTransform);
    return transform;
}

MTransformationMatrix calcRegistrationRT(long numPoints, std::vector<MPoint>::const_iterator ps, std::vector<MPoint>::const_iterator pd)
{
    MTransformationMatrix transform;
    MPoint cs = MVector::zero, cd = MVector::zero;
    std::vector<MPoint>::const_iterator sit = ps;
    std::vector<MPoint>::const_iterator dit = pd;
    for (long i = 0; i < numPoints; ++i, ++sit, ++dit)
    {
        cs += *sit;
        cd += *dit;
    }
    cs = cs / static_cast<double>(numPoints);
    cd = cd / static_cast<double>(numPoints);
    if (numPoints < 3)
    {
        transform.setTranslation(cd - cs, MSpace::kTransform);
        return transform;
    }

    Eigen::Matrix<double, 4, 4> moment;
    double sxx = 0, sxy = 0, sxz = 0, syx = 0, syy = 0, syz = 0, szx = 0, szy = 0, szz = 0;
    sit = ps;
    dit = pd;
    for (long i = 0; i < numPoints; ++i, ++sit, ++dit)
    {
        sxx += (sit->x - cs.x) * (dit->x - cd.x);
        sxy += (sit->x - cs.x) * (dit->y - cd.y);
        sxz += (sit->x - cs.x) * (dit->z - cd.z);
        syx += (sit->y - cs.y) * (dit->x - cd.x);
        syy += (sit->y - cs.y) * (dit->y - cd.y);
        syz += (sit->y - cs.y) * (dit->z - cd.z);
        szx += (sit->z - cs.z) * (dit->x - cd.x);
        szy += (sit->z - cs.z) * (dit->y - cd.y);
        szz += (sit->z - cs.z) * (dit->z - cd.z);
    }
    moment(0, 0) = sxx + syy + szz;
    moment(0, 1) = syz - szy;        moment(1, 0) = moment(0, 1);
    moment(0, 2) = szx - sxz;        moment(2, 0) = moment(0, 2);
    moment(0, 3) = sxy - syx;        moment(3, 0) = moment(0, 3);
    moment(1, 1) = sxx - syy - szz;
    moment(1, 2) = sxy + syx;        moment(2, 1) = moment(1, 2);
    moment(1, 3) = szx + sxz;        moment(3, 1) = moment(1, 3);
    moment(2, 2) = -sxx + syy - szz;
    moment(2, 3) = syz + szy;        moment(3, 2) = moment(2, 3);
    moment(3, 3) = -sxx - syy + szz;

    if (moment.norm() > 0)
    {
        Eigen::EigenSolver<Eigen::Matrix<double, 4, 4>> es(moment);
        long maxi = 0;
        for (long i = 1; i < 4; ++i)
        {
            if (es.eigenvalues()(maxi).real() < es.eigenvalues()(i).real())
            {
                maxi = i;
            }
        }
        transform.setRotationQuaternion(
            es.eigenvectors()(1, maxi).real(),
            es.eigenvectors()(2, maxi).real(),
            es.eigenvectors()(3, maxi).real(),
            es.eigenvectors()(0, maxi).real());
    }
    MPoint cs0 = cs * transform.asMatrix();
    transform.setTranslation(cd - cs0, MSpace::kTransform);
    return transform;
}

MTransformationMatrix calcRegistrationSRT(long numPoints, std::vector<MPoint>::const_iterator ps, std::vector<MPoint>::const_iterator pd)
{
    const long numIterations = 50;
    const double threshold = 1.0e-8;
    const double scaleLowerBound = 1.0e-3;
    MTransformationMatrix transform = calcRegistrationRT(numPoints, ps, pd);

    MVector scale(1.0, 1.0, 1.0);
    MVector denom(0.0, 0.0, 0.0);
    std::vector<MPoint>::const_iterator sit = ps;
    for (long i = 0; i < numPoints; ++i, ++sit)
    {
        denom[0] += (*sit)[0] * (*sit)[0];
        denom[1] += (*sit)[1] * (*sit)[1];
        denom[2] += (*sit)[2] * (*sit)[2];
    }
    std::vector<MPoint> sstm(numPoints);
    sit = ps;
    for (long i = 0; i < numPoints; ++i, ++sit)
    {
        sstm[i][0] = (*sit)[0];
        sstm[i][1] = (*sit)[1];
        sstm[i][2] = (*sit)[2];
    }
    std::vector<MPoint> dstm(numPoints);
    for (long l = 0; l < numIterations; ++l)
    {
        MMatrix im = transform.asMatrix().inverse();
        std::vector<MPoint>::const_iterator dit = pd;
        for (long i = 0; i < numPoints; ++i, ++dit)
        {
            dstm[i] = *dit * im;
        }
        for (long d = 0; d < 3; ++d)
        {
            double ds = 0.0;
            sit = ps;
            for (long i = 0; i < numPoints; ++i, ++sit)
            {
                ds += (*sit)[d] * dstm[i][d];
            }
            scale[d] = denom[d] < threshold ? 1.0 : ds / denom[d];
            sit = ps;
            for (long i = 0; i < numPoints; ++i, ++sit)
            {
                sstm[i][d] = scale[d] * (*sit)[d];
            }
        }
        transform = calcRegistrationRT(numPoints, sstm.begin(), pd);
    }
    double sv[3] = { scale[0], scale[1], scale[2] };
    transform.setScale(sv, MSpace::kTransform);
    return transform;
}

typedef MTransformationMatrix(*RegistrationFunc)(long numPoints, std::vector<MPoint>::const_iterator ps, std::vector<MPoint>::const_iterator pd);

RegistrationFunc registrationFuncs[3] = {
    calcRegistrationT,
    calcRegistrationRT,
    calcRegistrationSRT };

void computeSamplePoints(std::vector<MPoint>& sample, long sid, long joint, const Output& output, const Input& input)
{
    const long numVertices = input.numVertices;
    const long& numIndices = output.numIndices;

    for (long v = 0; v < numVertices; ++v)
    {
        sample[v] = input.sample[sid * numVertices + v];
        const MPoint& s = input.restShape[v];
        for (long i = 0; i < numIndices; ++i)
        {
            const long jnt = output.skinIndex[v * numIndices + i];
            if (jnt >= 0 && jnt != joint)
            {
                const double w = output.skinWeight[v * numIndices + i];
                const MTransformationMatrix& at = output.skinMatrix[sid * input.numJoints + jnt];
                sample[v] -= w * (s * at.asMatrix());
            }
        }
    }
}

void subtractCentroid(std::vector<MPoint>& model, std::vector<MPoint>& sample, MPoint& corModel, MPoint& corSample, const Eigen::VectorXd& weight, const Output& output, const Input& input)
{
    const long numVertices = input.numVertices;

    double wsqsum = 0;
    corModel = MVector::zero;
    corSample = MVector::zero;
    for (long v = 0; v < numVertices; ++v)
    {
        const double w = weight[v];
        corModel += w * w * input.restShape[v];
        corSample += w * sample[v];
        wsqsum += w * w;
    }
    corModel = corModel / wsqsum;
    corSample = corSample / wsqsum;
    for (long v = 0; v < numVertices; ++v)
    {
        model[v] = weight[v] * (input.restShape[v] - corModel);
        sample[v] -= weight[v] * corSample;
    }
}
#pragma endregion

#pragma region Bone transformation
class JointTransformUpdator
{
private:
    Output* output;
    const Input* input;
    const Eigen::VectorXd* weight;
    int transformType;
    long joint;

public:
    JointTransformUpdator(Output* output_, const Input* input_,
        const Eigen::VectorXd* weight_, int joint_, int transformType_)
        : output(output_), input(input_),
        weight(weight_), joint(joint_), transformType(transformType_)
    {
    }
    void operator ()(const tbb::blocked_range<int>& range) const
    {
        std::vector<MPoint> model(input->numVertices);
        std::vector<MPoint> sample(input->numVertices);
        for (long s = range.begin(); s != range.end(); ++s)
        {
            if (s == 0)
            {
                output->skinMatrix[joint] = MTransformationMatrix::identity;
                continue;
            }
            computeSamplePoints(sample, s, joint, *output, *input);
            MPoint corModel(0, 0, 0), corSample(0, 0, 0);
            subtractCentroid(model, sample, corModel, corSample, *weight, *output, *input);
            MTransformationMatrix transform = registrationFuncs[transformType](model.size(), model.begin(), sample.begin());
            MVector d = corSample - corModel * transform.asMatrix();
            transform.setTranslation(d + transform.getTranslation(MSpace::kTransform), MSpace::kTransform);
            output->skinMatrix[s * input->numJoints + joint] = transform;
        }
    }
};

void updateJointTransformProc(Output& output, int transformType, const Input& input, long selection = -1)
{
    const long numVertices = input.numVertices;
    const long numSamples  = input.numSamples;
    const long numJoints    = output.numJoints;
    const long numIndices  = output.numIndices;

    Eigen::VectorXd weight = Eigen::VectorXd::Zero(numVertices);
    long begin = selection < 0 ? 0 : selection;
    long end = selection < 0 ? numJoints : selection + 1;
    for (long joint = begin; joint < end; ++joint)
    {
        for (long v = 0; v < numVertices; ++v)
        {
            weight[v] = 0.0;
            for (long i = 0; i < numIndices; ++i)
            {
                long jnt = output.skinIndex[v * numIndices + i];
                if (jnt == joint)
                {
                    weight[v] = output.skinWeight[v * numIndices + i];
                    break;
                }
            }
        }
        double wsqsum = weight.dot(weight);
        if (wsqsum > 1.0e-8)
        {
            tbb::blocked_range<int> blockedRange(0, numSamples);
            JointTransformUpdator transformUpdator(&output, &input, &weight, joint, transformType);
            tbb::parallel_for(blockedRange, transformUpdator);
        }
        else
        {
            for (int s = 0; s < numSamples; ++s)
            {
                output.skinMatrix[s * input.numJoints + joint] = MTransformationMatrix::identity;
            }
            char buf[256];
            sprintf(buf, "Fixed joint #%d", joint);
            MGlobal::displayInfo(buf);
        }
    }
}

static PyObject* updateJointTransform(PyObject *self, PyObject *args)
{
    PyObject* pin            = PyTuple_GET_ITEM(args, 0);
    PyObject* pout           = PyTuple_GET_ITEM(args, 1);
    const long transformType = PyInt_AsLong(PyTuple_GET_ITEM(args, 2));
    Input* const input   = reinterpret_cast<Input*>(PyCapsule_GetPointer(pin, "SSDSInput"));
    Output* const output = reinterpret_cast<Output*>(PyCapsule_GetPointer(pout, "SSDSOutput"));
    updateJointTransformProc(*output, transformType, *input);
    return Py_None;
}
#pragma endregion

#pragma region Skinning weight
Eigen::VectorXd solveWeightQP(Eigen::SparseMatrix<c_float, 0, c_int>& pm, Eigen::VectorXd& qv)
{
    const int n = static_cast<int>(pm.rows());
    Eigen::SparseMatrix<c_float, 0, c_int> a(n + 1, n);
    std::vector<c_float> q(n), l(n + 1, 0.0), u(n + 1, 1.0);
    for (int i = 0; i < n; ++i)
    {
        a.coeffRef(0, i) = 1.0;
        a.coeffRef(i + 1, i) = 1.0;
        q[i] = qv[i];
    }
    l[0] = u[0] = 1.0;
    a.makeCompressed();
    pm.makeCompressed();

    // Problem settings
    OSQPSettings settings;
    osqp_set_default_settings(&settings);
    settings.polish = 0;
    settings.adaptive_rho = 0;
    settings.warm_start = 0;
    settings.eps_abs = 1.0e-12; // experimentally set
    settings.eps_rel = 1.0e-12; // experimentally set

    OSQPData osqpData;
    osqpData.n = n;
    osqpData.m = n + 1;
    osqpData.P = csc_matrix(n, n, pm.nonZeros(), pm.valuePtr(), pm.innerIndexPtr(), pm.outerIndexPtr());
    osqpData.q = q.data();
    osqpData.A = csc_matrix(n + 1, n, a.nonZeros(), a.valuePtr(), a.innerIndexPtr(), a.outerIndexPtr());
    osqpData.l = l.data();
    osqpData.u = u.data();

    OSQPWorkspace* osqpWork = osqp_setup(&osqpData, &settings);
    c_int retval = osqp_solve(osqpWork);
    settings.rho = compute_rho_estimate(osqpWork);
    Eigen::VectorXd weight(n);
    for (int i = 0; i < n; ++i)
    {
        weight[i] = osqpWork->solution->x[i];
    }
    osqp_cleanup(osqpWork);

    settings.warm_start = 1;
    osqpWork = osqp_setup(&osqpData, &settings);
    osqp_warm_start_x(osqpWork, weight.data());
    retval = osqp_solve(osqpWork);
    for (int i = 0; i < n; ++i)
    {
        weight[i] = osqpWork->solution->x[i];
    }
    osqp_cleanup(osqpWork);
    c_free(osqpData.A);
    c_free(osqpData.P);
    return weight;
}

class SkinWeightUpdator
{
private:
    Output* output;
    const Input* input;
    const int numIndices;
    const double smoothness;
    std::vector<double> prevWeight;
    std::vector<long> prevIndex;
public:
    SkinWeightUpdator(Output* output_, const Input* input_,
        int numMaxInfluences_, double smoothness_)
        : output(output_), input(input_),
        numIndices(numMaxInfluences_), smoothness(smoothness_ == 0 ? 0 : std::pow(10.0, smoothness_ - 8))
    {
        prevWeight = output->skinWeight;
        prevIndex = output->skinIndex;
    }
    void operator ()(const tbb::blocked_range<int>& range) const
    {
        const int numVertices = input->numVertices;
        const int numSamples  = input->numSamples;
        const int numJoints   = output->numJoints;
        const double epsilon  = 1.0e-3;

        Eigen::MatrixXd sbasis   = Eigen::MatrixXd::Zero(numIndices, numSamples * 3);
        Eigen::VectorXd sample   = Eigen::VectorXd::Zero(numSamples * 3);
        Eigen::VectorXd slaplacm = Eigen::VectorXd::Zero(numIndices);
        Eigen::VectorXd slaplacv = Eigen::VectorXd::Zero(numIndices);
        for (int v = range.begin(); v != range.end(); ++v)
        {
            const std::vector<int>& vertCluster = output->vertCluster[v];
            const int numActJoints = vertCluster.size();
            Eigen::MatrixXd basis   = Eigen::MatrixXd::Zero(numActJoints, numSamples * 3);
            Eigen::VectorXd laplacm = Eigen::VectorXd::Zero(numActJoints);
            Eigen::VectorXd laplacv = Eigen::VectorXd::Zero(numActJoints);

            for (int i = 0; i < numIndices; ++i)
            {
                output->skinIndex[v * numIndices + i] = -1;
                output->skinWeight[v * numIndices + i] = 0;
            }
            // Laplacian term
            double lsum = 0.0;
            laplacv.setZero();

            for (auto it = input->neighbor[v].begin(); it != input->neighbor[v].end(); ++it)
            {
                const int nvid = it->first;
                lsum -= it->second;
                for (int i = 0; i < numIndices; ++i)
                {
                    const int nbid = prevIndex[nvid * numIndices + i];
                    auto nbit = std::find(output->vertCluster[v].begin(), output->vertCluster[v].end(), nbid);
                    if (nbid >= 0 && nbit != output->vertCluster[v].end())
                    {
                        laplacv[nbit - output->vertCluster[v].begin()] -= it->second * prevWeight[nvid * numIndices + i];
                    }
                }
            }
            laplacm.setConstant(smoothness * lsum);
            laplacv *= -smoothness;

            // variables
            const MPoint& restVertex = input->restShape[v];
            for (int s = 0; s < numSamples; ++s)
            {
                for (int j = 0; j < numActJoints; ++j)
                {
                    const int jid = vertCluster[j];
                    const MTransformationMatrix& rt = output->skinMatrix[s * numJoints + jid];
                    MPoint tv = restVertex * rt.asMatrix();
                    basis(j, s * 3 + 0) = tv.x;
                    basis(j, s * 3 + 1) = tv.y;
                    basis(j, s * 3 + 2) = tv.z;
                }
                sample[s * 3 + 0] = input->sample[s * numVertices + v].x;
                sample[s * 3 + 1] = input->sample[s * numVertices + v].y;
                sample[s * 3 + 2] = input->sample[s * numVertices + v].z;
            }
            // optimization
            Eigen::SparseMatrix<c_float, 0, c_int> gm =
                (basis * basis.transpose()
                    + laplacm.asDiagonal().toDenseMatrix()
                    ).sparseView();
            gm.makeCompressed();
            Eigen::VectorXd gv = -basis * sample + laplacv;
            double scale = std::max(std::abs(gv.minCoeff()), std::abs(gv.maxCoeff()));

            Eigen::VectorXd weight = solveWeightQP(gm, gv);
            double weightSum = 0;
            for (int i = 0; i < numIndices; ++i)
            {
                double maxw = epsilon;
                int bestjoint = -1;
                for (int j = 0; j < numActJoints; ++j)
                {
                    if (weight[j] > maxw)
                    {
                        maxw = weight[j];
                        bestjoint = j;
                    }
                }
                if (bestjoint < 0)
                {
                    break;
                }
                output->skinIndex[v * numIndices + i] = bestjoint;
                output->skinWeight[v * numIndices + i] = maxw;
                weightSum += maxw;
                weight[bestjoint] = -1.0;
            }
            // subproblem
            if (weightSum < 1.0)
            {
                sbasis.setZero();
                slaplacm.setZero();
                slaplacv.setZero();
                for (int i = 0; i < numIndices; ++i)
                {
                    int j = output->skinIndex[v * numIndices + i];
                    if (j >= 0)
                    {
                        for (int s = 0; s < numSamples * 3; ++s)
                        {
                            sbasis(i, s) = basis(j, s);
                        }
                        slaplacm[i] = laplacm[j];
                        slaplacv[i] = laplacv[j];
                    }
                }
                Eigen::SparseMatrix<c_float, 0, c_int> sgm =
                    (sbasis * sbasis.transpose()
                        + slaplacm.asDiagonal().toDenseMatrix()
                        ).sparseView();
                Eigen::VectorXd sgv = -sbasis * sample + slaplacv;
                sgm.makeCompressed();
                weight = solveWeightQP(sgm, sgv);
                weightSum = 0;
                for (int i = 0; i < numIndices; ++i)
                {
                    if (weight[i] < epsilon)
                    {
                        output->skinWeight[v * numIndices + i] = 0;
                        output->skinIndex[v * numIndices + i] = -1;
                    }
                    else
                    {
                        output->skinWeight[v * numIndices + i] = weight[i];
                        weightSum += weight[i];
                    }
                }
            }
            for (int i = 0; i < numIndices; ++i)
            {
                output->skinIndex[v * numIndices + i] = output->skinIndex[v * numIndices + i] < 0
                    ? -1 : vertCluster[output->skinIndex[v * numIndices + i]];
                output->skinWeight[v * numIndices + i] /= weightSum;
            }
        }
    }
};

static PyObject* updateSkinWeight(PyObject *self, PyObject *args)
{
    PyObject* pin        = PyTuple_GET_ITEM(args, 0);
    PyObject* pout       = PyTuple_GET_ITEM(args, 1);
    double smoothness   = PyFloat_AsDouble(PyTuple_GET_ITEM(args, 2));
    Input* const input   = reinterpret_cast<Input*>(PyCapsule_GetPointer(pin, "SSDSInput"));
	Output* const output = reinterpret_cast<Output*>(PyCapsule_GetPointer(pout, "SSDSOutput"));
    SkinWeightUpdator weightUpdater(output, input, output->numIndices, smoothness);
    tbb::blocked_range<int> blockedRange(0, input->numVertices);
	tbb::parallel_for(blockedRange, weightUpdater);
    return Py_None;
}
#pragma endregion

#pragma region Clustering
// see https://sites.google.com/view/fumiyanarita/project/la_ssdr_mdmc
static void detectNeighborClusters(const Input& input, Output& output)
{
    MGlobal::displayInfo("Detecting neighbor clusters");
    const long& numIndices = output.numIndices;

    std::vector<std::set<int>> clusters(output.numJoints);
    for (int v = 0; v < input.numVertices; ++v)
    {
        const int primjid = output.skinIndex[v * numIndices];
        if (input.numRings == 0)
        {
            for (int j = 0; j < output.numJoints; ++j)
            {
                for (int v = 0; v < output.numJoints; ++v)
                {
                    clusters[j].insert(v);
                }
            }
        }
        else
        {
            clusters[primjid].insert(0);
            clusters[primjid].insert(primjid);
            // one-ring neighbors
            for (auto nvit = input.neighbor[v].begin(); nvit != input.neighbor[v].end(); ++nvit)
            {
                const int nvid = nvit->first;
                clusters[primjid].insert(output.skinIndex[nvid * numIndices]);
                if (input.numRings == 1)
                {
                    continue;
                }
                // two-ring neighbors
                for (auto nnvit = input.neighbor[nvid].begin(); nnvit != input.neighbor[nvid].end(); ++nnvit)
                {
                    const int nnvid = nnvit->first;
                    clusters[primjid].insert(output.skinIndex[nnvid * numIndices]);
                }
            }
        }
    }
    for (int v = 0; v < input.numVertices; ++v)
    {
        const int primjid = output.skinIndex[v * numIndices];
        output.vertCluster[v] = std::vector<int>(clusters[primjid].begin(), clusters[primjid].end());
    }
}

long bindVertexToJoint(Output& output, const Input& input)
{
    const long numVertices = input.numVertices;
    const long numSamples  = input.numSamples;
    const long numIndices  = output.numIndices;

    std::vector<int> numJointVertices(output.numJoints, 0);
    for (long v = 0; v < numVertices; ++v)
    {
        long bestJoint = 0;
        double minErr = std::numeric_limits<double>::max();
        const MPoint& restShapePos = input.restShape[v];
        for (long j = 0; j < output.numJoints; ++j)
        {
            double errsq = 0;
            for (long s = 0; s < numSamples; ++s)
            {
                MMatrix am = output.skinMatrix[s * input.numJoints + j].asMatrix();
                MVector diff = input.sample[s * numVertices + v] - restShapePos * am;
                errsq += diff * diff;
            }
            errsq *= (input.restShape[v] - output.restJointPos[j]).length();
            if (errsq < minErr)
            {
                bestJoint = j;
                minErr = errsq;
            }
        }
        ++numJointVertices[bestJoint];
        output.skinIndex[v * numIndices + 0] = bestJoint;
    }

    std::vector<int>::iterator smallestBoneSize = std::min_element(numJointVertices.begin(), numJointVertices.end());
    while (*smallestBoneSize <= 0)
    {
        const long smallestBone = static_cast<int>(smallestBoneSize - numJointVertices.begin());
        numJointVertices.erase(numJointVertices.begin() + smallestBone);
        output.restJointPos.erase(output.restJointPos.begin() + smallestBone);
        for (long s = 0; s < input.numSamples; ++s)
        {
            for (int j = output.numJoints - 2; j >= smallestBone; --j)
            {
                output.skinMatrix[s * input.numJoints + j] = output.skinMatrix[s * input.numJoints + j + 1];
            }
        }
        for (long v = 0; v < numVertices; ++v)
        {
            if (output.skinIndex[v * numIndices + 0] >= smallestBone)
            {
                --output.skinIndex[v * numIndices + 0];
            }
        }
        --output.numJoints;
        smallestBoneSize = std::min_element(numJointVertices.begin(), numJointVertices.end());
    }
    return static_cast<int>(numJointVertices.size());
}

long findMostStableVertex(const Input& input)
{
    const long numVertices = input.numVertices;
    const long numSamples = input.numSamples;
    double minErrorSq = std::numeric_limits<double>::max();
    long mostStableVertex = -1;
    for (int v = 0; v < numVertices; ++v)
    {
        double errSq = 0;
        for (int s = 0; s < numSamples; ++s)
        {
            MVector diff = input.sample[s * numVertices + v] - input.restShape[v];
            errSq += diff * diff;
        }
        if (errSq < minErrorSq)
        {
            minErrorSq = errSq;
            mostStableVertex = v;
        }
    }
    return mostStableVertex;
}

long findDistantVertex(const Input& input, const Output& output, const std::set<long>& covered)
{
    const long numVertices = input.numVertices;
    const long numSamples  = input.numSamples;
    const long numIndices  = output.numIndices;
    double maxErrorSq      = -std::numeric_limits<double>::max();
    long mostDistantVertex = -1;
    for (int v = 0; v < numVertices; ++v)
    {
        if (covered.find(v) != covered.end())
        {
            continue;
        }
        long index = output.skinIndex[v * numIndices + 0];
        double errSq = 0;
        for (int s = 0; s < numSamples; ++s)
        {
            MMatrix bm = output.skinMatrix[s * input.numJoints + index].asMatrix();
            MVector diff = input.sample[s * numVertices + v] - input.restShape[v] * bm;
            errSq += diff * diff;
        }
        if (errSq > maxErrorSq)
        {
            maxErrorSq = errSq;
            mostDistantVertex = v;
        }
    }
    return mostDistantVertex;
}

// solving p-center problem
static PyObject* clusterVerticesPcenter(PyObject* self, PyObject* args)
{
    PyObject* pin = PyTuple_GET_ITEM(args, 0);
    PyObject* pout = PyTuple_GET_ITEM(args, 1);
    const long transformType = PyInt_AsLong(PyTuple_GET_ITEM(args, 2));
    const Input& input = *reinterpret_cast<Input*>(PyCapsule_GetPointer(pin, "SSDSInput"));
    Output& output = *reinterpret_cast<Output*>(PyCapsule_GetPointer(pout, "SSDSOutput"));
    const long& numIndices = output.numIndices;

    output.numJoints = input.numJoints;
    std::fill(output.skinIndex.begin(), output.skinIndex.end(), -1);
    std::fill(output.skinWeight.begin(), output.skinWeight.end(), 0.0);
    std::fill(output.skinMatrix.begin(), output.skinMatrix.end(), MTransformationMatrix::identity);
    std::fill(output.restJointPos.begin(), output.restJointPos.end(), MPoint::origin);

    char buf[512];
    long stableVertex = findMostStableVertex(input);
    output.skinIndex[stableVertex * numIndices] = 0;
    output.skinWeight[stableVertex * numIndices] = 1.0;
    output.restJointPos[0] = input.restShape[stableVertex];
    std::vector<long> jointVertices(output.numJoints);
    jointVertices[0] = stableVertex;
    sprintf(buf, "Added to Vertex %d (stable)", stableVertex);
    MGlobal::displayInfo(buf);

    for (int jid = 1; jid < input.numJoints; ++jid)
    {
        double maxDist = -1.0;
        int distantVertex = -1;
        for (int i = 0; i < input.numVertices; ++i)
        {
            if (std::find(jointVertices.begin(), jointVertices.end(), i) != jointVertices.end())
            {
                continue;
            }
            double mindist = std::numeric_limits<double>::max();
            for (int j = 0; j < jid; ++j)
            {
                const int joint = jointVertices[j];
                MVector v = input.restShape[i] - input.restShape[joint];
                if (mindist > v.length())
                {
                    mindist = v.length();
                }
            }
            if (mindist > maxDist)
            {
                maxDist = mindist;
                distantVertex = i;
            }
        }
        output.restJointPos[jid] = input.restShape[distantVertex];
        jointVertices[jid] = distantVertex;
        sprintf(buf, "Added to Vertex %d, %f", distantVertex, maxDist);
        MGlobal::displayInfo(buf);
    }

    for (int i = 0; i < input.numVertices; ++i)
    {
        double minDist = std::numeric_limits<double>::max();
        int nearestJoint = -1;
        for (int j = 0; j < output.numJoints; ++j)
        {
            const int joint = jointVertices[j];
            MVector v = input.restShape[i] - input.restShape[joint];
            if (v.length() < minDist)
            {
                minDist = v.length();
                nearestJoint = j;
            }
        }
        output.skinIndex[i * numIndices] = nearestJoint;
        output.skinWeight[i * numIndices] = 1.0;
    }
    detectNeighborClusters(input, output);
    updateJointTransformProc(output, transformType, input);
    return PyInt_FromLong(output.numJoints);
}

static PyObject* clusterVerticesAdaptive(PyObject *self, PyObject *args)
{
    PyObject* pin            = PyTuple_GET_ITEM(args, 0);
    PyObject* pout           = PyTuple_GET_ITEM(args, 1);
    const long transformType = PyInt_AsLong(PyTuple_GET_ITEM(args, 2));
    const Input& input = *reinterpret_cast<Input*>(PyCapsule_GetPointer(pin, "SSDSInput"));
    Output& output     = *reinterpret_cast<Output*>(PyCapsule_GetPointer(pout, "SSDSOutput"));
    const long& numIndices = output.numIndices;

    output.numJoints   = 0;
    std::fill(output.skinIndex.begin(), output.skinIndex.end(), -1);
    std::fill(output.skinWeight.begin(), output.skinWeight.end(), 0.0);
    std::fill(output.skinMatrix.begin(), output.skinMatrix.end(), MTransformationMatrix::identity);
    output.restJointPos.clear();

    std::set<long> covered;
    char buf[512];

    long stableVertex = findMostStableVertex(input);
    output.skinIndex[stableVertex * numIndices] = 0;
    output.skinWeight[stableVertex * numIndices] = 1.0;
    output.restJointPos.push_back(input.restShape[stableVertex]);
    covered.insert(stableVertex);
    sprintf(buf, "Added to Vertex %d (stable)", stableVertex);
    MGlobal::displayInfo(buf);

    for (int i = 0; i < input.neighbor[stableVertex].size(); ++i)
    {
        long neighbor = input.neighbor[stableVertex][i].first;
        if (neighbor < 0)
        {
            continue;
        }
        output.skinIndex[neighbor * numIndices] = 0;
        output.skinWeight[neighbor * numIndices] = 1.0;
        covered.insert(neighbor);
    }
    updateJointTransformProc(output, transformType, input, 0);
    for (int v = 0; v < input.numVertices; ++v)
    {
        output.skinIndex[v * numIndices] = 0;
        output.skinWeight[v * numIndices] = 1.0;
    }
    output.numJoints = 1;

    for (int iteration = 0; iteration < input.numJoints + 10; ++iteration)
    {
        if (output.numJoints >= input.numJoints || covered.size() >= input.numVertices)
        {
            break;
        }
        long distantVertex = findDistantVertex(input, output, covered);
        output.skinIndex[distantVertex * numIndices] = output.numJoints;
        output.restJointPos.push_back(input.restShape[distantVertex]);
        covered.insert(distantVertex);
        for (int i = 0; i < input.neighbor[distantVertex].size(); ++i)
        {
            long neighbor = input.neighbor[distantVertex][i].first;
            if (neighbor < 0)
            {
                continue;
            }
            output.skinIndex[neighbor * numIndices] = output.numJoints;
            covered.insert(neighbor);
        }
        updateJointTransformProc(output, transformType, input, output.numJoints);
        ++output.numJoints;
        output.numJoints = bindVertexToJoint(output, input);

        sprintf(buf, "Added to Vertex %d", distantVertex);
        MGlobal::displayInfo(buf);
    }
    detectNeighborClusters(input, output);
    return PyInt_FromLong(output.numJoints);
}
#pragma endregion

#pragma region Initialization
static PyObject* initialize(PyObject *self, PyObject *args)
{
    PyArrayObject* initPos          = reinterpret_cast<PyArrayObject*>(PyTuple_GET_ITEM(args, 0));
    PyArrayObject* shapeSample      = reinterpret_cast<PyArrayObject*>(PyTuple_GET_ITEM(args, 1));
    PyArrayObject* neighborVertices = reinterpret_cast<PyArrayObject*>(PyTuple_GET_ITEM(args, 2));
    const long numJoints            = PyInt_AsLong(PyTuple_GET_ITEM(args, 3));
    const long numIndices           = PyInt_AsLong(PyTuple_GET_ITEM(args, 4));
    const long numRings             = PyInt_AsLong(PyTuple_GET_ITEM(args, 5));
    const long numVertices = static_cast<long>(initPos->dimensions[0]);
    const long numSamples  = static_cast<long>(shapeSample->dimensions[0]);
    // allocation
    Input* input       = new Input();
    input->numVertices = numVertices;
    input->numSamples  = numSamples;
    input->numJoints   = numJoints;
    input->numRings    = numRings;
    input->restShape   = std::vector<MPoint>(numVertices);
    input->sample      = std::vector<MPoint>(numSamples * numVertices);
    input->neighbor    = std::vector<std::vector<std::pair<int, double>>>(numVertices);
    Output* output       = new Output();
    output->numJoints    = 0;
    output->numIndices   = numIndices;
    output->skinIndex    = std::vector<long>(numVertices * numIndices, -1);
    output->skinWeight   = std::vector<double>(numVertices * numIndices, 0.0);
    output->skinMatrix   = std::vector<MTransformationMatrix>(numSamples * numJoints, MTransformationMatrix::identity);
    output->restJointPos = std::vector<MPoint>(numJoints);
    output->vertCluster  = std::vector<std::vector<int>>(numVertices);
    // size normalization
    MVector centroid(0, 0, 0);
    MVector minPos(1.0e10, 1.0e10, 1.0e10), maxPos(-1.0e10, -1.0e10, -1.0e10);
    for (int v = 0; v < numVertices; ++v)
    {
        MVector p(GET_DOUBLE2(initPos, v, 0),
                  GET_DOUBLE2(initPos, v, 1),
                  GET_DOUBLE2(initPos, v, 2));
        input->restShape[v] = p;
        centroid += p;
        for (int i = 0; i < 3; ++i)
        {
            minPos[i] = std::min(minPos[i], p[i]);
            maxPos[i] = std::max(maxPos[i], p[i]);
        }
    }
    centroid /= numVertices;
    double scale = std::max(std::max(maxPos.x - minPos.x, maxPos.y - minPos.y), maxPos.z - minPos.z);
    MMatrix m = MMatrix::identity;
    m(0, 0) = m(1, 1) = m(2, 2) = 1.0 / scale;
    m(3, 0) = -centroid.x / scale;
    m(3, 1) = -centroid.y / scale;
    m(3, 2) = -centroid.z / scale;
    input->normalizer = m;
    // normalized samples
    for (int v = 0; v < numVertices; ++v)
    {
        input->restShape[v] = input->restShape[v] * input->normalizer;
        for (int s = 0; s < numSamples; ++s)
        {
            MPoint p(GET_DOUBLE3(shapeSample, s, v, 0),
                     GET_DOUBLE3(shapeSample, s, v, 1),
                     GET_DOUBLE3(shapeSample, s, v, 2));
            input->sample[s * numVertices + v] = p * input->normalizer;
        }
    }
    // one-ring neighbor [Le and Deng 2014]
    for (int v = 0; v < numVertices; ++v)
    {
		double lv = 0.0;
        for (int i = 0; i < neighborVertices->dimensions[1]; ++i)
        {
            long n = GET_LONG2(neighborVertices, v, i);
            if (n < 0)
            {
                continue;
            }
            double len = (input->restShape[v] - input->restShape[n]).length();
            if (len > 0)
            {
				// [Le and Deng 2014]
				double dsqsum = 0;
				for (int s = 0; s < input->numSamples; ++s)
				{
					float diff = (input->restShape[v] - input->restShape[n]).length()
						- (input->sample[s * numVertices + v] - input->sample[s * numVertices + n]).length();
					dsqsum += diff * diff;
				}
				dsqsum += 1.0e-10;
				double dvn = 1.0 / std::sqrt(dsqsum / input->numSamples);
                input->neighbor[v].push_back(std::make_pair(n, dvn));
                lv += dvn;
            }
        }
		for (int n = 0; n < input->neighbor[v].size(); ++n)
		{
			input->neighbor[v][n].second /= -lv;
		}
	}
    PyObject* retval = PyTuple_New(2);
    PyTuple_SetItem(retval, 0, PyCapsule_New(input,  "SSDSInput", nullptr));
    PyTuple_SetItem(retval, 1, PyCapsule_New(output, "SSDSOutput", nullptr));
    return retval;
}

static PyObject* release(PyObject *self, PyObject *args)
{
	PyObject* pin = PyTuple_GET_ITEM(args, 0);
	PyObject* pout = PyTuple_GET_ITEM(args, 1);
	delete reinterpret_cast<Input*>(PyCapsule_GetPointer(pin, "SSDSInput"));
	delete reinterpret_cast<Output*>(PyCapsule_GetPointer(pout, "SSDSOutput"));
	return Py_None;
}

static PyObject* retrieveResult(PyObject *self, PyObject *args)
{
	PyObject* pin                = PyTuple_GET_ITEM(args, 0);
	PyObject* pout               = PyTuple_GET_ITEM(args, 1);
	PyArrayObject* skinIndex     = reinterpret_cast<PyArrayObject*>(PyTuple_GET_ITEM(args, 2));
	PyArrayObject* skinWeight    = reinterpret_cast<PyArrayObject*>(PyTuple_GET_ITEM(args, 3));
	PyArrayObject* skinMatrix    = reinterpret_cast<PyArrayObject*>(PyTuple_GET_ITEM(args, 4));
    PyArrayObject* restJointPos = reinterpret_cast<PyArrayObject*>(PyTuple_GET_ITEM(args, 5));
    Input* const input     = reinterpret_cast<Input*>(PyCapsule_GetPointer(pin, "SSDSInput"));
	Output* const output   = reinterpret_cast<Output*>(PyCapsule_GetPointer(pout, "SSDSOutput"));
    const long& numIndices = output->numIndices;

    MMatrix inm = input->normalizer.inverse();
	for (int s = 0; s < input->numSamples; ++s)
	{
		for (int b = 0; b < output->numJoints; ++b)
		{
			MMatrix m = input->normalizer * output->skinMatrix[s * output->numJoints + b].asMatrix() * inm;
			for (int i = 0; i < 4; ++i)
			{
				for (int j = 0; j < 4; ++j)
				{
					SET_DOUBLE4(skinMatrix, b, s, i, j, m(i, j));
				}
			}
        }
	}
    for (int v = 0; v < input->numVertices; ++v)
    {
        for (long i = 0; i < numIndices; ++i)
        {
            SET_LONG2(skinIndex, v, i, output->skinIndex[v * numIndices + i]);
            SET_DOUBLE2(skinWeight, v, i, output->skinWeight[v * numIndices + i]);
        }
    }
    for (int j = 0; j < output->numJoints; ++j)
    {
        MPoint cp = output->restJointPos[j] * inm;
        SET_DOUBLE2(restJointPos, j, 0, cp.x);
        SET_DOUBLE2(restJointPos, j, 1, cp.y);
        SET_DOUBLE2(restJointPos, j, 2, cp.z);
    }
	return Py_None;
}
#pragma endregion

static PyMethodDef methods[] = {
    { "initialize", initialize, METH_VARARGS, "initialize" },
    { "clusterVerticesPcenter", clusterVerticesPcenter, METH_VARARGS, "clusterVerticesPcenter" },
    { "clusterVerticesAdaptive", clusterVerticesAdaptive, METH_VARARGS, "clusterVerticesAdaptive" },
    { "updateJointTransform", updateJointTransform, METH_VARARGS, "updateJointTransform" },
    { "updateSkinWeight", updateSkinWeight, METH_VARARGS, "updateSkinWeight" },
	{ "retrieveResult", retrieveResult, METH_VARARGS, "retrieveResult" },
	{ "release", release, METH_VARARGS, "release" },
	{ NULL, NULL, 0, NULL }
};

PyMODINIT_FUNC initnative(void)
{
    Py_InitModule("native", methods);
    import_array();
	Eigen::initParallel();
}


//=============================================================================================================
/**
 * @file     fiff_cov.cpp
 * @author   Lorenz Esch <lesch@mgh.harvard.edu>;
 *           Matti Hamalainen <msh@nmr.mgh.harvard.edu>;
 *           Christoph Dinh <chdinh@nmr.mgh.harvard.edu>
 * @since    0.1.0
 * @date     July, 2012
 *
 * @section  LICENSE
 *
 * Copyright (C) 2012, Lorenz Esch, Matti Hamalainen, Christoph Dinh. All rights reserved.
 *
 * Redistribution and use in source and binary forms, with or without modification, are permitted provided that
 * the following conditions are met:
 *     * Redistributions of source code must retain the above copyright notice, this list of conditions and the
 *       following disclaimer.
 *     * Redistributions in binary form must reproduce the above copyright notice, this list of conditions and
 *       the following disclaimer in the documentation and/or other materials provided with the distribution.
 *     * Neither the name of MNE-CPP authors nor the names of its contributors may be used
 *       to endorse or promote products derived from this software without specific prior written permission.
 *
 * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED
 * WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A
 * PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT,
 * INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
 * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION)
 * HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING
 * NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
 * POSSIBILITY OF SUCH DAMAGE.
 *
 *
 * @brief    Definition of the FiffCov Class.
 *
 */

//=============================================================================================================
// INCLUDES
//=============================================================================================================

#include "fiff_cov.h"
#include "fiff_stream.h"
#include "fiff_info_base.h"
#include "fiff_dir_node.h"

#include <utils/mnemath.h>

//=============================================================================================================
// QT INCLUDES
//=============================================================================================================

#include <QPair>

//=============================================================================================================
// EIGEN INCLUDES
//=============================================================================================================

#include <Eigen/SVD>

//=============================================================================================================
// USED NAMESPACES
//=============================================================================================================

using namespace FIFFLIB;
using namespace UTILSLIB;
using namespace Eigen;

//=============================================================================================================
// DEFINE MEMBER METHODS
//=============================================================================================================

FiffCov::FiffCov()
: kind(-1)
, diag(false)
, dim(-1)
, nfree(-1)
{
    qRegisterMetaType<QSharedPointer<FIFFLIB::FiffCov> >("QSharedPointer<FIFFLIB::FiffCov>");
    qRegisterMetaType<FIFFLIB::FiffCov>("FIFFLIB::FiffCov");
}

//=============================================================================================================

FiffCov::FiffCov(QIODevice &p_IODevice)
: kind(-1)
, diag(false)
, dim(-1)
, nfree(-1)
{
    FiffStream::SPtr t_pStream(new FiffStream(&p_IODevice));

    if(!t_pStream->open())
    {
        printf("\tNot able to open IODevice.\n");//ToDo Throw here
        return;
    }

    if(!t_pStream->read_cov(t_pStream->dirtree(), FIFFV_MNE_NOISE_COV, *this))
        printf("\tFiff covariance not found.\n");//ToDo Throw here

    qRegisterMetaType<QSharedPointer<FIFFLIB::FiffCov> >("QSharedPointer<FIFFLIB::FiffCov>");
    qRegisterMetaType<FIFFLIB::FiffCov>("FIFFLIB::FiffCov");
}

//=============================================================================================================

FiffCov::FiffCov(const FiffCov &p_FiffCov)
: QSharedData(p_FiffCov)
, kind(p_FiffCov.kind)
, diag(p_FiffCov.diag)
, dim(p_FiffCov.dim)
, names(p_FiffCov.names)
, data(p_FiffCov.data)
, projs(p_FiffCov.projs)
, bads(p_FiffCov.bads)
, nfree(p_FiffCov.nfree)
, eig(p_FiffCov.eig)
, eigvec(p_FiffCov.eigvec)
{
    qRegisterMetaType<QSharedPointer<FIFFLIB::FiffCov> >("QSharedPointer<FIFFLIB::FiffCov>");
    qRegisterMetaType<FIFFLIB::FiffCov>("FIFFLIB::FiffCov");
}

//=============================================================================================================

FiffCov::~FiffCov()
{
}

//=============================================================================================================

void FiffCov::clear()
{
    kind = -1;
    diag = false;
    dim = -1;
    names.clear();
    data = MatrixXd();
    projs.clear();
    bads.clear();
    nfree = -1;
    eig = VectorXd();
    eigvec = MatrixXd();
}

//=============================================================================================================

FiffCov FiffCov::pick_channels(const QStringList &p_include, const QStringList &p_exclude)
{
    RowVectorXi sel = FiffInfoBase::pick_channels(this->names, p_include, p_exclude);
    FiffCov res;//No deep copy here - since almost everything else is adapted anyway

    res.kind = this->kind;
    res.diag = this->diag;
    res.dim = sel.size();

    for(qint32 k = 0; k < sel.size(); ++k)
        res.names << this->names[sel(k)];

    res.data.resize(res.dim, res.dim);
    for(qint32 i = 0; i < res.dim; ++i)
        for(qint32 j = 0; j < res.dim; ++j)
            res.data(i, j) = this->data(sel(i), sel(j));
    res.projs = this->projs;

    for(qint32 k = 0; k < this->bads.size(); ++k)
        if(res.names.contains(this->bads[k]))
            res.bads << this->bads[k];
    res.nfree = this->nfree;

    return res;
}

//=============================================================================================================

FiffCov FiffCov::prepare_noise_cov(const FiffInfo &p_Info, const QStringList &p_ChNames) const
{
    FiffCov p_NoiseCov(*this);

    VectorXi C_ch_idx = VectorXi::Zero(p_NoiseCov.names.size());
    qint32 count = 0;
    for(qint32 i = 0; i < p_ChNames.size(); ++i)
    {
        qint32 idx = p_NoiseCov.names.indexOf(p_ChNames[i]);
        if(idx > -1)
        {
            C_ch_idx[count] = idx;
            ++count;
        }
    }
    C_ch_idx.conservativeResize(count);

    MatrixXd C(count, count);

    if(!p_NoiseCov.diag)
        for(qint32 i = 0; i < count; ++i)
            for(qint32 j = 0; j < count; ++j)
                C(i,j) = p_NoiseCov.data(C_ch_idx(i), C_ch_idx(j));
    else
    {
        qWarning("Warning in FiffCov::prepare_noise_cov: This has to be debugged - not done before!");
        C = MatrixXd::Zero(count, count);
        for(qint32 i = 0; i < count; ++i)
            C.diagonal()[i] = p_NoiseCov.data(C_ch_idx(i),0);
    }

    MatrixXd proj;
    qint32 ncomp = p_Info.make_projector(proj, p_ChNames);

    //Create the projection operator
    if (ncomp > 0 && proj.rows() == count)
    {
        printf("Created an SSP operator (subspace dimension = %d)\n", ncomp);
        C = proj * (C * proj.transpose());
    } else {
        qWarning("Warning in FiffCov::prepare_noise_cov: No projections applied since no projectors specified or projector dimensions do not match!");
    }

    RowVectorXi pick_meg = p_Info.pick_types(true, false, false, defaultQStringList, p_Info.bads);
    RowVectorXi pick_eeg = p_Info.pick_types(false, true, false, defaultQStringList, p_Info.bads);

    QStringList meg_names, eeg_names;

    for(qint32 i = 0; i < pick_meg.size(); ++i)
        meg_names << p_Info.chs[pick_meg[i]].ch_name;
    VectorXi C_meg_idx = VectorXi::Zero(p_NoiseCov.names.size());
    count = 0;
    for(qint32 k = 0; k < C.rows(); ++k)
    {
        if(meg_names.indexOf(p_ChNames[k]) > -1)
        {
            C_meg_idx[count] = k;
            ++count;
        }
    }
    if(count > 0)
        C_meg_idx.conservativeResize(count);
    else
        C_meg_idx = VectorXi();

    //
    for(qint32 i = 0; i < pick_eeg.size(); ++i)
        eeg_names << p_Info.chs[pick_eeg(0,i)].ch_name;
    VectorXi C_eeg_idx = VectorXi::Zero(p_NoiseCov.names.size());
    count = 0;
    for(qint32 k = 0; k < C.rows(); ++k)
    {
        if(eeg_names.indexOf(p_ChNames[k]) > -1)
        {
            C_eeg_idx[count] = k;
            ++count;
        }
    }

    if(count > 0)
        C_eeg_idx.conservativeResize(count);
    else
        C_eeg_idx = VectorXi();

    bool has_meg = C_meg_idx.size() > 0;
    bool has_eeg = C_eeg_idx.size() > 0;

    MatrixXd C_meg, C_eeg;
    VectorXd C_meg_eig, C_eeg_eig;
    MatrixXd C_meg_eigvec, C_eeg_eigvec;
    if (has_meg)
    {
        count = C_meg_idx.rows();
        C_meg = MatrixXd(count,count);
        for(qint32 i = 0; i < count; ++i)
            for(qint32 j = 0; j < count; ++j)
                C_meg(i,j) = C(C_meg_idx(i), C_meg_idx(j));
        MNEMath::get_whitener(C_meg, false, QString("MEG"), C_meg_eig, C_meg_eigvec);
    }

    if (has_eeg)
    {
        count = C_eeg_idx.rows();
        C_eeg = MatrixXd(count,count);
        for(qint32 i = 0; i < count; ++i)
            for(qint32 j = 0; j < count; ++j)
                C_eeg(i,j) = C(C_eeg_idx(i), C_eeg_idx(j));
        MNEMath::get_whitener(C_eeg, false, QString("EEG"), C_eeg_eig, C_eeg_eigvec);
    }

    qint32 n_chan = p_ChNames.size();
    p_NoiseCov.eigvec = MatrixXd::Zero(n_chan, n_chan);
    p_NoiseCov.eig = VectorXd::Zero(n_chan);

    if(has_meg)
    {
        for(qint32 i = 0; i < C_meg_idx.rows(); ++i)
            for(qint32 j = 0; j < C_meg_idx.rows(); ++j)
                p_NoiseCov.eigvec(C_meg_idx[i], C_meg_idx[j]) = C_meg_eigvec(i, j);
        for(qint32 i = 0; i < C_meg_idx.rows(); ++i)
            p_NoiseCov.eig(C_meg_idx[i]) = C_meg_eig[i];
    }
    if(has_eeg)
    {
        for(qint32 i = 0; i < C_eeg_idx.rows(); ++i)
            for(qint32 j = 0; j < C_eeg_idx.rows(); ++j)
                p_NoiseCov.eigvec(C_eeg_idx[i], C_eeg_idx[j]) = C_eeg_eigvec(i, j);
        for(qint32 i = 0; i < C_eeg_idx.rows(); ++i)
            p_NoiseCov.eig(C_eeg_idx[i]) = C_eeg_eig[i];
    }

    if (C_meg_idx.size() + C_eeg_idx.size() != n_chan)
    {
        printf("Error in FiffCov::prepare_noise_cov: channel sizes do no match!\n");//ToDo Throw here
        return FiffCov();
    }

    p_NoiseCov.data = C;
    p_NoiseCov.dim = p_ChNames.size();
    p_NoiseCov.diag = false;
    p_NoiseCov.names = p_ChNames;

    return p_NoiseCov;
}

//=============================================================================================================

FiffCov FiffCov::regularize(const FiffInfo& p_info, double p_fRegMag, double p_fRegGrad, double p_fRegEeg, bool p_bProj, QStringList p_exclude) const
{
    FiffCov cov(*this);

    if(p_exclude.size() == 0)
    {
        p_exclude = p_info.bads;
        for(qint32 i = 0; i < cov.bads.size(); ++i)
            if(!p_exclude.contains(cov.bads[i]))
                p_exclude << cov.bads[i];
    }

    //Allways exclude all STI channels from covariance computation
    int iNoStimCh = 0;

    for(int i=0; i<p_info.chs.size(); i++) {
        if(p_info.chs[i].kind == FIFFV_STIM_CH) {
            p_exclude << p_info.chs[i].ch_name;
            iNoStimCh++;
        }
    }

    RowVectorXi sel_eeg = p_info.pick_types(false, true, false, defaultQStringList, p_exclude);
    RowVectorXi sel_mag = p_info.pick_types(QString("mag"), false, false, defaultQStringList, p_exclude);
    RowVectorXi sel_grad = p_info.pick_types(QString("grad"), false, false, defaultQStringList, p_exclude);

    QStringList info_ch_names = p_info.ch_names;
    QStringList ch_names_eeg, ch_names_mag, ch_names_grad;
    for(qint32 i = 0; i < sel_eeg.size(); ++i)
        ch_names_eeg << info_ch_names[sel_eeg(i)];
    for(qint32 i = 0; i < sel_mag.size(); ++i)
        ch_names_mag << info_ch_names[sel_mag(i)];
    for(qint32 i = 0; i < sel_grad.size(); ++i)
        ch_names_grad << info_ch_names[sel_grad(i)];

    // This actually removes bad channels from the cov, which is not backward
    // compatible, so let's leave all channels in
    FiffCov cov_good = cov.pick_channels(info_ch_names, p_exclude);
    QStringList ch_names = cov_good.names;

    std::vector<qint32> idx_eeg, idx_mag, idx_grad;
    for(qint32 i = 0; i < ch_names.size(); ++i)
    {
        if(ch_names_eeg.contains(ch_names[i]))
            idx_eeg.push_back(i);
        else if(ch_names_mag.contains(ch_names[i]))
            idx_mag.push_back(i);
        else if(ch_names_grad.contains(ch_names[i]))
            idx_grad.push_back(i);
    }

    MatrixXd C(cov_good.data);

    //Subtract number of found stim channels because they are still in C but not the idx_eeg, idx_mag or idx_grad
    if((unsigned) C.rows() - iNoStimCh != idx_eeg.size() + idx_mag.size() + idx_grad.size()) {
        printf("Error in FiffCov::regularize: Channel dimensions do not fit.\n");//ToDo Throw
    }

    QList<FiffProj> t_listProjs;
    if(p_bProj)
    {
        t_listProjs = p_info.projs + cov_good.projs;
        FiffProj::activate_projs(t_listProjs);
    }

    //Build regularization MAP
    QMap<QString, QPair<double, std::vector<qint32> > > regData;
    regData.insert("EEG", QPair<double, std::vector<qint32> >(p_fRegEeg, idx_eeg));
    regData.insert("MAG", QPair<double, std::vector<qint32> >(p_fRegMag, idx_mag));
    regData.insert("GRAD", QPair<double, std::vector<qint32> >(p_fRegGrad, idx_grad));

    //
    //Regularize
    //
    QMap<QString, QPair<double, std::vector<qint32> > >::Iterator it;
    for(it = regData.begin(); it != regData.end(); ++it)
    {
        QString desc(it.key());
        double reg = it.value().first;
        std::vector<qint32> idx = it.value().second;

        if(idx.size() == 0 || reg == 0.0)
            printf("\tNothing to regularize within %s data.\n", desc.toUtf8().constData());
        else
        {
            printf("\tRegularize %s: %f\n", desc.toUtf8().constData(), reg);
            MatrixXd this_C(idx.size(), idx.size());
            for(quint32 i = 0; i < idx.size(); ++i)
                for(quint32 j = 0; j < idx.size(); ++j)
                    this_C(i,j) = cov_good.data(idx[i], idx[j]);

            MatrixXd U;
            qint32 ncomp;
            if(p_bProj)
            {
                QStringList this_ch_names;
                for(quint32 k = 0; k < idx.size(); ++k)
                    this_ch_names << ch_names[idx[k]];

                MatrixXd P;
                ncomp = FiffProj::make_projector(t_listProjs, this_ch_names, P); //ToDo: Synchronize with mne-python and debug

                JacobiSVD<MatrixXd> svd(P, ComputeFullU);
                //Sort singular values and singular vectors
                VectorXd t_s = svd.singularValues();
                MatrixXd t_U = svd.matrixU();
                MNEMath::sort<double>(t_s, t_U);

                U = t_U.block(0,0, t_U.rows(), t_U.cols()-ncomp);

                if (ncomp > 0)
                {
                    printf("\tCreated an SSP operator for %s (dimension = %d).\n", desc.toUtf8().constData(), ncomp);
                    this_C = U.transpose() * (this_C * U);
                }
            }

            double sigma = this_C.diagonal().mean();
            this_C.diagonal() = this_C.diagonal().array() + reg * sigma;  // modify diag inplace
            if(p_bProj && ncomp > 0)
                this_C = U * (this_C * U.transpose());

            for(qint32 i = 0; i < this_C.rows(); ++i)
                for(qint32 j = 0; j < this_C.cols(); ++j)
                    C(idx[i],idx[j]) = this_C(i,j);
        }
    }

    // Put data back in correct locations
    RowVectorXi idx = FiffInfo::pick_channels(cov.names, info_ch_names, p_exclude);
    for(qint32 i = 0; i < idx.size(); ++i)
        for(qint32 j = 0; j < idx.size(); ++j)
            cov.data(idx[i], idx[j]) = C(i, j);

    return cov;
}

//=============================================================================================================

FiffCov& FiffCov::operator= (const FiffCov &rhs)
{
    if (this != &rhs) // protect against invalid self-assignment
    {
        kind = rhs.kind;
        diag = rhs.diag;
        dim = rhs.dim;
        names = rhs.names;
        data = rhs.data;
        projs = rhs.projs;
        bads = rhs.bads;
        nfree = rhs.nfree;
        eig = rhs.eig;
        eigvec = rhs.eigvec;
    }
    // to support chained assignment operators (a=b=c), always return *this
    return *this;
}


/*
 @file
 @author Toshihide Hara, Tokyo university of science

 @brief MTRAP: (Pairwise) Sequence alignment algorithm by a new measure based on transition probability between two consecutive pairs of residues

 === HOW TO USE ===
 [1] Sequence Alignment
 mtrap [option] inName outName
 [2] Output Binary format Transition-quantity Matrix
 mtrap -tm inName -outbintm outName
 or
 bash> for i in `ls *.tq` ; do mtrap -tm $i -outbintm ${i%.*}.btq  ; done

=== TODO ===
20120316
= general =
・new が失敗したときNULLがかえってくるとはかぎらない
・staticは，初期化においてマルチスレッド+VCで問題を起こす
20120311
・OPEN MPをOnにすると，Profile作成は高速化するがIterativeRefinementsはたまに遅くなる．原因不明．HyperThreadやキャッシュが原因？
= genetic distance & guide tree =
・案内木構築法としてNJ,WSPS等を実装
・案内木のための遺伝距離行列としてEER以外にもp-distance(clustal互換)等を実装すべき

 $Revision: $
*/

#include <iostream>
#include <iomanip>
#include <cmath>
#include <cassert>
#include <cctype>
#include <map>
#include <string>
#include <fstream>
#include <cmath>
#include <ctime>
#include <Eigen/Dense>
#include <Eigen/StdVector>
#include <shand/format.hpp>
#include "distance.h"
#include "scorematrix.h"
#include "tqmatrix.h"
#include "mtrap.h"
#include "pairwise.h"
#include "probalign.h"
#include "utility.h"
#include "entropy.h"
#include "primarylibrary.h"
#include "search.h"
#include "sequences.h"
#include "globaloption.h"
#include "tree.h"
#include "estimatefamily.h"
#include "config.h"

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

using namespace std;

void help_line(const string& option, const string& value, const string& explain){
    cerr << left << setw(10) << option;
    if(value != "")
    {
        cerr << right << setw(17) << ("[" + value + "]  ");
    }else{
        cerr << right << setw(17) << " ";
    }
    cerr << explain;
    cerr << endl;
}

void printHelp(){
    cerr << "==================" << endl;
    cerr << "  MTRAP Ver. " << string(VERSION);
#ifdef _OPENMP
    cerr << " [USE OPENMP]";
#endif
#ifdef USE_PROBALIGN
    cerr << " [USE PROBALIGN]";
#endif
#ifdef DEBUG
    cerr << " [DEBUG VERSION]";
#endif
    cerr << endl;

    cerr << "=== HOW TO USE ===" << endl;
    cerr << "mtrap [OPTIONS] inputfile outputfile" << endl;
    cerr << endl;
    cerr << "===   OPTIONS  === " << endl;
    cerr << "= GENERAL =" << endl;
    help_line("-I","...","add search path e.g. -I ~/mymatrix");
    help_line("-filelist","...","");
    help_line("-i","...","input file (FASTA format)");
    help_line("-o","...","output file");
    help_line("-nosort","","the results are not sorted");
    help_line("-noestimation","","do not estimate a family");
    help_line("-cds","GENETICCODES (comma separate)","\"protein\" coding DNA sequence");
    cerr << "             GENETICCODE 0: The Standard Code (transl_table=1)" << endl;
    cerr << "                         1: The Vertebrate Mitochondrial Code (transl_table=2)" << endl;
    cerr << "                         2: The Invertebrate Mitochondrial Code (transl_table=5)" << endl;
    cerr << "                         3: The Bacterial, Archaeal and Plant Plastid Code (transl_table=11)" << endl;
    cerr << endl;
    cerr << "= GAP =" << endl;
    help_line("-go",DS(globalOption.gap_open),"gap open [score]");
    help_line("-ge",DS(globalOption.gap_ext),"gap extension [score]");
    cerr << endl;
    cerr << "= MULTIPLE ALIGNMENT =" << endl;
    help_line("-itr",IS(globalOption.numIterativeRefReps),"set the number of iterations");
    help_line("-c",IS(globalOption.numConsistencyReps),"set the number of consistency transformations");
    cerr << endl;
    cerr << "= SUBSTITUTION MATRIX =" << endl;
    help_line("-m",globalOption.smatrixName,"substitution matrix e.g. CGONNET250, EBLOSUM62");
    help_line("-rm",globalOption.rmatrixName,"ramachandran matrix e.g. MRAMA1");
    help_line("-averageExtendedAA","","average extended amino acids (B,Z and X) substitutions");
    cerr << endl;
    cerr << "= TRANSITION QUANTITY MATRIX =" << endl;
    help_line("-tm",globalOption.tmatrixName,"transition-quantity matrix");
    help_line("-e",DS(ScoreMatrix::epsilon),"epsilon, weight for transition quantity");
    cerr << endl;
    cerr << "= RAMACHANDRAN QUANTITY =" << endl;
    help_line("-gamma",DS(ScoreMatrix::gamma),"gamma, weight for ramachandran matrix");
    cerr << endl;
    cerr << "= PARTITION FUNCTION POSTERIOR PROBABILITY =" << endl;
    help_line("-pf",DS(1 - globalOption.mtrap_degree),"the degree of partition function");
    help_line("-pm",globalOption.pmatrixName,"substitution matrix used for partition function");
    help_line("-pgo",DS(globalOption.gap_open_part),"gap open [score] for partition function");
    help_line("-pge",DS(globalOption.gap_ext_part),"gap extension [score] for partition function");
    help_line("-beta",DS(globalOption.beta),"beta, weight for partition function");
#ifndef USE_PROBALIGN
    help_line("-beta2",DS(globalOption.betaTQ),"beta2, weight for a part of TQ of partition function");
#endif
    cerr << endl;
    cerr << "= GENERATING MODE =" << endl;
    help_line("-primarylibrary","ID","output T-Coffee primary library, ID: sequence identity (same as T-Coffee)");
    cerr << endl;
    help_line("-count","MODES FILELIST WEIGHTLIST","count the frequency of transition");
    cerr << "           MODES 0: normal 1:skip gap-gap site 2:set zero gap site" << endl;
    cerr << "                 3: direct prob 4: without lower case 5: global TQ" << endl;
    cerr << "                 f: symmetrize the frequencies not the tq" << endl;
    cerr << "                 n: do not round the substitution matrix" << endl;
    help_line("-outbintm","FILENAME","output binary format transition-quantity matrix");
    cerr << endl;
//    cerr << "= DATABASE SEARCH MODE (UNDER CONSTRUCTION) =" << endl;
//    cerr << "-search query database : query and database must be multiple fasta format. The option -o should be used together." << endl;
//    cerr << "-search_allpair database : database must be ASTRAL SCOP fasta format. The option -o should be used together." << endl;
//    cerr << "-searchlimit query database" << endl;
//    cerr << endl;
//  cerr << "= DEBUG =" << endl;
//	cerr << "-pass output pass matrix filename" << endl;
//	cerr << "-oldeer use traditional entropy evolution rate" << endl;
//	cerr << "-gmode value[0] 0: A,B 1: width=1(B=0)  2: mirror only(A=-1,B=0)" << endl;
//	cerr << "!!! Set -gmode before -m -tm !!!" << endl;
//	cerr << "-withoutT: Gap cost without transition value" << endl;
}

//! load the list for batch style aligment.
/*!

    \param fileName comma separated style, i.e. input filename, output filename[ret].
*/
void loadFilelist(const string& fileName)
{
    ifstream ifs(fileName.c_str());
    string buff;

    while(ifs && getline(ifs, buff))
    {
        vector<string> vs;
        mySplit(buff, ", ", vs);
        if(vs.size() == 0)break;
        globalOption.inputFilenameList.push_back(vs[0]);
        globalOption.outputFilenameList.push_back(vs[1]);
    }

    ifs.close();
}

void setOptionDefaultValue()
{
    // set primary default value
    globalOption.tmatrixName = "SABmark1.63_sup_weighted.btq";
    globalOption.smatrixName = "VTML200I"; // "CGONNET250";
    globalOption.pmatrixName = "VTML200I"; // CGONNET160";
    globalOption.rmatrixName = "MRAMA1";
    globalOption.fmatrixDir = "fsmatrix-sabmark1.65";

    globalOption.csvSCOP_SAB_ID_Name = "scop1.75-id-with-sab1.65-group.csv";
    globalOption.csvSAB_seq_group_Name = "fsmatrix-sabmark1.65/summary_sequence.csv";

    globalOption.familyEstimation = false;
    globalOption.familyDatabaseName = "fsmatrix-sabmark1.65/SABmark1.65.fasta";

    globalOption.CDSMode = false;
    globalOption.CDS_type.clear();

    globalOption.gap_open = -11;
    globalOption.gap_ext = -0.3;
    globalOption.gap_open_part = -22;
    globalOption.gap_ext_part = -1;
    globalOption.mtrap_degree = 0.3;
    globalOption.beta = 0.2;
    globalOption.betaTQ = 1.5;
    ScoreMatrix::gapWithoutTscore = false;
    ScoreMatrix::epsilon = 0.775;
    ScoreMatrix::gamma = 0.0;
    globalOption.primaryLibraryDistance = Distance::GID;
    globalOption.numConsistencyReps = 2;
    globalOption.numIterativeRefReps = 10;
    globalOption.sortByInputOrder = true;

    // set secondary default value

}

int setOption(int argc, char ** argv, int priority)
{
    vec_string vecARGV;
    for(int i=1; i<argc; i++)vecARGV.push_back(string(argv[i]));

    // First priority options
    if(priority==0)
    {
        // default value
        setOptionDefaultValue();

        for(vector<string>::iterator i=vecARGV.begin(); i!=vecARGV.end(); ++i)
        {
            if(i->size() > 0 && i->at(0) == '-')
            {
                if(*i == "-I"){
                    i++; fileSearchDirList.push_back(*i + "/");
                }
            }
        }
        return 0;
    }else if(priority==1){
        // Second priority options
        for(vector<string>::iterator i=vecARGV.begin(); i!=vecARGV.end(); ++i)
        {
            if(i->at(0) == '-')
            {
                if(*i == "-o"){i++; globalOption.outputFilenameList.push_back(*i);
                }else if(*i == "-i"){i++; globalOption.inputFilenameList.push_back(*i);
                }else if(*i == "-nosort"){ globalOption.sortByInputOrder = false;
                }else if(*i == "-noestimation"){ globalOption.familyEstimation = false;
                }else if(*i == "-cds"){
                    i++;
                    globalOption.CDSMode = true;
                    VS tmpVec;
                    mySplit(*i, ",", tmpVec);
                    globalOption.CDS_type.clear();
                    for(VS::const_iterator itr=tmpVec.begin(); itr!=tmpVec.end(); ++itr) globalOption.CDS_type.push_back(SI(*itr));
                }else if(*i == "-pass"){i++; globalOption.outputFilename_pass = *i;
                }else if(*i == "-go"){i++; globalOption.gap_open = string2binary<double>(*i);
                }else if(*i == "-ge"){i++; globalOption.gap_ext = string2binary<double>(*i);
                }else if(*i == "-pgo"){i++; globalOption.gap_open_part = string2binary<double>(*i);
                }else if(*i == "-pge"){i++; globalOption.gap_ext_part = string2binary<double>(*i);
                }else if(*i == "-itr"){i++; globalOption.numIterativeRefReps = string2binary<int>(*i);
                }else if(*i == "-c"){i++; globalOption.numConsistencyReps = string2binary<int>(*i);
                }else if(*i == "-gmode"){i++; ScoreMatrix::TRANS_MODE = string2binary<int>(*i);
                }else if(*i == "-oldeer"){globalOption.optOLDEER = true;
                }else if(*i == "-m"){i++; globalOption.smatrixName = *i;
                }else if(*i == "-tm"){i++; globalOption.tmatrixName = *i;
                }else if(*i == "-rm"){i++; globalOption.rmatrixName = *i;
                }else if(*i == "-pm"){i++; globalOption.pmatrixName = *i;
                }else if(*i == "-averageExtendedAA"){ScoreMatrix::averageExtendedAA = true;
                }else if(*i == "-e"){i++; ScoreMatrix::epsilon = string2binary<double>(*i);
                }else if(*i == "-beta"){i++; globalOption.beta = string2binary<double>(*i);
                }else if(*i == "-beta2"){i++; globalOption.betaTQ = string2binary<double>(*i);
                }else if(*i == "-gamma"){i++; ScoreMatrix::gamma = string2binary<double>(*i);
                }else if(*i == "-pf"){i++; globalOption.mtrap_degree = 1.0 - string2binary<double>(*i);
                }else if(*i == "-withoutT"){ScoreMatrix::gapWithoutTscore = true;
                }else if(*i == "-filelist"){i++; loadFilelist(*i);
                }else if(*i == "-outbintm"){i++; globalOption.outputFilenameList.push_back(*i); globalOption.binoutMode = true;
                }else if(*i == "-search"){i++; globalOption.inputFilenameList.push_back(*i); i++; globalOption.inputFilenameList.push_back(*i); globalOption.searchMode = true;
                }else if(*i == "-search_allpair"){i++; globalOption.inputFilenameList.push_back(*i); globalOption.search_allPairMode = true;
                }else if(*i == "-searchlimit"){i++; globalOption.limitQuery = string2binary<int>(*i); i++; globalOption.limitDatabase = string2binary<int>(*i);
                }else if(*i == "-primarylibrary"){
                    i++; globalOption.primaryLibraryMode = true;
                    if(*i == "gid"){globalOption.primaryLibraryDistance = Distance::GID;
                    }else if(*i == "lid"){globalOption.primaryLibraryDistance = Distance::LID;
                    }else if(*i == "pdist"){globalOption.primaryLibraryDistance = Distance::PDIST;
                    }else if(*i == "eer"){globalOption.primaryLibraryDistance = Distance::EER;
                    }else if(*i == "eer2"){globalOption.primaryLibraryDistance = Distance::EER2;
                    }else if(*i == "ecd"){globalOption.primaryLibraryDistance = Distance::ECD;
                    }else if(*i == "ecdr"){globalOption.primaryLibraryDistance = Distance::ECDR;
                    }else{
                        putError(*i + " mode is not implemented", false);
                        return 1;
                    }
                }else if(*i == "-count"){
                    i++; if(i==vecARGV.end()) { putError("command line may have some mistakes.", true); }
                    for(int j=0; j<(int)i->size(); j++)
                    {
                        if(i->at(j) == '0'){
                            ScoreMatrix::countModeSkipGapGapSite = false;
                            ScoreMatrix::countModeSetZeroGapSite = false;
                            ScoreMatrix::countModeDirectProb = false;
                        }else if(i->at(j) == '1'){
                            ScoreMatrix::countModeSkipGapGapSite = true;
                        }else if(i->at(j) == '2'){
                            ScoreMatrix::countModeSetZeroGapSite = true;
                        }else if(i->at(j) == '3'){
                            ScoreMatrix::countModeDirectProb = true;
                        }else if(i->at(j) == '4'){
                            ScoreMatrix::countModeWithoutLowerCase = true;
                        }else if(i->at(j) == '5'){
                            ScoreMatrix::countModeGlobalTQ = true;
                        }else if(i->at(j) == 'f'){
                            ScoreMatrix::countModeSymmetrizeFreq = true;
                        }else if(i->at(j) == 'n'){
                            ScoreMatrix::countModeNoRound = true;
                        }
                    }
                    // input: database
                    i++; if(i==vecARGV.end()) { putError("command line may have some mistakes.", true); }
                    globalOption.inputFilenameList.push_back(*i);
                    // input: sequence weight
                    i++; if(i==vecARGV.end()) { putError("command line may have some mistakes.", true); }
                    globalOption.inputFilenameList.push_back(*i);
                    // output
                    i++; if(i==vecARGV.end()) { putError("command line may have some mistakes.", true); }
                    globalOption.outputFilenameList.push_back(*i);
                    globalOption.countMode = true;
                }else if(*i == "-I"){
                    ++i;
                    // reservation for first priority.
                }else{
                    return 1;
                }
            }else{
                if(globalOption.inputFilenameList.size() == 0){
                    globalOption.inputFilenameList.push_back(*i);
                }else if(globalOption.outputFilenameList.size() == 0){
                    globalOption.outputFilenameList.push_back(*i);
                }else{
                    putError("too many filenames.", false);
                    return 1;
                }
            }
        }
    } // end of priority

    // === error check ===
    if(globalOption.searchMode)
    {
        assert(globalOption.inputFilenameList.size() == 2 && globalOption.outputFilenameList.size() == 1);
    }else if(globalOption.search_allPairMode){
        assert(globalOption.inputFilenameList.size() == 1 && globalOption.outputFilenameList.size() == 1);
    }else if(globalOption.primaryLibraryMode){
        assert(globalOption.inputFilenameList.size() == 1 && globalOption.outputFilenameList.size() == 1);
    }else if(globalOption.countMode){
        assert(globalOption.inputFilenameList.size() == 2 && globalOption.outputFilenameList.size() == 1);
    }else if(globalOption.binoutMode){
        assert(globalOption.inputFilenameList.size() == 0 && globalOption.outputFilenameList.size() == 1);
    }else{
        if(globalOption.inputFilenameList.size() == 0)
        {
            cerr << "no input filename." << endl;
            return 1;
        }
        if(globalOption.outputFilenameList.size() == 0)
        {
            cerr << "no output filename." << endl;
            return 1;
        }
        assert(globalOption.inputFilenameList.size() == globalOption.outputFilenameList.size());
    }

    if(globalOption.familyEstimation)
    {
        // --- load { SCOP sccs: SABmark group } information ---
        globalOption.setMapSccsGroup();
        // --- load { SABmark sequence: SABmark group } information ---
        globalOption.setMapSeqGroup();
    }

    // --- load matrices ---
    setMatrixConfiguration();

    return 0;
}

void setMatrixConfiguration()
{
    static string previous_smatrix = "";
    static string previous_rmatrix = "";
    static string previous_pmatrix = "";
    static string previous_tmatrix = "";

    bool shouldRecalcDist = false;

    // load ramachandran matrix
    if(previous_rmatrix != globalOption.rmatrixName && ScoreMatrix::gamma > 0.0)
    {
        ScoreMatrix::setRamachandran(fixFilePath(globalOption.rmatrixName));
        previous_rmatrix = globalOption.rmatrixName;
    }

    // load substitution matrix used for partition function
    // this matrix should NOT be normalized
    if(previous_pmatrix != globalOption.pmatrixName)
    {
        ScoreMatrix::loadMatrix(fixFilePath(globalOption.pmatrixName), false, ScoreMatrix::pmatrix);
        previous_pmatrix = globalOption.pmatrixName;
        shouldRecalcDist = true;
    }

    // load substitution matrix
    if(previous_smatrix != globalOption.smatrixName)
    {
        ScoreMatrix::loadMatrix(fixFilePath(globalOption.smatrixName), true, ScoreMatrix::amatrix);
        previous_smatrix = globalOption.smatrixName;
        shouldRecalcDist = true;

        // normalize gap scores
        if(ScoreMatrix::normalized)
        {
            globalOption.gap_open_metric = ScoreMatrix::score_to_dist(globalOption.gap_open, 1);
            globalOption.gap_ext_metric = ScoreMatrix::score_to_dist(globalOption.gap_ext, 1);
            shouldRecalcDist = true;
        }
    }

    // load transition quantity matrix
    if(globalOption.tmatrixName != "" && previous_tmatrix != globalOption.tmatrixName)
    {
        ScoreMatrix::loadTMatrix(fixFilePath(globalOption.tmatrixName));
        previous_tmatrix = globalOption.tmatrixName;
        shouldRecalcDist = true;
    }

    if(shouldRecalcDist)
    {
        ScoreMatrix::calcDistTQplusDiffTable(globalOption);
    }

    // --- DEBUG ---
#ifdef DEBUG
    putLog("::setOption():");
    putLog("score metric matrix (amatrix):");
    ScoreMatrix::showScoreMatrix(ScoreMatrix::amatrix);
    putLog("score metric matrix (pmatrix):");
    ScoreMatrix::showScoreMatrix(ScoreMatrix::pmatrix);
    putLog("gap open = " + DS(globalOption.gap_open));
    putLog("gap ext  = " + DS(globalOption.gap_ext));
    putLog("gap open metric = " + DS(globalOption.gap_open_metric));
    putLog("gap ext  metric = " + DS(globalOption.gap_ext_metric));
    putLog("beta = " + DS(globalOption.beta));
    putLog("betaTQ = " + DS(globalOption.betaTQ));
    putLog("exp(beta * gap open) = " + DS(exp(globalOption.beta * globalOption.gap_open)));
    putLog("exp(beta * gap ext)  = " + DS(exp(globalOption.beta * globalOption.gap_ext)));
    putLog("exp(-beta * gap open) = " + DS(exp(-globalOption.beta * globalOption.gap_open_metric)));
    putLog("exp(-beta * gap ext)  = " + DS(exp(-globalOption.beta * globalOption.gap_ext_metric)));
    putLog("TRANS_A=" + DS(ScoreMatrix::TRANS_A));
    putLog("TRANS_B=" + DS(ScoreMatrix::TRANS_B));
#endif
}

/* Profile version of Multiple Sequence Alignment
* outFile == "" の時は標準出力にアライメントを出力する
* profileは２次元配列を１次元で管理したもの．サイズは(seq1len+1)*(seq2len+1)．また，profile[i][0]=profile[0][j]=0である．
* sparseMatrixはprofileの２次元管理版．０に近い成分は切り捨てたもの
* TODO: ペアワイズアライメント時はProfileを生成せずに直接アライメントを出力し高速化をすべき
*/
void runMSA(const string& inFile, const string& outFile)
{
    clock_t t1, t2, t3, t4;
    t1 = clock();

    // --- read the input sequences ---
    const Sequences* pInputSeqs = new Sequences(inFile, true);
    const int numSeqs = pInputSeqs->seqs.size();

    // --- translate protein coding DNA sequence to AAs
    const Sequences* pCDSSeqs = ((globalOption.CDSMode) ? pInputSeqs->genTranslatedSeqs(globalOption.CDS_type) : NULL);
    if(globalOption.CDSMode) putLog("CDS translated.");

    const Sequences& refSeqs = ((pCDSSeqs != NULL) ? *pCDSSeqs : *pInputSeqs);

    // --- generate the sparse profiles and distances ---
    // 右上三角成分のみ利用(numSeqs*(numSeqs-1)/2だけ余分にメモリを確保している)
    //putLog("generate the sparse profiles and distances");
    VVPROFILE distances(numSeqs, VPROFILE(numSeqs, PROFILE_ZERO));
    VVpSM sparseProfiles(numSeqs, VpSM(numSeqs, (SparseMatrix*)NULL));

    // --- estimate family ---
    if(globalOption.familyEstimation)
    {
        string fmatrixName = globalOption.smatrixName;
        VS estimatedGroups;
        const string mostFamousGroup = ScoreMatrix::estimateGroups(estimatedGroups, refSeqs, true);
        if(mostFamousGroup != "")
        {
            fmatrixName = fixFilePath(globalOption.fmatrixDir + "/" + mostFamousGroup + ".mat");
            putLog("family is estimated. use family specific matrix [" + fmatrixName + "]");
        }else{
            putLog("family is not estimated. use default matrix [" + fmatrixName + "]");
        }

        // --- load family specific matrix ---
        const string old_smatrix = globalOption.smatrixName;
        globalOption.smatrixName = fmatrixName;
        setMatrixConfiguration();
        globalOption.smatrixName = old_smatrix;
    }

    // --- output the parameters ---
//    putLog("=== ALIGNMENT PARAMETERS ===");
//    putLog("-gap_open = " + DS(globalOption.gap_open) + "(" + DS(globalOption.gap_open_metric) + ")");
//    putLog("-gap_ext  = " + DS(globalOption.gap_ext) + "(" + DS(globalOption.gap_ext_metric) + ")");
//    putLog("");

    // --- Open MP ---
#ifdef _OPENMP
    globalOption.numProfilePairs = numSeqs * (numSeqs - 1) / 2;
    globalOption.profilePairs = new pair<int,int>[globalOption.numProfilePairs];
    int pairIndex = 0;
    for(int i=0; i<numSeqs-1; i++)
    {
        for(int j=i+1; j<numSeqs; j++)
        {
            globalOption.profilePairs[pairIndex].first = i;
            globalOption.profilePairs[pairIndex].second = j;
            pairIndex++;
        }
    }
#endif
#ifdef _OPENMP
#pragma omp parallel for private(pairIndex) default(shared) schedule(dynamic)
    for(pairIndex=0; pairIndex<globalOption.numProfilePairs; pairIndex++)
    {
        int i = globalOption.profilePairs[pairIndex].first;
        int j = globalOption.profilePairs[pairIndex].second;
#else
    for (int i=0; i<numSeqs-1; i++)
    {
        for (int j=i+1; j<numSeqs; j++)
        {
#endif
            VPROFILE* profile_mtrap(NULL);
            VPROFILE* profile_postprob(NULL);

            string& seq1 = *(refSeqs.seqs[i]);
            string& seq2 = *(refSeqs.seqs[j]);
            string& name1 = *(refSeqs.names[i]);
            string& name2 = *(refSeqs.names[j]);
            // --- compute profile and distance ---
            // MTRAP
            AlignerMTRAP aligner_mtrap(&name1, &seq1, &name2, &seq2, &globalOption);
            PROFILE_TYPE distance_mtrap = aligner_mtrap.calcPathMatrix();
            profile_mtrap = aligner_mtrap.genProfile();

            // MTRAP with partition function
            if(fabs(globalOption.mtrap_degree - 1.0) > 1.0e-10)
            {
#ifdef USE_PROBALIGN
                // Probalign
                Probalign aligner_probalign;
                profile_postprob = aligner_probalign.genProfile(&seq1, &seq2);
#else
                // Transition quantity with partition function posterior probability
                // To use beta2=0.0 is similar to Probalign
                profile_postprob = aligner_mtrap.genPartProfile();
#endif
            }

#ifdef DEBUG
            aligner_mtrap.showPathMatrix();
            string* dAlignStr = aligner_mtrap.genAlignment();
            putLog("alignment: " + *dAlignStr);
            delete dAlignStr;
            putLog("profile (mtrap):");
            showProfile(profile_mtrap, seq1.size(), seq2.size());
            if(fabs(globalOption.mtrap_degree - 1.0) > 1.0e-10)
            {
                putLog("profile (postprob):");
                showProfile(profile_postprob, seq1.size(), seq2.size());
            }
#endif
            // --- combine all profiles and generate the sparse profile ---
            if(fabs(globalOption.mtrap_degree - 1.0) > 1.0e-10)
            {
                VPROFILE::iterator ptr_mtrap = profile_mtrap->begin();
                VPROFILE::iterator ptr_postprob = profile_postprob->begin();
                // TODO: Eigenにマップして合成すれば高速化になるはず
                for(size_t ii=0; ii<seq1.size(); ii++)
                {
                    for(size_t jj=0; jj<seq2.size(); jj++)
                    {
                        *ptr_mtrap = globalOption.mtrap_degree * (*ptr_mtrap) + (1.0 - globalOption.mtrap_degree) * (*ptr_postprob);
                        ptr_mtrap++;
                        ptr_postprob++;
                    }
                }
            }

            // consider the size of prefix '@'
            sparseProfiles[i][j] = new SparseMatrix(seq1.size()-1, seq2.size()-1, *profile_mtrap);
            // --- set the distance ---
            distances[i][j] = distance_mtrap;
            // --- finalize ---
            delete profile_mtrap;
            if(fabs(globalOption.mtrap_degree - 1.0) > 1.0e-10)
            {
                delete profile_postprob;
            }
#ifndef _OPENMP
        }
#endif
    }

    // --- create the guide tree ---
    t2 = clock();
#ifdef DEBUG
    putLog("create the guide tree");
#endif
    tr_node* pGuideTree = genTreeUPGMA(distances);
#ifdef DEBUG
    putLog(pGuideTree->toString());
#endif
    // --- calculate the sequeces weights ---
#ifdef DEBUG
    putLog("calculate the sequences weights");
#endif
    VPROFILE seqsWeights(numSeqs);
    pGuideTree->setClustalWeight(seqsWeights);
#ifdef DEBUG
    for(size_t i=0; i<seqsWeights.size(); i++) putLog("weight of leaf" + binary2string<int>(i) + "=" + binary2string<double>(seqsWeights[i]));
#endif

    // --- perform the consistency transformation ---
#ifdef DEBUG
    putLog("perform the consistency tranformation");
#endif
    for (int r=0; r<globalOption.numConsistencyReps; r++)
    {
#ifdef DEBUG
        putLog("consistency transform " + IS(r+1) + ".");
#endif
        doConsistencyTrans(seqsWeights, &refSeqs, sparseProfiles);
    }

    // --- compute the final multiple alignment ---
    t3 = clock();
    //putLog("compute the final multiple alignment");
    Sequences* finalAlignment = genFinalAlignment(pGuideTree, &refSeqs, sparseProfiles, &seqsWeights, globalOption);

    // --- output ---
    // console
    //finalAlignment->output(std::cout, ((globalOption.CDSMode) ? const_cast<Sequences*>(pInputSeqs) : NULL));
    // file
    if(outFile != "")
    {
        ofstream ofs(outFile.c_str());
        finalAlignment->output(ofs, ((globalOption.CDSMode) ? const_cast<Sequences*>(pInputSeqs) : NULL));
        ofs.close();
    }

    // --- finalize ---
    for (int i=0; i<numSeqs-1; i++)
    {
        for (int j=i+1; j<numSeqs; j++)
        {
            delete sparseProfiles[i][j];
        }
    }

    delete pInputSeqs;
    delete pCDSSeqs;
    delete pGuideTree;

    delete finalAlignment;

#ifdef _OPENMP
    delete [] globalOption.profilePairs;
#endif

    t4 = clock();
    cerr << "generate the profiles  :" << (double)(t2 - t1) / CLOCKS_PER_SEC << endl;
    cerr << "generate the guide tree:" << (double)(t3 - t2) / CLOCKS_PER_SEC << endl;
    cerr << "iterative refinements  :" << (double)(t4 - t3) / CLOCKS_PER_SEC << endl;
    cerr << "all                    :" << (double)(t4 - t1) / CLOCKS_PER_SEC << endl;
}

int main(int argc, char **argv)
{
    // === INITIALIZE ===
#ifdef _OPENMP
    Eigen::initParallel();
    Eigen::setNbThreads(0);
#endif

    initFileSearchDir();
    ScoreMatrix::init();
    setOption(argc, argv, 0);

    // === OPTION ===
    if(setOption(argc, argv, 1) > 0)
    {
        printHelp();
        return 1;
    }

    // === OUTPUT THE INPUT PARAMETERS ===
    putLog("=== SEARCH DIRS ===");
    for(VS::iterator i=fileSearchDirList.begin(); i!=fileSearchDirList.end(); ++i) putLog(*i);
    putLog("");

    if(globalOption.searchMode)
    {
        cout << "query: " << globalOption.inputFilenameList.at(0) << endl;
        cout << "database: " << globalOption.inputFilenameList.at(1) << endl;
        cout << "output: " << globalOption.outputFilenameList.at(0) << endl;
        search(globalOption.inputFilenameList.at(0), globalOption.inputFilenameList.at(1), globalOption.outputFilenameList.at(0));
    }else if(globalOption.search_allPairMode){
        cout << "database: " << globalOption.inputFilenameList.at(0) << endl;
        cout << "output: " << globalOption.outputFilenameList.at(0) << endl;
        search_allPair(globalOption.inputFilenameList.at(0), globalOption.outputFilenameList.at(0));
    }else if(globalOption.primaryLibraryMode){
        cout << "database: " << globalOption.inputFilenameList.at(0) << endl;
        cout << "output: " << globalOption.outputFilenameList.at(0) << endl;
        makePrimaryLibrary(globalOption.inputFilenameList.at(0), globalOption.outputFilenameList.at(0));
    }else if(globalOption.countMode){
        cout << "database: " << globalOption.inputFilenameList.at(0) << endl;
        cout << "sequence weight: " << globalOption.inputFilenameList.at(1) << endl;
        cout << "output: " << globalOption.outputFilenameList.at(0) << endl;
        ScoreMatrix::generateMatrix(globalOption.inputFilenameList.at(0), globalOption.inputFilenameList.at(1), globalOption.outputFilenameList.at(0));
    }else if(globalOption.binoutMode){
        cout << "generate binary format Transition-quantity matrix." << endl;
        ScoreMatrix::outputBinaryTQ( globalOption.outputFilenameList.at(0) );
    }else{
        for(int i=0; i<(int)globalOption.inputFilenameList.size(); i++)
        {
#ifdef _OPENMP
            //set OpenMP to use dynamic number of threads which is equal to the number of processor cores on the host
            //omp_set_num_threads(2);
            cerr << "OpenMP : Enabled (Max # of threads = " << omp_get_max_threads() << ")" << endl;
#endif
            runMSA(globalOption.inputFilenameList.at(i), globalOption.outputFilenameList.at(i));
        }
    }

    return 0;
}