/*******************************************************************************
shapeAlignment.cpp - Shape-it
 
Copyright 2012 by Silicos-it, a division of Imacosi BVBA
 
This file is part of Shape-it.

	Shape-it is free software: you can redistribute it and/or modify
	it under the terms of the GNU Lesser General Public License as published 
	by the Free Software Foundation, either version 3 of the License, or
	(at your option) any later version.

	Shape-it is distributed in the hope that it will be useful,
	but WITHOUT ANY WARRANTY; without even the implied warranty of
	MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
	GNU Lesser General Public License for more details.

	You should have received a copy of the GNU Lesser General Public License
	along with Shape-it.  If not, see <http://www.gnu.org/licenses/>.

Shape-it is linked against OpenBabel version 2.

	OpenBabel is free software; you can redistribute it and/or modify
	it under the terms of the GNU General Public License as published by
	the Free Software Foundation version 2 of the License.

***********************************************************************/



#include <Shape/shapeAlignment.h>



ShapeAlignment::ShapeAlignment(GaussianVolume & gRef, GaussianVolume & gDb)
:  _gRef(&gRef)
    , _gDb(&gDb)
    , _rAtoms(0)
    , _rGauss(0)
    , _dAtoms(0)
    , _dGauss(0)
    , _maxSize(0)
    , _maxIter(50)
    , _matrixMap()
{
    // Loop over the single atom volumes of both molecules and make the combinations
    _rAtoms = gRef.levels[0];
    _dAtoms = gDb.levels[0];
    _rGauss = gRef.gaussians.size();
    _dGauss = gDb.gaussians.size();
    _maxSize = _rGauss * _dGauss + 1;
}



ShapeAlignment::~ShapeAlignment(void)
{
    _gRef = NULL;
    _gDb = NULL;

    // Clear the matrix map
    for (MatIter mi = _matrixMap.begin(); mi != _matrixMap.end(); ++mi) {
	if (mi->second != NULL) {
	    delete mi->second;
	    mi->second = NULL;
	}
    }
}



AlignmentInfo ShapeAlignment::gradientAscent(SiMath::Vector rotor)
{
    // Create a queue to hold the pairs to process
    std::queue < std::pair < unsigned int, unsigned int > >processQueue;

    // Helper variables
    // Overlap matrix is stored as a double error
    double *Aij;
    double Aq[4];

    double Vij(0.0);
    double qAq(0.0);

    std::vector < unsigned int >*d1(NULL);
    std::vector < unsigned int >*d2(NULL);
    std::vector < unsigned int >::iterator it1;

    // Gradient information
    SiMath::Vector overGrad(4, 0.0);

    // Hessian
    SiMath::Matrix overHessian(4, 4, 0.0);

    // Overlap volume
    double atomOverlap(0.0);
    double pharmOverlap(0.0);

    // Solution info
    AlignmentInfo res;

    double oldVolume(0.0);
    unsigned int iterations(0);
    unsigned int mapIndex(0);

//      unsigned int DBSize = _gDb->pharmacophores.size();

    while (iterations < 20) {
	// reset volume
	atomOverlap = 0.0;
	pharmOverlap = 0.0;
	iterations++;

	// reset gradient
	overGrad = 0.0;

	// reset hessian
	overHessian = 0.0;

	double lambda(0.0);

	// iterator over the matrix map
	MatIter matIter;

	// create atom-atom overlaps 
	for (unsigned int i(0); i < _rAtoms; ++i) {
	    for (unsigned int j(0); j < _dAtoms; ++j) {
		mapIndex = (i * _dGauss) + j;

		if ((matIter =
		     _matrixMap.find(mapIndex)) == _matrixMap.end()) {
		    Aij =
			_updateMatrixMap(_gRef->gaussians[i],
					 _gDb->gaussians[j]);

		    // add to map
		    _matrixMap[mapIndex] = Aij;
		} else {
		    Aij = matIter->second;
		}

		// rotor product
		Aq[0] =
		    Aij[0] * rotor[0] + Aij[1] * rotor[1] +
		    Aij[2] * rotor[2] + Aij[3] * rotor[3];
		Aq[1] =
		    Aij[4] * rotor[0] + Aij[5] * rotor[1] +
		    Aij[6] * rotor[2] + Aij[7] * rotor[3];
		Aq[2] =
		    Aij[8] * rotor[0] + Aij[9] * rotor[1] +
		    Aij[10] * rotor[2] + Aij[11] * rotor[3];
		Aq[3] =
		    Aij[12] * rotor[0] + Aij[13] * rotor[1] +
		    Aij[14] * rotor[2] + Aij[15] * rotor[3];

		qAq =
		    rotor[0] * Aq[0] + rotor[1] * Aq[1] +
		    rotor[2] * Aq[2] + rotor[3] * Aq[3];

		// compute overlap volume
		Vij = Aij[16] * exp(-qAq);

		// check if overlap is sufficient enough, should be more than 0.1 for atom - atom overlap
		if (Vij /
		    (_gRef->gaussians[i].volume +
		     _gDb->gaussians[j].volume - Vij) < EPS) {
		    continue;
		}
		// add to overlap volume
		atomOverlap += Vij;

		// update gradient -2Vij (Aijq);
		double v2 = 2.0 * Vij;

		lambda -= v2 * qAq;

		overGrad[0] -= v2 * Aq[0];
		overGrad[1] -= v2 * Aq[1];
		overGrad[2] -= v2 * Aq[2];
		overGrad[3] -= v2 * Aq[3];

		// overHessian += 2*Vij(2*Aijq'qAij-Aij); (only upper triangular part)
		overHessian[0][0] += v2 * (2.0 * Aq[0] * Aq[0] - Aij[0]);
		overHessian[0][1] += v2 * (2.0 * Aq[0] * Aq[1] - Aij[1]);
		overHessian[0][2] += v2 * (2.0 * Aq[0] * Aq[2] - Aij[2]);
		overHessian[0][3] += v2 * (2.0 * Aq[0] * Aq[3] - Aij[3]);
		overHessian[1][1] += v2 * (2.0 * Aq[1] * Aq[1] - Aij[5]);
		overHessian[1][2] += v2 * (2.0 * Aq[1] * Aq[2] - Aij[6]);
		overHessian[1][3] += v2 * (2.0 * Aq[1] * Aq[3] - Aij[7]);
		overHessian[2][2] += v2 * (2.0 * Aq[2] * Aq[2] - Aij[10]);
		overHessian[2][3] += v2 * (2.0 * Aq[2] * Aq[3] - Aij[11]);
		overHessian[3][3] += v2 * (2.0 * Aq[3] * Aq[3] - Aij[15]);

		// loop over child nodes and add to queue
		d1 = _gRef->childOverlaps[i];
		d2 = _gDb->childOverlaps[j];

		// first add (i,child(j))
		if (d2 != NULL) {
		    for (it1 = d2->begin(); it1 != d2->end(); ++it1) {
			processQueue.push(std::make_pair < unsigned int,
					  unsigned int >(i, *it1));
		    }
		}
		// second add (child(i),j)
		if (d1 != NULL) {
		    for (it1 = d1->begin(); it1 != d1->end(); ++it1) {
			processQueue.push(std::make_pair < unsigned int,
					  unsigned int >(*it1, j));
		    }
		}
	    }
	}

	while (!processQueue.empty()) {
	    // get next element from queue
	    std::pair < unsigned int, unsigned int >nextPair =
		processQueue.front();
	    processQueue.pop();

	    unsigned int i = nextPair.first;
	    unsigned int j = nextPair.second;

	    // check cache
	    mapIndex = (i * _dGauss) + j;
	    if ((matIter = _matrixMap.find(mapIndex)) == _matrixMap.end()) {
		Aij =
		    _updateMatrixMap(_gRef->gaussians[i],
				     _gDb->gaussians[j]);
		_matrixMap[mapIndex] = Aij;
	    } else {
		Aij = matIter->second;
	    }

	    // rotor product
	    Aq[0] =
		Aij[0] * rotor[0] + Aij[1] * rotor[1] + Aij[2] * rotor[2] +
		Aij[3] * rotor[3];
	    Aq[1] =
		Aij[4] * rotor[0] + Aij[5] * rotor[1] + Aij[6] * rotor[2] +
		Aij[7] * rotor[3];
	    Aq[2] =
		Aij[8] * rotor[0] + Aij[9] * rotor[1] +
		Aij[10] * rotor[2] + Aij[11] * rotor[3];
	    Aq[3] =
		Aij[12] * rotor[0] + Aij[13] * rotor[1] +
		Aij[14] * rotor[2] + Aij[15] * rotor[3];

	    qAq =
		rotor[0] * Aq[0] + rotor[1] * Aq[1] + rotor[2] * Aq[2] +
		rotor[3] * Aq[3];

	    // compute overlap volume
	    Vij = Aij[16] * exp(-qAq);

	    // check if overlap is sufficient enough
	    if (fabs(Vij) /
		(_gRef->gaussians[i].volume + _gDb->gaussians[j].volume -
		 fabs(Vij)) < EPS) {
		continue;
	    }

	    atomOverlap += Vij;

	    double v2 = 2.0 * Vij;

	    lambda -= v2 * qAq;

	    // update gradient -2Vij (Cij*Aij*q);
	    overGrad[0] -= v2 * Aq[0];
	    overGrad[1] -= v2 * Aq[1];
	    overGrad[2] -= v2 * Aq[2];
	    overGrad[3] -= v2 * Aq[3];

	    // hessian 2*Vij(2*AijqqAij-Aij); (only upper triangular part)
	    overHessian[0][0] += v2 * (2.0 * Aq[0] * Aq[0] - Aij[0]);
	    overHessian[0][1] += v2 * (2.0 * Aq[0] * Aq[1] - Aij[1]);
	    overHessian[0][2] += v2 * (2.0 * Aq[0] * Aq[2] - Aij[2]);
	    overHessian[0][3] += v2 * (2.0 * Aq[0] * Aq[3] - Aij[3]);
	    overHessian[1][1] += v2 * (2.0 * Aq[1] * Aq[1] - Aij[5]);
	    overHessian[1][2] += v2 * (2.0 * Aq[1] * Aq[2] - Aij[6]);
	    overHessian[1][3] += v2 * (2.0 * Aq[1] * Aq[3] - Aij[7]);
	    overHessian[2][2] += v2 * (2.0 * Aq[2] * Aq[2] - Aij[10]);
	    overHessian[2][3] += v2 * (2.0 * Aq[2] * Aq[3] - Aij[11]);
	    overHessian[3][3] += v2 * (2.0 * Aq[3] * Aq[3] - Aij[15]);


	    // loop over child nodes and add to queue
	    d1 = _gRef->childOverlaps[i];
	    d2 = _gDb->childOverlaps[j];
	    if (d1 != NULL
		&& _gRef->gaussians[i].nbr > _gDb->gaussians[j].nbr) {
		for (it1 = d1->begin(); it1 != d1->end(); ++it1) {
		    // add (child(i),j)
		    processQueue.push(std::make_pair < unsigned int,
				      unsigned int >(*it1, j));
		}
	    } else {
		// first add (i,child(j))
		if (d2 != NULL) {
		    for (it1 = d2->begin(); it1 != d2->end(); ++it1) {
			processQueue.push(std::make_pair < unsigned int,
					  unsigned int >(i, *it1));
		    }
		}
		if (d1 != NULL
		    && _gDb->gaussians[j].nbr - _gRef->gaussians[i].nbr <
		    2) {
		    for (it1 = d1->begin(); it1 != d1->end(); ++it1) {
			// add (child(i),j)
			processQueue.push(std::make_pair < unsigned int,
					  unsigned int >(*it1, j));
		    }
		}
	    }
	}


	// check if the new volume is better than the previously found one
	// if not quit the loop
	if (iterations > 6 && atomOverlap < oldVolume + 0.0001) {
	    break;
	}
	// store latest volume found
	oldVolume = atomOverlap;

	// no measurable overlap between two volumes
	if (std::isnan(lambda) || std::isnan(oldVolume) || oldVolume == 0) {
	    break;
	}
	// update solution 
	if (oldVolume > res.overlap) {
	    res.overlap = atomOverlap;
	    res.rotor = rotor;
	    if (res.overlap /
		(_gRef->overlap + _gDb->overlap - res.overlap) > 0.99) {
		break;
	    }
	}
	// update the gradient and hessian 
	overHessian -= lambda;

	// fill lower triangular of the hessian matrix
	overHessian[1][0] = overHessian[0][1];
	overHessian[2][0] = overHessian[0][2];
	overHessian[2][1] = overHessian[1][2];
	overHessian[3][0] = overHessian[0][3];
	overHessian[3][1] = overHessian[1][3];
	overHessian[3][2] = overHessian[2][3];

	// update gradient to make h
	overGrad[0] -= lambda * rotor[0];
	overGrad[1] -= lambda * rotor[1];
	overGrad[2] -= lambda * rotor[2];
	overGrad[3] -= lambda * rotor[3];

	// update gradient based on inverse hessian
	SiMath::Vector tmp = rowProduct(overHessian, overGrad);
	double h =
	    overGrad[0] * tmp[0] + overGrad[1] * tmp[1] +
	    overGrad[2] * tmp[2] + overGrad[3] * tmp[3];

	// small scaling of the gradient
	h = 1 / h * (overGrad[0] * overGrad[0] +
		     overGrad[1] * overGrad[1] +
		     overGrad[2] * overGrad[2] +
		     overGrad[3] * overGrad[3]);
	overGrad *= h;

	// update rotor based on gradient information
	rotor -= overGrad;

	// normalise rotor such that it has unit norm
	double nr =
	    sqrt(rotor[0] * rotor[0] + rotor[1] * rotor[1] +
		 rotor[2] * rotor[2] + rotor[3] * rotor[3]);

	rotor[0] /= nr;
	rotor[1] /= nr;
	rotor[2] /= nr;
	rotor[3] /= nr;

    }				// end of endless while loop
    return res;
}



AlignmentInfo ShapeAlignment::simulatedAnnealing(SiMath::Vector rotor)
{
    // create a queue to hold the pairs to process
    std::queue < std::pair < unsigned int, unsigned int > >processQueue;

    // helper variables
    // overlap matrix is stored as a double error
    double *Aij;
    double Aq[4];

    // map store store already computed matrices
    MatIter matIter;

    double Vij(0.0), qAq(0.0);

    std::vector < unsigned int >*d1(NULL);
    std::vector < unsigned int >*d2(NULL);
    std::vector < unsigned int >::iterator it1;

    // overlap volume
    double atomOverlap(0.0);
    double pharmOverlap(0.0);

    double dTemperature = 1.1;

    // solution info
    AlignmentInfo res;

    SiMath::Vector oldRotor(rotor);
    SiMath::Vector bestRotor(rotor);

    double oldVolume(0.0);
    double bestVolume(0.0);
    unsigned int iterations(0);
    unsigned int sameCount(0);
    unsigned int mapIndex(0);
//      unsigned int DBSize = _gDb->pharmacophores.size();

    while (iterations < _maxIter) {
	// reset volume
	atomOverlap = 0.0;
	pharmOverlap = 0.0;

	++iterations;

	// temperature of the simulated annealing step
	double T = sqrt((1.0 + iterations) / dTemperature);

	// create atom-atom overlaps 
	for (unsigned int i = 0; i < _rAtoms; ++i) {
	    for (unsigned int j = 0; j < _dAtoms; ++j) {
		mapIndex = (i * _dGauss) + j;

		if ((matIter =
		     _matrixMap.find(mapIndex)) == _matrixMap.end()) {
		    Aij =
			_updateMatrixMap(_gRef->gaussians[i],
					 _gDb->gaussians[j]);
		    _matrixMap[mapIndex] = Aij;
		} else {
		    Aij = matIter->second;
		}

		// rotor product
		Aq[0] =
		    Aij[0] * rotor[0] + Aij[1] * rotor[1] +
		    Aij[2] * rotor[2] + Aij[3] * rotor[3];
		Aq[1] =
		    Aij[4] * rotor[0] + Aij[5] * rotor[1] +
		    Aij[6] * rotor[2] + Aij[7] * rotor[3];
		Aq[2] =
		    Aij[8] * rotor[0] + Aij[9] * rotor[1] +
		    Aij[10] * rotor[2] + Aij[11] * rotor[3];
		Aq[3] =
		    Aij[12] * rotor[0] + Aij[13] * rotor[1] +
		    Aij[14] * rotor[2] + Aij[15] * rotor[3];

		qAq =
		    rotor[0] * Aq[0] + rotor[1] * Aq[1] +
		    rotor[2] * Aq[2] + rotor[3] * Aq[3];

		// compute overlap volume
		Vij = Aij[16] * exp(-qAq);

		// check if overlap is sufficient enough, should be more than 0.1 for atom - atom overlap
		if (Vij /
		    (_gRef->gaussians[i].volume +
		     _gDb->gaussians[j].volume - Vij) < EPS) {
		    continue;
		}
		// add to overlap volume
		atomOverlap += Vij;

		// loop over child nodes and add to queue
		d1 = _gRef->childOverlaps[i];
		d2 = _gDb->childOverlaps[j];

		// first add (i,child(j))
		if (d2 != NULL) {
		    for (it1 = d2->begin(); it1 != d2->end(); ++it1) {
			processQueue.push(std::make_pair < unsigned int,
					  unsigned int >(i, *it1));
		    }
		}
		// second add (child(i),j)
		if (d1 != NULL) {
		    for (it1 = d1->begin(); it1 != d1->end(); ++it1) {
			processQueue.push(std::make_pair < unsigned int,
					  unsigned int >(*it1, j));
		    }
		}
	    }
	}

	while (!processQueue.empty()) {
	    // get next element from queue
	    std::pair < unsigned int, unsigned int >nextPair =
		processQueue.front();
	    processQueue.pop();

	    unsigned int i = nextPair.first;
	    unsigned int j = nextPair.second;

	    // check cache
	    mapIndex = (i * _dGauss) + j;
	    if ((matIter = _matrixMap.find(mapIndex)) == _matrixMap.end()) {
		Aij =
		    _updateMatrixMap(_gRef->gaussians[i],
				     _gDb->gaussians[j]);
		_matrixMap[mapIndex] = Aij;
	    } else {
		Aij = matIter->second;
	    }

	    // rotor product
	    Aq[0] =
		Aij[0] * rotor[0] + Aij[1] * rotor[1] + Aij[2] * rotor[2] +
		Aij[3] * rotor[3];
	    Aq[1] =
		Aij[4] * rotor[0] + Aij[5] * rotor[1] + Aij[6] * rotor[2] +
		Aij[7] * rotor[3];
	    Aq[2] =
		Aij[8] * rotor[0] + Aij[9] * rotor[1] +
		Aij[10] * rotor[2] + Aij[11] * rotor[3];
	    Aq[3] =
		Aij[12] * rotor[0] + Aij[13] * rotor[1] +
		Aij[14] * rotor[2] + Aij[15] * rotor[3];

	    qAq =
		rotor[0] * Aq[0] + rotor[1] * Aq[1] + rotor[2] * Aq[2] +
		rotor[3] * Aq[3];

	    // compute overlap volume
	    Vij = Aij[16] * exp(-qAq);

	    // check if overlap is sufficient enough
	    if (fabs(Vij) /
		(_gRef->gaussians[i].volume + _gDb->gaussians[j].volume -
		 fabs(Vij)) < EPS) {
		continue;
	    }
	    // even number of overlap atoms => addition to volume
	    atomOverlap += Vij;

	    // loop over child nodes and add to queue
	    d1 = _gRef->childOverlaps[i];
	    d2 = _gDb->childOverlaps[j];
	    if (d1 != NULL
		&& _gRef->gaussians[i].nbr > _gDb->gaussians[j].nbr) {
		for (it1 = d1->begin(); it1 != d1->end(); ++it1) {
		    // add (child(i),j)
		    processQueue.push(std::make_pair < unsigned int,
				      unsigned int >(*it1, j));
		}
	    } else {
		// first add (i,child(j))
		if (d2 != NULL) {
		    for (it1 = d2->begin(); it1 != d2->end(); ++it1) {
			processQueue.push(std::make_pair < unsigned int,
					  unsigned int >(i, *it1));
		    }
		}
		if (d1 != NULL
		    && _gDb->gaussians[j].nbr - _gRef->gaussians[i].nbr <
		    2) {
		    for (it1 = d1->begin(); it1 != d1->end(); ++it1) {
			// add (child(i),j)
			processQueue.push(std::make_pair < unsigned int,
					  unsigned int >(*it1, j));
		    }
		}
	    }
	}

	// check if the new volume is better than the previously found one
	double overlapVol = atomOverlap;
	if (overlapVol < oldVolume) {
	    double D = exp(-sqrt(oldVolume - overlapVol)) / T;
	    if (SiMath::randD(0, 1) < D) {
		oldRotor = rotor;
		oldVolume = overlapVol;
		sameCount = 0;
	    } else {
		++sameCount;
		if (sameCount == 30) {
		    iterations = _maxIter;
		}
	    }
	} else {
	    // store latest volume found
	    oldVolume = overlapVol;
	    oldRotor = rotor;

	    // update best found so far
	    bestRotor = rotor;
	    bestVolume = overlapVol;
	    sameCount = 0;

	    // check if it is better than the best solution found so far
	    if (overlapVol > res.overlap) {
		res.overlap = atomOverlap;
		res.rotor = rotor;
		if ((res.overlap /
		     (_gRef->overlap + _gDb->overlap - res.overlap)) >
		    0.99) {
		    break;
		}
	    }

	}

	// make random permutation
	// double range = 0.05;
	double range = 0.1 / T;
	rotor[0] = oldRotor[0] + SiMath::randD(-range, range);
	rotor[1] = oldRotor[1] + SiMath::randD(-range, range);
	rotor[2] = oldRotor[2] + SiMath::randD(-range, range);
	rotor[3] = oldRotor[3] + SiMath::randD(-range, range);

	// normalise rotor such that it has unit norm
	double nr =
	    sqrt(rotor[0] * rotor[0] + rotor[1] * rotor[1] +
		 rotor[2] * rotor[2] + rotor[3] * rotor[3]);
	rotor[0] /= nr;
	rotor[1] /= nr;
	rotor[2] /= nr;
	rotor[3] /= nr;

    }				// end of endless while loop

    return res;
}



void ShapeAlignment::setMaxIterations(unsigned int i)
{
    _maxIter = i;
    return;
}



double *ShapeAlignment::_updateMatrixMap(AtomGaussian & a,
					 AtomGaussian & b)
{
    double *A = new double[17];

    // variables to store sum and difference of components
    double dx = (a.center.x - b.center.x);
    double dx2 = dx * dx;
    double dy = (a.center.y - b.center.y);
    double dy2 = dy * dy;
    double dz = (a.center.z - b.center.z);
    double dz2 = dz * dz;
    double sx = (a.center.x + b.center.x);
    double sx2 = sx * sx;
    double sy = (a.center.y + b.center.y);
    double sy2 = sy * sy;
    double sz = (a.center.z + b.center.z);
    double sz2 = sz * sz;

    // update overlap matrix
    double C = a.alpha * b.alpha / (a.alpha + b.alpha);
    A[0] = C * (dx2 + dy2 + dz2);
    A[1] = C * (dy * sz - dz * sy);
    A[2] = C * (dz * sx - dx * sz);
    A[3] = C * (dx * sy - dy * sx);
    A[4] = A[1];
    A[5] = C * (dx2 + sy2 + sz2);
    A[6] = C * (dx * dy - sx * sy);
    A[7] = C * (dx * dz - sx * sz);
    A[8] = A[2];
    A[9] = A[6];
    A[10] = C * (sx2 + dy2 + sz2);
    A[11] = C * (dy * dz - sy * sz);
    A[12] = A[3];
    A[13] = A[7];
    A[14] = A[11];
    A[15] = C * (sx2 + sy2 + dz2);

    // last elements holds overlap scaling constant
    // even number of overlap atoms => addition to volume
    // odd number => substraction
    if ((a.nbr + b.nbr) % 2 == 0) {
	A[16] = a.C * b.C * pow(PI / (a.alpha + b.alpha), 1.5);
    } else {
	A[16] = -a.C * b.C * pow(PI / (a.alpha + b.alpha), 1.5);
    }

    return A;
}