/**
 * 
 * Written by Márton Münz and Philip C Biggin
 * Copyright (c) University of Oxford, United Kingdom
 * Visit http://sbcb.bioch.ox.ac.uk/jgromacs/
 * 
 * This source code file is part of JGromacs v1.0.
 * 
 * JGromacs v1.0 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, either version 3 of the License, or
 * (at your option) any later version.
 * 
 * JGromacs v1.0. 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 General Public License for more details.
 * 
 * You should have received a copy of the GNU General Public License
 * along with JGromacs v1.0. If not, see <http://www.gnu.org/licenses/>.
 * 
 */

package jgromacs.analysis;

import java.util.ArrayList;

import jgromacs.data.IndexSet;
import jgromacs.data.Point3D;
import jgromacs.data.PointList;
import jgromacs.data.Structure;
import jgromacs.data.Trajectory;
import jama.Matrix;
import jama.SingularValueDecomposition;

/**
 * Collection of methods for superposing structures
 *
 */
public class Superposition {

	// Non Weighted
	
	/**
	 * Calculates the superposition of a point list to another
	 * @param toBeSuperposed the point list to be superposed
	 * @param reference the reference point list
	 * @return superposed points
	 */
	public static PointList superposeTo(PointList toBeSuperposed, PointList reference){
		PointList referencecopy = (PointList)reference.clone();
		PointList forfitcopy = (PointList)toBeSuperposed.clone();
		referencecopy.centerPoints();
		forfitcopy.centerPoints();
		Matrix Y = referencecopy.getAsMatrix();
		Matrix X = forfitcopy.getAsMatrix();
		Matrix C = X.times(Y.transpose());
		SingularValueDecomposition svn = new SingularValueDecomposition(C);
		Matrix U = svn.getU();
		Matrix V = svn.getV();
		double d = Math.signum(C.det());
		Matrix middle = new Matrix(3,3,0);
		middle.set(0, 0, 1);
		middle.set(1, 1, 1);
		middle.set(2, 2, d);
		Matrix R = V.times(middle.times(U.transpose()));
		forfitcopy.rotate(R);
		forfitcopy.translate(reference.getCentroid());
		return forfitcopy;
	}
	
	/**
	 * Calculates the superposition of a structure to another
	 * @param toBeSuperposed the structure to be superposed
	 * @param reference the reference structure
	 * @return superposed structure
	 */
	public static Structure superposeTo(Structure toBeSuperposed, Structure reference){
		PointList fitted = superposeTo(toBeSuperposed.getAllAtomCoordinates(), reference.getAllAtomCoordinates());
		Structure ret = (Structure)toBeSuperposed.clone();
		ret.setAllAtomCoordinates(fitted);
		return ret;
	}
	
	/**
	 * Calculates the superposition of a structure to another using a subset of atoms for fitting
	 * @param toBeSuperposed the structure to be superposed
	 * @param indices1 first index set
	 * @param reference the reference structure
	 * @param indices2 second index set
	 * @return superposed structure
	 */
	public static Structure superposeTo(Structure toBeSuperposed, IndexSet indices1, Structure reference, IndexSet indices2){
		PointList superposed = (PointList)toBeSuperposed.getAllAtomCoordinates().clone();
		PointList referencecopy = (PointList)reference.getSubStructure(indices2).getAllAtomCoordinates().clone();
		PointList forfitcopy = (PointList)toBeSuperposed.getSubStructure(indices1).getAllAtomCoordinates().clone();
		Structure ret = (Structure)toBeSuperposed.clone();
		
		Point3D centroid =  toBeSuperposed.getSubStructure(indices1).getAllAtomCoordinates().getCentroid();
		
		for (int i = 0; i < superposed.getNumberOfPoints(); i++){
			superposed.getPoint(i).setX(superposed.getPoint(i).getX()-centroid.getX());
			superposed.getPoint(i).setY(superposed.getPoint(i).getY()-centroid.getY());
			superposed.getPoint(i).setZ(superposed.getPoint(i).getZ()-centroid.getZ());
		}	
		
		referencecopy.centerPoints();
		forfitcopy.centerPoints();
		Matrix Y = referencecopy.getAsMatrix();
		Matrix X = forfitcopy.getAsMatrix();
		Matrix C = X.times(Y.transpose());
		SingularValueDecomposition svn = new SingularValueDecomposition(C);
		Matrix U = svn.getU();
		Matrix V = svn.getV();
		double d = Math.signum(C.det());
		Matrix middle = new Matrix(3,3,0);
		middle.set(0, 0, 1);
		middle.set(1, 1, 1);
		middle.set(2, 2, d);
		Matrix R = V.times(middle.times(U.transpose()));
		
		superposed.rotate(R);
		superposed.translate(reference.getSubStructure(indices2).getAllAtomCoordinates().getCentroid());
		
		ret.setAllAtomCoordinates(superposed);
		return ret;
	}
	

	/**
	 * Calculates the superposition of each frame of a trajectory to a common reference frame
	 * @param t trajectory to be superposed
	 * @param reference reference point list
	 * @return superposed trajectory
	 */
	public static Trajectory superposeTo(Trajectory t, PointList reference){
		Trajectory ret = new Trajectory();
		for (int i = 0; i < t.getNumberOfFrames(); i++) {
			PointList frame = t.getFrameAsPointList(i);
			PointList fittedframe = superposeTo(frame, reference);
			ret.addFrame(fittedframe);
		}
		return ret;
	}
	
	/**
	 * Calculates the superposition of each frame of a trajectory to a common reference frame
	 * @param t trajectory to be superposed
	 * @param reference reference structure
	 * @return superposed trajectory
	 */
	public static Trajectory superposeTo(Trajectory t, Structure reference){
		return superposeTo(t, reference.getAllAtomCoordinates());
	}
	
	/**
	 * Calculates the superposition of each frame of a trajectory to a common reference frame 
	 * using a subset of atoms for fitting
	 * @param t trajectory to be superposed
	 * @param reference reference structure
	 * @param indices index set
	 * @return superposed trajectory
	 */
	public static Trajectory superposeTo(Trajectory t, Structure reference, IndexSet indices){
		Trajectory ret = new Trajectory();
		for (int i = 0; i < t.getNumberOfFrames(); i++) {
			Structure frame = t.getFrameAsStructure(i);
			Structure fittedframe = superposeTo(frame, indices, reference, indices);
			ret.addFrame(fittedframe.getAllAtomCoordinates());
		}
		return ret;
	}
	
	// Weighted
	
	/**
	 * Calculates the weighted superposition of a point list to another
	 * @param toBeSuperposed the point list to be superposed
	 * @param reference the reference point list
	 * @param weights vector of weights
	 * @return superposed points
	 */
	public static PointList weightedSuperposeTo(PointList toBeSuperposed, PointList reference, ArrayList<Double> weights){
		int numofpoints = toBeSuperposed.getNumberOfPoints();
		// Computing weighted center of mass:
		Point3D muA = new Point3D();
		Point3D muB = new Point3D();
		double sumw = 0;
		for (int i = 0; i < numofpoints; i++) {
			muA = muA.plus(toBeSuperposed.getPoint(i).multiplyByScalar(weights.get(i)));
			muB = muB.plus(reference.getPoint(i).multiplyByScalar(weights.get(i)));
			sumw += weights.get(i);
		}
		muA = muA.multiplyByScalar(1.0/sumw);
		muB = muB.multiplyByScalar(1.0/sumw);
	
		
		// Weighted covariance matrix:
		Matrix C = new Matrix(3,3,0);
		for (int i = 0; i < 3; i++) {
			for (int j = 0; j < 3; j++) {
				double sum = 0;
				for (int k = 0; k < numofpoints; k++) {
					double aki = 0;
					double muAi = 0;
					double bkj = 0;
					double muBj = 0;
					if (i==0) { 
						aki = toBeSuperposed.getPoint(k).getX();
						muAi = muA.getX();
					}
					if (i==1) { 
						aki = toBeSuperposed.getPoint(k).getY();
						muAi = muA.getY();
					}
					if (i==2) { 
						aki = toBeSuperposed.getPoint(k).getZ();
						muAi = muA.getZ();
					}
					
					if (j==0) { 
						bkj = reference.getPoint(k).getX();
						muBj = muB.getX();
					}
					if (j==1) { 
						bkj = reference.getPoint(k).getY();
						muBj = muB.getY();
					}
					if (j==2) { 
						bkj = reference.getPoint(k).getZ();
						muBj = muB.getZ();
					}
					sum+=weights.get(k)*(aki-muAi)*(bkj-muBj); 
				}
				C.set(i, j, sum);
			}
		}
				
		// SVD of C
		SingularValueDecomposition svn = new SingularValueDecomposition(C);
		Matrix U = svn.getU();
		Matrix V = svn.getV();
		double lambda = Math.signum(C.det());
		
		// Computing optimal rotation:
		Matrix middle = new Matrix(3,3,0);
		middle.set(0, 0, 1);
		middle.set(1, 1, 1);
		middle.set(2, 2, lambda);
		Matrix Rmin = V.times(middle.times(U.transpose()));
			
		// Computing optimal translation:
		Point3D Tmin = muB.minus(muA.transformByMatrix(Rmin));
	
		// Rotate and translate
		PointList forfitcopy = (PointList)toBeSuperposed.clone();
		forfitcopy.rotate(Rmin);
	
		forfitcopy.translate(Tmin);
		return forfitcopy;
		
	}

	/**
	 * Calculates the weighted superposition of a structure to another
	 * @param toBeSuperposed the structure to be superposed
	 * @param reference the reference structure
	 * @param weights vector of weights
	 * @return superposed structure
	 */
	public static Structure weightedSuperposeTo(Structure toBeSuperposed, Structure reference, ArrayList<Double> weights){
		PointList fitted = weightedSuperposeTo(toBeSuperposed.getAllAtomCoordinates(), reference.getAllAtomCoordinates(), weights);
		Structure ret = (Structure)toBeSuperposed.clone();
		ret.setAllAtomCoordinates(fitted);
		return ret;
	}
	
	/**
	 * Calculates the weighted superposition of each frame of a trajectory to a common reference frame
	 * @param t trajectory to be superposed
	 * @param reference reference point list
	 * @param weights vector of weights
	 * @return superposed trajectory
	 */
	public static Trajectory weightedSuperposeTo(Trajectory t, PointList reference, ArrayList<Double> weights){
		Trajectory ret = new Trajectory();
		for (int i = 0; i < t.getNumberOfFrames(); i++) {
			PointList frame = t.getFrameAsPointList(i);
			PointList fittedframe = weightedSuperposeTo(frame, reference, weights);
			ret.addFrame(fittedframe);
		}
		return ret;
	}
	
	/**
	 * Calculates the weighted superposition of each frame of a trajectory to a common reference frame
	 * @param t trajectory to be superposed
	 * @param reference reference structure
	 * @param weights vector of weights
	 * @return superposed trajectory
	 */
	public static Trajectory weightedSuperposeTo(Trajectory t, Structure reference, ArrayList<Double> weights){
		return weightedSuperposeTo(t, reference.getAllAtomCoordinates(), weights);
	}
	
}
