/*
 *                    BioJava development code
 *
 * This code may be freely distributed and modified under the
 * terms of the GNU Lesser General Public Licence.  This should
 * be distributed with the code.  If you do not have a copy,
 * see:
 *
 *      http://www.gnu.org/copyleft/lesser.html
 *
 * Copyright for this code is held jointly by the individual
 * authors.  These should be listed in @author doc comments.
 *
 * For more information on the BioJava project and its aims,
 * or to join the biojava-l mailing list, visit the home page
 * at:
 *
 *      http://www.biojava.org/
 *
 */
package org.biojava.nbio.structure.align.multiple.util;

import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;

import org.biojava.nbio.structure.Atom;
import org.biojava.nbio.structure.Calc;
import org.biojava.nbio.structure.SVDSuperimposer;
import org.biojava.nbio.structure.StructureException;
import org.biojava.nbio.structure.align.multiple.MultipleAlignment;
import org.biojava.nbio.structure.jama.Matrix;

/**
 * Utility class for calculating common scores of {@link MultipleAlignment}s.
 *
 * @author Spencer Bliven
 * @author Aleix Lafita
 * @since 4.1.0
 *
 */
public class MultipleAlignmentScorer {

	// Names for commonly used scores
	public static final String PROBABILITY = "Probability";
	public static final String CE_SCORE = "CE-score";
	public static final String RMSD = "RMSD";
	public static final String AVGTM_SCORE = "AvgTM-score";
	public static final String MC_SCORE = "MC-score";
	public static final String REF_RMSD = "Ref-RMSD";
	public static final String REFTM_SCORE = "RefTM-score";

	/**
	 * Calculates and puts the RMSD and the average TM-Score of the
	 * MultipleAlignment.
	 *
	 * @param alignment
	 * @throws StructureException
	 * @see #getAvgTMScore(MultipleAlignment)
	 * @see #getRMSD(MultipleAlignment)
	 */
	public static void calculateScores(MultipleAlignment alignment)
			throws StructureException {

		// Put RMSD
		List<Atom[]> trans = MultipleAlignmentTools.transformAtoms(alignment);
		alignment.putScore(RMSD, getRMSD(trans));

		// Put AvgTM-Score
		List<Integer> lengths = new ArrayList<Integer>(alignment.size());
		for (Atom[] atoms : alignment.getAtomArrays()) {
			lengths.add(atoms.length);
		}
		alignment.putScore(AVGTM_SCORE, getAvgTMScore(trans, lengths));
	}

	/**
	 * Calculates the RMSD of all-to-all structure comparisons (distances) of
	 * the given MultipleAlignment.
	 * <p>
	 * Complexity: T(n,l) = O(l*n^2), if n=number of structures and l=alignment
	 * length.
	 * <p>
	 * The formula used is just the sqroot of the average distance of all
	 * possible pairs of atoms. Thus, for R structures aligned over C columns
	 * without gaps, we have
	 *
	 * <pre>
	 * RMSD = \sqrt{1/(C*(R*(R-1)/2)) * \sum_{r1=1}^{R-1}
	 * \sum_{r2=r1+1}^{R-1} \sum_{j=0}^{C-1} (atom[r1][c]-atom[r2][c])^2}
	 * </pre>
	 *
	 * @param alignment
	 * @return double RMSD
	 */
	public static double getRMSD(MultipleAlignment alignment) {
		List<Atom[]> trans = MultipleAlignmentTools.transformAtoms(alignment);
		return getRMSD(trans);
	}

	/**
	 * Calculates the RMSD of all-to-all structure comparisons (distances),
	 * given a set of superimposed atoms.
	 *
	 * @param transformed
	 * @return double RMSD
	 * @see #getRMSD(MultipleAlignment)
	 */
	public static double getRMSD(List<Atom[]> transformed) {

		double sumSqDist = 0;
		int comparisons = 0;

		for (int r1 = 0; r1 < transformed.size(); r1++) {
			for (int c = 0; c < transformed.get(r1).length; c++) {
				Atom refAtom = transformed.get(r1)[c];
				if (refAtom == null)
					continue;

				double nonNullSqDist = 0;
				int nonNullLength = 0;
				for (int r2 = r1 + 1; r2 < transformed.size(); r2++) {
					Atom atom = transformed.get(r2)[c];
					if (atom != null) {
						nonNullSqDist += Calc.getDistanceFast(refAtom, atom);
						nonNullLength++;
					}
				}

				if (nonNullLength > 0) {
					comparisons++;
					sumSqDist += nonNullSqDist / nonNullLength;
				}
			}
		}
		return Math.sqrt(sumSqDist / comparisons);
	}

	/**
	 * /** Calculates the average RMSD from all structures to a reference s
	 * tructure, given a set of superimposed atoms.
	 * <p>
	 * Complexity: T(n,l) = O(l*n), if n=number of structures and l=alignment
	 * length.
	 * <p>
	 * For ungapped alignments, this is just the sqroot of the average distance
	 * from an atom to the aligned atom from the reference. Thus, for R
	 * structures aligned over C columns (with structure 0 as the reference), we
	 * have:
	 *
	 * <pre>
	 * RefRMSD = \sqrt{ 1/(C*(R-1)) * \sum_{r=1}^{R-1} \sum_{j=0}^{C-1}
	 * (atom[1][c]-atom[r][c])^2 }
	 * </pre>
	 * <p>
	 * For gapped alignments, null atoms are omitted from consideration, so that
	 * the RMSD is the average over all columns with non-null reference of the
	 * average RMSD within the non-null elements of the column.
	 *
	 * @param alignment
	 *            MultipleAlignment
	 * @param ref
	 *            reference structure index
	 * @return
	 */
	public static double getRefRMSD(MultipleAlignment alignment, int ref) {
		List<Atom[]> trans = MultipleAlignmentTools.transformAtoms(alignment);
		return getRefRMSD(trans, ref);
	}

	/**
	 * Calculates the average RMSD from all structures to a reference s
	 * tructure, given a set of superimposed atoms.
	 * <p>
	 * Complexity: T(n,l) = O(l*n), if n=number of structures and l=alignment
	 * length.
	 * <p>
	 * For ungapped alignments, this is just the sqroot of the average distance
	 * from an atom to the aligned atom from the reference. Thus, for R
	 * structures aligned over C columns (with structure 0 as the reference), we
	 * have:
	 *
	 * <pre>
	 * RefRMSD = \sqrt{ 1/(C*(R-1)) * \sum_{r=1}^{R-1} \sum_{j=0}^{C-1}
	 * (atom[1][c]-atom[r][c])^2 }
	 * </pre>
	 * <p>
	 * For gapped alignments, null atoms are omitted from consideration, so that
	 * the RMSD is the average over all columns with non-null reference of the
	 * average RMSD within the non-null elements of the column.
	 *
	 * @param transformed
	 * @param reference
	 * @return
	 */
	public static double getRefRMSD(List<Atom[]> transformed, int reference) {

		double sumSqDist = 0;
		int totalLength = 0;
		for (int c = 0; c < transformed.get(reference).length; c++) {
			Atom refAtom = transformed.get(reference)[c];
			if (refAtom == null)
				continue;

			double nonNullSqDist = 0;
			int nonNullLength = 0;
			for (int r = 0; r < transformed.size(); r++) {
				if (r == reference)
					continue;
				Atom atom = transformed.get(r)[c];
				if (atom != null) {
					nonNullSqDist += Calc.getDistanceFast(refAtom, atom);
					nonNullLength++;
				}
			}

			if (nonNullLength > 0) {
				totalLength++;
				sumSqDist += nonNullSqDist / nonNullLength;
			}
		}
		return Math.sqrt(sumSqDist / totalLength);
	}

	/**
	 * Calculates the average TMScore of all the possible pairwise structure
	 * comparisons of the given alignment.
	 * <p>
	 * Complexity: T(n,l) = O(l*n^2), if n=number of structures and l=alignment
	 * length.
	 *
	 * @param alignment
	 * @return double Average TMscore
	 * @throws StructureException
	 */
	public static double getAvgTMScore(MultipleAlignment alignment)
			throws StructureException {

		List<Atom[]> trans = MultipleAlignmentTools.transformAtoms(alignment);

		List<Integer> lengths = new ArrayList<Integer>(alignment.size());
		for (Atom[] atoms : alignment.getAtomArrays()) {
			lengths.add(atoms.length);
		}
		return getAvgTMScore(trans, lengths);
	}

	/**
	 * Calculates the average TMScore all the possible pairwise structure
	 * comparisons of the given a set of superimposed Atoms and the original
	 * structure lengths.
	 * <p>
	 * Complexity: T(n,l) = O(l*n^2), if n=number of structures and l=alignment
	 * length.
	 *
	 * @param transformed
	 *            aligned Atoms transformed
	 * @param lengths
	 *            lengths of the structures in residue number
	 * @return double Average TMscore
	 * @throws StructureException
	 */
	public static double getAvgTMScore(List<Atom[]> transformed,
			List<Integer> lengths) throws StructureException {

		if (transformed.size() != lengths.size())
			throw new IllegalArgumentException("Input sizes differ.");

		double sumTM = 0;
		int comparisons = 0;

		for (int r1 = 0; r1 < transformed.size(); r1++) {
			for (int r2 = r1 + 1; r2 < transformed.size(); r2++) {
				int len = transformed.get(r1).length;
				// Remove nulls from both arrays
				Atom[] ref = new Atom[len];
				Atom[] aln = new Atom[len];
				int nonNullLen = 0;
				for (int c = 0; c < len; c++) {
					if (transformed.get(r1)[c] != null
							&& transformed.get(r2)[c] != null) {
						ref[nonNullLen] = transformed.get(r1)[c];
						aln[nonNullLen] = transformed.get(r2)[c];
						nonNullLen++;
					}
				}
				// Truncate nulls
				if (nonNullLen < len) {
					ref = Arrays.copyOf(ref, nonNullLen);
					aln = Arrays.copyOf(aln, nonNullLen);
				}
				sumTM += SVDSuperimposer.getTMScore(ref, aln, lengths.get(r1),
						lengths.get(r2));
				comparisons++;
			}
		}
		return sumTM / comparisons;
	}

	/**
	 * Calculates the average TMScore from all structures to a reference
	 * structure, given a set of superimposed atoms.
	 * <p>
	 * Complexity: T(n,l) = O(l*n), if n=number of structures and l=alignment
	 * length.
	 *
	 * @param alignment
	 * @param reference
	 *            Index of the reference structure
	 * @return
	 * @throws StructureException
	 */
	public static double getRefTMScore(MultipleAlignment alignment, int ref)
			throws StructureException {

		List<Atom[]> trans = MultipleAlignmentTools.transformAtoms(alignment);

		List<Integer> lengths = new ArrayList<Integer>(alignment.size());
		for (Atom[] atoms : alignment.getAtomArrays()) {
			lengths.add(atoms.length);
		}
		return getRefTMScore(trans, lengths, ref);
	}

	/**
	 * Calculates the average TMScore from all structures to a reference
	 * structure, given a set of superimposed atoms.
	 * <p>
	 * Complexity: T(n,l) = O(l*n^2), if n=number of structures and l=alignment
	 * length.
	 *
	 * @param transformed
	 *            Arrays of aligned atoms, after superposition
	 * @param lengths
	 *            lengths of the full input structures
	 * @param reference
	 *            Index of the reference structure
	 * @return
	 * @throws StructureException
	 */
	public static double getRefTMScore(List<Atom[]> transformed,
			List<Integer> lengths, int reference) throws StructureException {

		if (transformed.size() != lengths.size())
			throw new IllegalArgumentException("Input sizes differ");

		double sumTM = 0;
		int comparisons = 0;

		int len = transformed.get(reference).length;
		for (int r = 0; r < transformed.size(); r++) {
			if (r == reference)
				continue;
			// remove nulls from both arrays
			Atom[] ref = new Atom[len];
			Atom[] aln = new Atom[len];
			int nonNullLen = 0;
			for (int c = 0; c < len; c++) {
				if (transformed.get(reference)[c] != null
						&& transformed.get(r)[c] != null) {
					ref[nonNullLen] = transformed.get(reference)[c];
					aln[nonNullLen] = transformed.get(r)[c];
					nonNullLen++;
				}
			}
			// truncate nulls
			if (nonNullLen < len) {
				ref = Arrays.copyOf(ref, nonNullLen);
				aln = Arrays.copyOf(aln, nonNullLen);
			}
			sumTM += SVDSuperimposer.getTMScore(ref, aln,
					lengths.get(reference), lengths.get(r));
			comparisons++;
		}
		return sumTM / comparisons;
	}

	/**
	 * Calculates the MC score, specific for the MultipleAlignment algorithm.
	 * The score function is modified from the original CEMC paper, making it
	 * continuous and differentiable.
	 * <p>
	 * The maximum score of a match is 20, and the penalties for gaps are part
	 * of the input. The half-score distance, d0, is chosen as in the TM-score.
	 * <p>
	 * Complexity: T(n,l) = O(l*n^2), if n=number of structures and l=alignment
	 * length.
	 *
	 * @param alignment
	 * @param gapOpen
	 *            penalty for gap opening (reasonable values are in the range
	 *            (1.0-20.0)
	 * @param gapExtension
	 *            penalty for extending a gap (reasonable values are in the
	 *            range (0.5-10.0)
	 * @param dCutoff
	 *            the distance cutoff
	 * @return the value of the score
	 * @throws StructureException
	 */
	public static double getMCScore(MultipleAlignment alignment,
			double gapOpen, double gapExtension, double dCutoff)
			throws StructureException {

		List<Atom[]> trans = MultipleAlignmentTools.transformAtoms(alignment);

		// Calculate d0: same as the one in TM score
		int minLen = Integer.MAX_VALUE;
		for (Atom[] atoms : alignment.getAtomArrays())
			if (atoms.length < minLen)
				minLen = atoms.length;
		double d0 = 1.24 * Math.cbrt((minLen) - 15.) - 1.8;

		// Calculate the distance cutoff penalization
		double A = 20.0 / (1 + (dCutoff * dCutoff) / (d0 * d0));

		return getMCScore(trans, d0, gapOpen, gapExtension, A);
	}

	/**
	 * Calculates the MC score, specific for the MultipleAlignment algorithm.
	 * The score function is modified from the original CEMC paper, making it
	 * continuous and differentiable.
	 * <p>
	 * The maximum score of a match is 20, and the penalties for gaps are part
	 * of the input.
	 * <p>
	 * Complexity: T(n,l) = O(l*n^2), if n=number of structures and l=alignment
	 * length.
	 *
	 * @param transformed
	 *            List of transformed Atom arrays
	 * @param d0
	 *            parameter for the half-score distance
	 * @param gapOpen
	 *            penalty for gap opening (reasonable values are in the range
	 *            (1.0-20.0)
	 * @param gapExtension
	 *            penalty for extending a gap (reasonable values are in the
	 *            range (0.5-10.0)
	 * @param A
	 *            the distance cutoff penalization
	 * @return the value of the score
	 * @throws StructureException
	 */
	private static double getMCScore(List<Atom[]> trans, double d0,
			double gapOpen, double gapExtension, double A)
			throws StructureException {

		int size = trans.size();
		int length = trans.get(0).length;
		Matrix residueDistances = new Matrix(size, length, -1);
		double scoreMC = 0.0;
		int openGaps = 0;
		int extensionGaps = 0;

		// Calculate the average residue distances
		for (int r1 = 0; r1 < size; r1++) {
			boolean gapped = false;
			for (int c = 0; c < trans.get(r1).length; c++) {
				Atom refAtom = trans.get(r1)[c];
				// Calculate the gap extension and opening on the fly
				if (refAtom == null) {
					if (gapped)
						extensionGaps++;
					else {
						gapped = true;
						openGaps++;
					}
					continue;
				} else
					gapped = false;

				for (int r2 = r1 + 1; r2 < size; r2++) {
					Atom atom = trans.get(r2)[c];
					if (atom != null) {
						double distance = Calc.getDistance(refAtom, atom);
						if (residueDistances.get(r1, c) == -1) {
							residueDistances.set(r1, c, 1 + distance);
						} else {
							residueDistances.set(r1, c,
									residueDistances.get(r1, c) + distance);
						}
						if (residueDistances.get(r2, c) == -1) {
							residueDistances.set(r2, c, 1 + distance);
						} else {
							residueDistances.set(r2, c,
									residueDistances.get(r2, c) + distance);
						}
					}
				}
			}
		}

		for (int c = 0; c < length; c++) {
			int nonNullRes = 0;
			for (int r = 0; r < size; r++) {
				if (residueDistances.get(r, c) != -1)
					nonNullRes++;
			}
			for (int r = 0; r < size; r++) {
				if (residueDistances.get(r, c) != -1) {
					residueDistances.set(r, c, residueDistances.get(r, c)
							/ nonNullRes);
				}
			}
		}

		// Sum all the aligned residue scores
		for (int row = 0; row < size; row++) {
			for (int col = 0; col < length; col++) {
				if (residueDistances.get(row, col) == -1)
					continue;
				double d1 = residueDistances.get(row, col);
				double resScore = 20.0 / (1 + (d1 * d1) / (d0 * d0));
				scoreMC += resScore - A;
			}
		}

		// Apply the Gap penalty and return
		return scoreMC - (openGaps * gapOpen + extensionGaps * gapExtension);
	}
}
