/** Copyright (c) 2006, Wieger Hofstra and Marnix Kok
* All rights reserved.

* Redistribution and use in source and binary forms, with or without 
* modification, are permitted provided that the following conditions
* are met:

* 1. Redistributions of source code must retain the above copyright 
*    notice, this list of conditions and the following disclaimer.
* 2. 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.
* 3. Neither the name of the <ORGANIZATION> 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.
*
*
*	This is example code how to implement LINCS in C++.
*	LINCS is a LINear Constraint Solver for molecular simulations with bond constraints.
*
*	LINCS is described by: Hess B, Bekker H, Berendsen HJC and Fraaije JGEM (1997)
*	The original paper can be found here: http://www.cs.rug.nl/~bekker/publications/lincs.pdf
*
*	This example is created by Marnix Kok and Wieger Hofstra.
*
*	Note:
*	The weight of an atom is not implemented in this example.
*/


#include "lincs.h"
#include <vector>
#include <cmath>
#include <cstdlib>
#include <cassert>
#include <iostream>

#include "recorderOptions.h"
#include "ropeEnviroment.h"

using namespace std;


LINCSEnvironment::LINCSEnvironment(CLOptions options)
        : d_truncateAfter(options.powerOrder), d_activeParticles(0), d_verbose(options.verbose)
{}

/**
 *	Free allocated resources
 */
LINCSEnvironment::~LINCSEnvironment()
{
    // release 2nd dimension allocated mem
    for (size_t i = 0; i < 2; i++)
        delete [] d_rhs[i];

    for (size_t i = 0; i < d_constraints.size(); i++)
        delete [] d_nccMatrix[i];

    for (size_t i = 0; i < d_constraints.size(); i++)
        delete [] d_constraints[i];

    for (size_t i =0 ; i < d_particles.size(); i++)
        delete [] d_particles[i];

    // release 1st dimension allocated mem
    delete [] d_solution;
    delete [] d_rhs;
    delete [] d_nccMatrix;

}

/**
 * Initialize arrays that are important to process
 */
void LINCSEnvironment::initializeArrays(size_t const maxConstraints)
{
    d_maxConstraints = maxConstraints;

    d_solution = new double[d_constraints.size()];
    for (size_t j = 0; j < d_constraints.size(); ++j)
        d_solution[j] = 0;

    // initialize rhs array
    d_rhs = new double*[2];
    for (size_t i = 0; i < 2; i++)
    {
        d_rhs[i] = new double[d_constraints.size()];

        for (size_t j = 0; j < d_constraints.size(); ++j)
            d_rhs[i][j] = 0;
    }

    d_nccMatrix = new double*[d_constraints.size()];
    for (size_t i = 0; i < d_constraints.size(); ++i)
    {
        d_nccMatrix[i] = new double[maxConstraints];

        for (size_t j = 0; j < maxConstraints; ++j)
            d_nccMatrix[i][j] = 0;
    }

}

void LINCSEnvironment::constrainParticles()
{

    calcConstraintDirection();

    calcNCCMatrix();
    solve();

    correctRotationalLengthening();
    solve();
}



// -----------------------------------------------------------
//		Private interface imlpementation
// -----------------------------------------------------------

/**
 *	Determine the direction of the constraint
 */
void LINCSEnvironment::calcConstraintDirection()
{

    // iterate over all constraints
    for (size_t i = 0; i < nConstraints(); i++)
    {
        Constraint_t *c = d_constraints[i];

        // determine position difference
        c->direction.x = c->from->pX - c->to->pX;
        c->direction.y = c->from->pY - c->to->pY;
        c->direction.z = c->from->pZ - c->to->pZ;


        //A * B = Ax*Bx + Ay*By + Az*Bz

        if (d_verbose && i == 5)
            cout << " X dir = " << c->direction.x << " from " << c->from->pX << " to " << c->to->pX << "  |  "
            << " Y dir = " << c->direction.y << " from " << c->from->pY << " to " << c->to->pY << "  |  "
            << " Z dir = " << c->direction.z << " from " << c->from->pZ << " to " << c->to->pZ << endl;

        // determine eucledian distance
        double cLength = sqrt(sqr(c->direction.x) +
                              sqr(c->direction.y) +
                              sqr(c->direction.z));


        // normalize direction vector
        c->direction.x /= cLength;
        c->direction.y /= cLength;
        c->direction.z /= cLength;

    }
}


/**
 *	Calculate the normalized constraint coupling matrix
 */
void LINCSEnvironment::calcNCCMatrix()
{

    for (size_t i = 0; i < nConstraints(); i++)
    {
        Constraint_t *c = d_constraints[i];
        vector< Constraint_t* > connected = connectedConstraints(c);
        for (size_t n = 0; n < connected.size(); n++)
        {
            Constraint_t *k = connected[n];
            d_nccMatrix[i][n] = coefficient(i, n) *
                                ((c->direction.x * k->direction.x) +
                                 (c->direction.y * k->direction.y) +
                                 (c->direction.z * k->direction.z));
        }

        d_rhs[0][i] = c->sdiag * (c->direction.x * (c->from->x - c->to->x) +
                                  c->direction.y * (c->from->y - c->to->y) +
                                  c->direction.z * (c->from->z - c->to->z) - c->length);
        d_solution[i] = d_rhs[0][i];
    }
}

/**
 *	Solve the constraints put on the particles and reset them
 *  to their proper positions.
 */
void LINCSEnvironment::solve()
{
    size_t w = 1;

    // repeat d_truncateAfter times, equals amount of other particles
    // are influenced.
    for (size_t nrec = 0; nrec < d_truncateAfter; nrec++)
    {
        // iterate through all constraints
        for (size_t i = 0; i < nConstraints(); i++)
        {
            // retrieve connected constraints
            vector< Constraint_t* > con = connectedConstraints(d_constraints[i]);

            d_rhs[w][i] = 0;
            for (size_t n = 0; n < con.size(); n++)
                d_rhs[w][i] += d_nccMatrix[i][n] * d_rhs[(w == 1? 0 : 1)][constraintIndex(con[n])];

            d_solution[i] += d_rhs[w][i];
        }
        w = (w == 1? 0 : 1);
    }
    /** @brief performAdjustments
     *
     */

    // adjust the particle positions
    for (size_t i = 0; i < nConstraints(); i++)
    {
        Constraint_t *c = d_constraints[i];

        if (DEBUG)
            cout << "adjust pos of " << i << " old f: " << c->from->x << " old t: " << c->to->x << endl;

		double resistance = 1;

		if(c->from->y <= 0)
			resistance = GROUNDFRICTION;

        c->from->x -= c->direction.x * c->sdiag * d_solution[i] / resistance;
        c->from->y -= c->direction.y * c->sdiag * d_solution[i] / resistance;
        c->from->z -= c->direction.z * c->sdiag * d_solution[i] / resistance;

		resistance = 1;
		if(c->to->y <= 0)
			resistance = GROUNDFRICTION;

        c->to->x += c->direction.x * c->sdiag * d_solution[i] / resistance;
        c->to->y += c->direction.y * c->sdiag * d_solution[i] / resistance;
        c->to->z += c->direction.z * c->sdiag * d_solution[i] / resistance;

        if (d_verbose && i == 5)
            cout << "adjust pos of " << i << " new f: " << c->from->z << " new t: " << c->to->z << " The adjustment was: "
            << (c->direction.z * c->sdiag * d_solution[i])
            << " (" <<  c->direction.z << " * " <<  c->sdiag << " * " << d_solution[i] << ")"
            << endl;
    }
}


/**
 *	Adjust the values in d_solution to account for rotational lengthening
 */
void LINCSEnvironment::correctRotationalLengthening()
{
    for (size_t i = 0; i < nConstraints(); i++)
    {
        Constraint_t *c = d_constraints[i];
        //double p = (2 * sqr(c->length)) - sqr(length(c->from, c->to));
        double p = (2 * sqr(c->length)) - sqr(c->from->x - c->to->x) -
                   sqr(c->from->y - c->to->y) -
                   sqr(c->from->z - c->to->z);

        // check < 0, else we get lots of nan
        if (p < 0)
        {
            if (d_verbose)
                cout << " p of " << i << " out of bounds: " << p << endl;
            p = 0;
        }

        d_rhs[0][i] = c->sdiag * (c->length - sqrt(p));
        d_solution[i] = d_rhs[0][i];
    }
}


/**
 *	Collect the constraints that border the
 *  constraint specified in `constraint`
 */
vector< Constraint_t * > LINCSEnvironment::connectedConstraints(Constraint_t const *constraint)
{
    vector< Constraint_t * > connected;

    // iterate through all constraints
    for (size_t i = 0; i < nConstraints(); i++)
    {
        // self? continue
        if (d_constraints[i] == constraint)
            continue;

        // if either the start == end or end == start of other
        // constraint, add it to the connected list
        if (constraint->from == d_constraints[i]->to || constraint->to == d_constraints[i]->from)
            connected.push_back(d_constraints[i]);
    }

    return connected;
}

/**
 *	Determine what particle is shared by both constraints, returns 0 if
 *  none are shared
 */
Particle_t *LINCSEnvironment::sharedParticle(Constraint_t const *c1, Constraint_t const *c2)
{
    if (c1->from == c2->to)
        return c1->from;

    if (c1->to == c2->from)
        return c1->to;

    return 0;
}

/**
 *	Return the index in the vector of the constraint given by `c`. Returns
 *  -1 if no such constraints is found. (should not happen)
 */
int LINCSEnvironment::constraintIndex(Constraint_t const *c) const
{
    for (size_t i = 0; i < nConstraints(); i++)
        if (d_constraints[i] == c)
            return i;

    return -1;
}

/**
 *	Supply the coefficient of the constraint
 */
double LINCSEnvironment::coefficient(size_t i, size_t j)
{
    // make sure we do not break boundaries
    assert(i < nConstraints());
    assert(j < d_maxConstraints);

    double sign = 1.0;
    vector<Constraint_t *> connections = connectedConstraints(d_constraints[i]);

    // -- if atom1[i] = atom1[con[i,j]] or atom2[i]=atom2[con[i,j]] sign = -1 else sign = 1
    if (connections[j]->from == d_constraints[i]->from || connections[j]->to == d_constraints[i]->to)
        sign = -1.0;

    //Particle_t *sParticle = sharedParticle(d_constraints[i], connections[j]);
    // cout << "sparticle: " << sParticle->x << ", " << sParticle->y << ", " << sParticle->z << endl;

    // -- where c is the atom coupling constraints i and con[i,j],
    // coef[i, j] = sign * invmass[c] * sDiag[i] * sDiag[con[i,j]]
    //return sign * (1.0 / sParticle->mass) * d_constraints[i]->sdiag * connections[j]->sdiag;
    return sign * d_constraints[i]->sdiag * connections[j]->sdiag;
}

// -----------------------------------------------------------
//		Protected methods implementation
// -----------------------------------------------------------

/**
 *	Old particle position will equal the current position
 */
void LINCSEnvironment::ageParticle(Particle_t *particle)
{

    particle->pX = particle->x;
    particle->pY = particle->y;
    particle->pZ = particle->z;
    particle->pVel = particle->vel;

    if (particle->y > 0)
        d_activeParticles++;
}

/**
 *	Determine eucledian distance between two particles
 */
double LINCSEnvironment::length(Particle_t const *p1, Particle_t const *p2) const
{
    return sqrt( sqr(p1->x - p2->x) +
                 sqr(p1->y - p2->y) +
                 sqr(p1->z - p2->z) );
}