/*
 *                  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/
 *
 * Created on May 21, 2006
 *
 */
package org.biojava.nbio.structure.align;

import org.biojava.nbio.structure.*;
import org.biojava.nbio.structure.align.ce.GuiWrapper;
import org.biojava.nbio.structure.align.helper.AlignUtils;
import org.biojava.nbio.structure.align.helper.JointFragments;
import org.biojava.nbio.structure.align.pairwise.*;
import org.biojava.nbio.structure.geometry.Matrices;
import org.biojava.nbio.structure.geometry.SuperPositions;
import org.biojava.nbio.structure.io.PDBFileParser;
import org.biojava.nbio.structure.io.PDBFileReader;
import org.biojava.nbio.structure.jama.Matrix;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.FileOutputStream;
import java.io.InputStream;
import java.io.PrintStream;
import java.util.ArrayList;
import java.util.Collections;
import java.util.Comparator;
import java.util.List;

import javax.vecmath.Matrix4d;

/**
 * Perform a pairwise protein structure superimposition.
 *
 * <p>
 * The algorithm is a distance matrix based, rigid body protein structure
 * superimposition. It is based on a variation of the PSC++ algorithm provided
 * by Peter Lackner (Peter.Lackner@sbg.ac.at, personal communication) .
 * </p>
 *
 *
 *
 * <h2>Example</h2>
 *
 * <pre>
 *  public void run(){
 *
 * 		// first load two example structures
 * 		{@link InputStream} inStream1 = this.getClass().getResourceAsStream("/files/5pti.pdb");
 * 		{@link InputStream} inStream2 = this.getClass().getResourceAsStream("/files/1tap.pdb");
 *
 * 		{@link Structure} structure1 = null;
 * 		{@link Structure} structure2 = null;
 *
 * 		{@link PDBFileParser} pdbpars = new {@link PDBFileParser}();
 * 		structure1 = pdbpars.parsePDBFile(inStream1) ;
 * 		structure2 = pdbpars.parsePDBFile(inStream2);
 *
 *
 * 		// calculate structure superimposition for two complete structures
 * 		{@link StructurePairAligner} aligner = new {@link StructurePairAligner}();
 *
 *
 * 			// align the full 2 structures with default parameters.
 * 			// see StructurePairAligner for more options and how to align
 * 			// any set of Atoms
 * 			aligner.align(structure1,structure2);
 *
 * 			{@link AlternativeAlignment}[] aligs = aligner.getAlignments();
 * 			{@link AlternativeAlignment} a = aligs[0];
 * 			System.out.println(a);
 *
 * 			//display the alignment in Jmol
 *
 * 			// first get an artificial structure for the alignment
 * 			{@link Structure} artificial = a.getAlignedStructure(structure1, structure2);
 *
 *
 * 			// and then send it to Jmol (only will work if Jmol is in the Classpath)
 *
 * 			BiojavaJmol jmol = new BiojavaJmol();
 * 			jmol.setTitle(artificial.getName());
 * 			jmol.setStructure(artificial);
 *
 * 			// color the two structures
 *
 *
 * 			jmol.evalString("select *; backbone 0.4; wireframe off; spacefill off; " +
 * 					"select not protein and not solvent; spacefill on;");
 * 			jmol.evalString("select *"+"/1 ; color red; model 1; ");
 *
 *
 * 			// now color the equivalent residues ...
 *
 * 			String[] pdb1 = a.getPDBresnum1();
 * 			for (String res : pdb1 ){
 * 				jmol.evalString("select " + res + "/1 ; backbone 0.6; color white;");
 * 			}
 *
 * 			jmol.evalString("select *"+"/2; color blue; model 2;");
 * 			String[] pdb2 = a.getPDBresnum2();
 * 			for (String res :pdb2 ){
 * 				jmol.evalString("select " + res + "/2 ; backbone 0.6; color yellow;");
 * 			}
 *
 *
 * 			// now show both models again.
 * 			jmol.evalString("model 0;");
 *
 * 	}
 * </pre>
 *
 *
 *
 * @author Andreas Prlic
 * @author Peter Lackner
 * @since 1.4
 * @version %I% %G%
 */
public class StructurePairAligner {

	private final static Logger logger = LoggerFactory
			.getLogger(StructurePairAligner.class);

	private AlternativeAlignment[] alts;
	private Matrix distanceMatrix;
	private StrucAligParameters params;
	private FragmentPair[] fragPairs;

	private final List<AlignmentProgressListener> listeners = new ArrayList<>();

	public StructurePairAligner() {
		super();
		params = StrucAligParameters.getDefaultParameters();
		reset();
		alts = new AlternativeAlignment[0];
		distanceMatrix = new Matrix(0, 0);
	}

	public void addProgressListener(AlignmentProgressListener li) {
		listeners.add(li);
	}

	public void clearListeners() {
		listeners.clear();
	}

	/**
	 * example usage of this class
	 *
	 * @param args
	 */
	public static void main(String[] args) throws Exception {
		// UPDATE THE FOLLOWING LINES TO MATCH YOUR SETUP

		PDBFileReader pdbr = new PDBFileReader();
		pdbr.setPath("/Users/andreas/WORK/PDB/");

		// String pdb1 = "1crl";
		// String pdb2 = "1ede";

		String pdb1 = "1buz";
		String pdb2 = "1ali";
		String outputfile = "/tmp/alig_" + pdb1 + "_" + pdb2 + ".pdb";

		// NO NEED TO DO CHANGE ANYTHING BELOW HERE...

		StructurePairAligner sc = new StructurePairAligner();

		// step1 : read molecules

		logger.info("aligning {} vs. {}", pdb1, pdb2);

		Structure s1 = pdbr.getStructureById(pdb1);
		Structure s2 = pdbr.getStructureById(pdb2);

		// step 2 : do the calculations
		sc.align(s1, s2);

		AlternativeAlignment[] aligs = sc.getAlignments();

		// cluster similar results together
		ClusterAltAligs.cluster(aligs);

		// print the result:
		// the AlternativeAlignment object gives also access to rotation
		// matrices / shift vectors.
		for (AlternativeAlignment aa : aligs) {
			logger.info("Alternative Alignment: {}", aa);
		}

		// convert AlternativeAlignemnt 1 to PDB file, so it can be opened with
		// a viewer (e.g. Jmol, Rasmol)

		if (aligs.length > 0) {
			AlternativeAlignment aa1 = aligs[0];
			String pdbstr = aa1.toPDB(s1, s2);

			logger.info("writing alignment to {}", outputfile);
			FileOutputStream out = new FileOutputStream(outputfile);
			PrintStream p = new PrintStream(out);

			p.println(pdbstr);

			p.close();
			out.close();
		}

		// display the alignment in Jmol
		// only will work if Jmol is in the Classpath

		if (aligs.length > 0) {

			if (!GuiWrapper.isGuiModuleInstalled()) {
				logger.error("Could not find structure-gui modules in classpath, please install first!");
			}

		}

	}

	private void reset() {
		alts = new AlternativeAlignment[0];
		distanceMatrix = new Matrix(0, 0);
		fragPairs = new FragmentPair[0];

	}

	/**
	 * get the results of step 1 - the FragmentPairs used for seeding the
	 * alignment
	 *
	 * @return a FragmentPair[] array
	 */

	public FragmentPair[] getFragmentPairs() {
		return fragPairs;
	}

	public void setFragmentPairs(FragmentPair[] fragPairs) {
		this.fragPairs = fragPairs;
	}

	/**
	 * return the alternative alignments that can be found for the two
	 * structures
	 *
	 * @return AlternativeAlignment[] array
	 */
	public AlternativeAlignment[] getAlignments() {
		return alts;
	}

	/**
	 * return the difference of distance matrix between the two structures
	 *
	 * @return a Matrix
	 */
	public Matrix getDistMat() {
		return distanceMatrix;
	}

	/**
	 * get the parameters.
	 *
	 * @return the Parameters.
	 */
	public StrucAligParameters getParams() {
		return params;
	}

	/**
	 * set the parameters to be used for the algorithm
	 *
	 * @param params
	 *            the Parameter object
	 */
	public void setParams(StrucAligParameters params) {
		this.params = params;
	}

	/**
	 * Calculate the alignment between the two full structures with default
	 * parameters
	 *
	 * @param s1
	 * @param s2
	 * @throws StructureException
	 */
	public void align(Structure s1, Structure s2) throws StructureException {

		align(s1, s2, params);
	}

	/**
	 * Calculate the alignment between the two full structures with user
	 * provided parameters
	 *
	 * @param s1
	 * @param s2
	 * @param params
	 * @throws StructureException
	 */
	public void align(Structure s1, Structure s2, StrucAligParameters params)
			throws StructureException {
		// step 1 convert the structures to Atom Arrays

		Atom[] ca1 = getAlignmentAtoms(s1);
		Atom[] ca2 = getAlignmentAtoms(s2);

		notifyStartingAlignment(s1.getName(), ca1, s2.getName(), ca2);
		align(ca1, ca2, params);
	}

	/**
	 * Align two chains from the structures. Uses default parameters.
	 *
	 * @param s1
	 * @param chainId1
	 * @param s2
	 * @param chainId2
	 */
	public void align(Structure s1, String chainId1, Structure s2,
			String chainId2) throws StructureException {
		align(s1, chainId1, s2, chainId2, params);
	}

	/**
	 * Aligns two chains from the structures using user provided parameters.
	 *
	 * @param s1
	 * @param chainId1
	 * @param s2
	 * @param chainId2
	 * @param params
	 * @throws StructureException
	 */
	public void align(Structure s1, String chainId1, Structure s2,
			String chainId2, StrucAligParameters params)
			throws StructureException {
		reset();
		this.params = params;

		Chain c1 = s1.getPolyChainByPDB(chainId1);
		Chain c2 = s2.getPolyChainByPDB(chainId2);

		Structure s3 = new StructureImpl();
		s3.addChain(c1);

		Structure s4 = new StructureImpl();
		s4.addChain(c2);

		Atom[] ca1 = getAlignmentAtoms(s3);
		Atom[] ca2 = getAlignmentAtoms(s4);

		notifyStartingAlignment(s1.getName(), ca1, s2.getName(), ca2);
		align(ca1, ca2, params);
	}

	/**
	 * Returns the atoms that are being used for the alignment. (E.g. Calpha
	 * only, etc.)
	 *
	 * @param s
	 * @return an array of Atoms objects
	 */
	public Atom[] getAlignmentAtoms(Structure s) {
		String[] atomNames = params.getUsedAtomNames();
		return StructureTools.getAtomArray(s, atomNames);
	}

	/**
	 * Calculate the protein structure superimposition, between two sets of
	 * atoms.
	 *
	 * @param ca1
	 *            set of Atoms of structure 1
	 * @param ca2
	 *            set of Atoms of structure 2
	 * @param params
	 *            the parameters to use for the alignment
	 * @throws StructureException
	 */
	public void align(Atom[] ca1, Atom[] ca2, StrucAligParameters params)
			throws StructureException {

		reset();
		this.params = params;

		long timeStart = System.currentTimeMillis();

		// step 1 get all Diagonals of length X that are similar between both
		// structures
		logger.debug(" length atoms1:" + ca1.length);
		logger.debug(" length atoms2:" + ca2.length);

		logger.debug("step 1 - get fragments with similar intramolecular distances ");

		int k = params.getDiagonalDistance();
		int k2 = params.getDiagonalDistance2();
		int fragmentLength = params.getFragmentLength();

		if (ca1.length < (fragmentLength + 1)) {
			throw new StructureException("structure 1 too short (" + ca1.length
					+ "), can not align");
		}
		if (ca2.length < (fragmentLength + 1)) {
			throw new StructureException("structure 2 too short (" + ca2.length
					+ "), can not align");
		}
		int rows = ca1.length - fragmentLength + 1;
		int cols = ca2.length - fragmentLength + 1;
		distanceMatrix = new Matrix(rows, cols, 0.0);

		double[] dist1 = AlignUtils.getDiagonalAtK(ca1, k);

		double[] dist2 = AlignUtils.getDiagonalAtK(ca2, k);
		double[] dist3 = new double[0];
		double[] dist4 = new double[0];
		if (k2 > 0) {
			dist3 = AlignUtils.getDiagonalAtK(ca1, k2);
			dist4 = AlignUtils.getDiagonalAtK(ca2, k2);
		}

		double[][] utmp = new double[][] { { 0, 0, 1 } };
		Atom unitvector = new AtomImpl();
		unitvector.setCoords(utmp[0]);

		List<FragmentPair> fragments = new ArrayList<>();

		for (int i = 0; i < rows; i++) {

			Atom[] catmp1 = AlignUtils.getFragment(ca1, i, fragmentLength);
			Atom center1 = AlignUtils.getCenter(ca1, i, fragmentLength);

			for (int j = 0; j < cols; j++) {

				double rdd1 = AlignUtils.rms_dk_diag(dist1, dist2, i, j,
						fragmentLength, k);
				double rdd2 = 0;
				if (k2 > 0)
					rdd2 = AlignUtils.rms_dk_diag(dist3, dist4, i, j,
							fragmentLength, k2);
				double rdd = rdd1 + rdd2;
				distanceMatrix.set(i, j, rdd);

				if (rdd < params.getFragmentMiniDistance()) {
					FragmentPair f = new FragmentPair(fragmentLength, i, j);
					Atom[] catmp2 = AlignUtils.getFragment(ca2, j,
							fragmentLength);
					Atom center2 = AlignUtils.getCenter(ca2, j,
							fragmentLength);

					f.setCenter1(center1);
					f.setCenter2(center2);

					Matrix4d t = SuperPositions.superpose(
							Calc.atomsToPoints(catmp1),
							Calc.atomsToPoints(catmp2));

					Matrix rotmat = Matrices.getRotationJAMA(t);
					f.setRot(rotmat);

					Atom aunitv = (Atom) unitvector.clone();
					Calc.rotate(aunitv, rotmat);
					f.setUnitv(aunitv);

					boolean doNotAdd = false;
					if (params.reduceInitialFragments()) {
						doNotAdd = FragmentJoiner.reduceFragments(
								fragments, f, distanceMatrix);

					}
					if (doNotAdd)
						continue;

					fragments.add(f);
				}
			}
		}

		notifyFragmentListeners(fragments);

		FragmentPair[] fp = fragments.toArray(new FragmentPair[0]);
		setFragmentPairs(fp);

		logger.debug(" got # fragment pairs: {}", fp.length);

		logger.debug("step 2 - join fragments");

		// step 2 combine them to possible models
		FragmentJoiner joiner = new FragmentJoiner();

		JointFragments[] frags;

		if (params.isJoinFast()) {
			// apply the quick alignment procedure.
			// less quality in alignments, better for DB searches...
			frags = joiner.approach_ap3(ca1, ca2, fp, params);

			joiner.extendFragments(ca1, ca2, frags, params);

		} else if (params.isJoinPlo()) {
			// this approach by StrComPy (peter lackner):
			frags = joiner.frag_pairwise_compat(fp, params.getAngleDiff(),
					params.getFragCompat(), params.getMaxrefine());

		} else {

			// my first implementation
			frags = joiner.approach_ap3(ca1, ca2, fp, params);
		}

		notifyJointFragments(frags);

		logger.debug(" number joint fragments: {}", frags.length);

		logger.debug("step 3 - refine alignments");

		List<AlternativeAlignment> aas = new ArrayList<>();
		for (int i = 0; i < frags.length; i++) {
			JointFragments f = frags[i];
			AlternativeAlignment a = new AlternativeAlignment();
			a.apairs_from_idxlst(f);
			a.setAltAligNumber(i + 1);
			a.setDistanceMatrix(distanceMatrix);

			try {
				if (params.getMaxIter() > 0) {

					a.refine(params, ca1, ca2);
				} else {

					a.finish(params, ca1, ca2);

				}
			} catch (StructureException e) {
				logger.error("Refinement of fragment {} failed", i, e);
			}

			a.calcScores(ca1, ca2);
			aas.add(a);
		}

		// sort the alternative alignments
		Comparator<AlternativeAlignment> comp = new AltAligComparator();
		aas.sort(comp);
		Collections.reverse(aas);

		alts = aas.toArray(new AlternativeAlignment[0]);
		// do final numbering of alternative solutions
		int aanbr = 0;
		for (AlternativeAlignment a : alts) {
			aanbr++;
			a.setAltAligNumber(aanbr);
		}

		logger.debug("total calculation time: {} ms.",
				(System.currentTimeMillis() - timeStart));
	}

	private void notifyStartingAlignment(String name1, Atom[] ca1,
			String name2, Atom[] ca2) {
		for (AlignmentProgressListener li : listeners) {
			li.startingAlignment(name1, ca1, name2, ca2);
		}
	}

	private void notifyFragmentListeners(List<FragmentPair> fragments) {

		for (AlignmentProgressListener li : listeners) {
			li.calculatedFragmentPairs(fragments);
		}

	}

	private void notifyJointFragments(JointFragments[] fragments) {
		for (AlignmentProgressListener li : listeners) {
			li.jointFragments(fragments);
		}
	}

}
