/* This class is based on the original FATCAT implementation by
 * <pre>
 * Yuzhen Ye & Adam Godzik (2003)
 * Flexible structure alignment by chaining aligned fragment pairs allowing twists.
 * Bioinformatics vol.19 suppl. 2. ii246-ii255.
 * https://www.ncbi.nlm.nih.gov/pubmed/14534198
 * </pre>
 *
 * Thanks to Yuzhen Ye and A. Godzik for granting permission to freely use and redistribute this 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.
 *
 *
 * Created on Jun 17, 2009
 * Created by Andreas Prlic - RCSB PDB
 *
 */

package org.biojava.nbio.structure.align;

import org.biojava.nbio.structure.*;
import org.biojava.nbio.structure.align.model.AFP;
import org.biojava.nbio.structure.align.model.AFPChain;
import org.biojava.nbio.structure.jama.Matrix;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

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

//import org.biojava.nbio.structure.align.gui.jmol.StructureAlignmentJmol;

public class AFPTwister
{
	private final static Logger logger = LoggerFactory.getLogger(AFPTwister.class);

//   private static final boolean showAlignmentSteps = false;

	/**
	 * calculate the total rmsd of the blocks
	 * output a merged pdb file for both proteins
	  protein 1, in chain A
	 protein 2 is twisted according to the twists detected, in chain B
	 * @return  twisted Groups
	 */
	public static Group[] twistPDB(AFPChain afpChain,Atom[] ca1, Atom[] ca2) throws StructureException {
		//--------------------------------------------------------

		if ( afpChain.isShortAlign())
			return new Group[0];

		List<AFP> afpSet = afpChain.getAfpSet();

		int blockNum = afpChain.getBlockNum();

		int     i,  b2, e2;


		//superimposing according to the initial AFP-chaining
		Atom[] origCA      = StructureTools.cloneAtomArray(ca2);
		Atom[] iniTwistPdb = StructureTools.cloneAtomArray(ca2);

		int[] blockResSize = afpChain.getBlockResSize();

		int[][][] blockResList = afpChain.getBlockResList();

		int[] afpChainList    = afpChain.getAfpChainList();
		int[] block2Afp       = afpChain.getBlock2Afp();
		int[] blockSize       = afpChain.getBlockSize();
		int[] focusAfpList    = afpChain.getFocusAfpList();

		if ( focusAfpList == null){
			focusAfpList = new int[afpChain.getMinLen()];
			afpChain.setFocusAfpList(focusAfpList);
		}

		int  focusAfpn = 0;
		e2 = 0;
		b2 = 0;


		logger.debug("blockNUm at twister: ", blockNum);

		for(int bk = 0; bk < blockNum; bk ++)       {

			// THIS IS TRANSFORMING THE ORIGINAL ca2 COORDINATES, NO CLONING...
			// copies the atoms over to iniTwistPdb later on in modifyCod

			transformOrigPDB(blockResSize[bk], blockResList[bk][0], blockResList[bk][1],  ca1, ca2, null,-1);

			//transform pro2 according to comparison of pro1 and pro2 at give residues
			if(bk > 0)      {
				b2 = e2;
			}

			logger.debug("b2 is {} before  modifyCon", b2);

			if(bk < (blockNum - 1))   {

				//bend at the middle of two consecutive AFPs
				int afpPos = afpChainList[block2Afp[bk] + blockSize[bk] - 1];
				AFP a1 = afpSet.get(afpPos);
				e2 = a1.getP2() ;

				int afpPos2 = afpChainList[block2Afp[bk + 1]];
				AFP a2 = afpSet.get(afpPos2);
				e2 = (a2.getP2() - e2)/ 2 + e2;
				logger.debug("e2 : {}", e2);
			} else{
				// last one is until the end...
				e2 = ca2.length;
			}

			// this copies the coordinates over into iniTwistPdb
			cloneAtomRange(iniTwistPdb, ca2, b2, e2);

			//bound[bk] = e2;
			for(i = 0; i < blockSize[bk]; i ++)     {
				focusAfpList[focusAfpn] = afpChainList[block2Afp[bk] + i];
				focusAfpn++;
			}
		}


		int focusResn = afp2Res(afpChain, focusAfpn,  focusAfpList, 0);

		afpChain.setTotalLenIni(focusResn);

		logger.debug(String.format("calrmsdini for %d residues", focusResn));

		double totalRmsdIni = calCaRmsd(ca1,iniTwistPdb, focusResn, afpChain.getFocusRes1() , afpChain.getFocusRes2() );


		afpChain.setTotalRmsdIni(totalRmsdIni);
		logger.debug("got iniRMSD: {}", totalRmsdIni);
		if ( totalRmsdIni == 5.76611141613097) {
			logger.debug("{}", afpChain.getAlnseq1());
			logger.debug("{}", afpChain.getAlnsymb());
			logger.debug("{}", afpChain.getAlnseq2());
		}

		afpChain.setFocusAfpList(focusAfpList);
		afpChain.setBlock2Afp(block2Afp);
		afpChain.setAfpChainList(afpChainList);

		return twistOptimized(afpChain,ca1,origCA);
	}

	/**       superimposing according to the optimized alignment
	 *
	 * @param afpChain
	 * @param ca1
	 * @param ca2
	 * @return Group array twisted.
	 * @throws StructureException
	 */
	public static Group[] twistOptimized(AFPChain afpChain,Atom[] ca1, Atom[] ca2) throws StructureException {

		Atom[] optTwistPdb = new Atom[ca2.length];

		int gPos = -1;
		for (Atom a : ca2){
			gPos++;
			optTwistPdb[gPos] = a;
		}

		int blockNum = afpChain.getBlockNum();

		int b2 =  0;
		int e2 =  0;
		int focusResn = 0;
		int[] focusRes1 = afpChain.getFocusRes1();
		int[] focusRes2 = afpChain.getFocusRes2();

		if(focusRes1 == null) {
			 focusRes1 = new int[afpChain.getCa1Length()];
			 afpChain.setFocusRes1(focusRes1);
		}
		if(focusRes2 == null) {
			 focusRes2 = new int[afpChain.getCa2Length()];
			 afpChain.setFocusRes2(focusRes2);
		}

		int[] optLen      = afpChain.getOptLen();
		int[][][] optAln  = afpChain.getOptAln();

		for(int bk = 0; bk < blockNum; bk ++)       {
			//  THIS IS TRANSFORMING THE ORIGINAL ca2 COORDINATES, NO CLONING...
			// copies the atoms over to iniTwistPdb later on in modifyCod
			transformOrigPDB(optLen[bk], optAln[bk][0], optAln[bk][1], ca1, ca2,afpChain,bk);

			//transform pro2 according to comparison of pro1 and pro2 at give residues
			if(bk > 0)      { b2 = e2; }
			if(bk < blockNum - 1)   { //bend at the middle of two consecutive blocks
				e2 = optAln[bk][1][optLen[bk] - 1];
				e2 = (optAln[bk + 1][1][0] - e2)/ 2 + e2;
			}
			else    { e2 = ca2.length; }
			cloneAtomRange(optTwistPdb, ca2, b2, e2);
			for(int i = 0; i < optLen[bk]; i ++)        {
				focusRes1[focusResn] = optAln[bk][0][i];
				focusRes2[focusResn] = optAln[bk][1][i];
				focusResn ++;
			}
		}
		int totalLenOpt = focusResn;
		logger.debug("calrmsdopt for {} residues", focusResn);
		double totalRmsdOpt = calCaRmsd(ca1, optTwistPdb, focusResn, focusRes1, focusRes2);
		logger.debug("got opt RMSD: {}", totalRmsdOpt);
		int optLength = afpChain.getOptLength();

		if(totalLenOpt != optLength)    {
			logger.warn("Final alignment length is different {} {}", totalLenOpt, optLength);
		}
		logger.debug("final alignment length {}, rmsd {}", focusResn, totalRmsdOpt);

		afpChain.setTotalLenOpt(totalLenOpt);
		afpChain.setTotalRmsdOpt(totalRmsdOpt);

		return StructureTools.cloneGroups(optTwistPdb);

	}


	/**
	 * transform the coordinates in the ca2 according to the superimposing of the given position pairs.
	 * No Cloning, transforms input atoms.
	 */
	// orig name: transPdb
	private static void transformOrigPDB(int n, int[] res1, int[] res2, Atom[] ca1, Atom[] ca2, AFPChain afpChain, int blockNr)
	throws StructureException
	{
		logger.debug("transforming original coordinates {} len1: {} res1: {} len2: {} res2: {}", n, ca1.length, res1.length, ca2.length, res2.length);

		Atom[] cod1 = getAtoms(ca1, res1, n,false);
		Atom[] cod2 = getAtoms(ca2, res2, n,false);

		//double  *cod1 = pro1->Cod4Res(n, res1);
		//double  *cod2 = pro2->Cod4Res(n, res2);

		Matrix r;
		Atom t;

		SVDSuperimposer svd = new SVDSuperimposer(cod1, cod2);

		r = svd.getRotation();
		t = svd.getTranslation();

		logger.debug("transPdb: transforming orig coordinates with matrix: {}", r);

		if ( afpChain != null ){
			Matrix[] ms = afpChain.getBlockRotationMatrix();
			if ( ms == null)
				ms = new Matrix[afpChain.getBlockNum()];

			ms[blockNr]= r;

			Atom[] shifts = afpChain.getBlockShiftVector();
			if ( shifts == null)
				shifts = new Atom[afpChain.getBlockNum()];
			shifts[blockNr] = t;

			afpChain.setBlockRotationMatrix(ms);
			afpChain.setBlockShiftVector(shifts);
		}

		for (Atom a : ca2){
			Calc.rotate(a.getGroup(),r);
			Calc.shift(a.getGroup(),t);
		}


	}

	// like Cod4Res
	// most likely the clone flag is not needed
	private static Atom[] getAtoms(Atom[] ca, int[] positions, int length, boolean clone){

		List<Atom> atoms = new ArrayList<Atom>();
		for ( int i = 0 ; i < length ; i++){
			int p = positions[i];
			Atom a;
			if ( clone ){
				a = (Atom)ca[p].clone();
				a.setGroup((Group)ca[p].getGroup().clone());
			}
			else {
				a = ca[p];
			}
			atoms.add(a);
		}
		return atoms.toArray(new Atom[atoms.size()]);
	}



	/** Clones and copies the Atoms from p2 into p1 range is between r1 and r2
	 *
	 * @param p1
	 * @param p2
	 * @param r1
	 * @param r2
	 */
	// orig name: modifyCod
	private static void cloneAtomRange(Atom[] p1, Atom[] p2, int r1, int r2)
	throws StructureException
	{

		logger.debug("modifyCod from: {} to: {}", r1, r2);

		// special clone method, can;t use StructureTools.cloneCAArray, since we access the data
		// slightly differently here.
		List<Chain> model = new ArrayList<Chain>();
		for(int i = r1; i < r2; i ++) {

			Group g = p2[i].getGroup();
			Group newG = (Group)g.clone();

			p1[i] = newG.getAtom(p2[i].getName());
			Chain parentC = g.getChain();

			Chain newChain = null;

			for( Chain c: model){
				if (c.getChainID().equals(parentC.getChainID())){
					newChain = c;
					break;
				}
			}

			if ( newChain == null){
				newChain = new ChainImpl();
				newChain.setChainID(parentC.getChainID());
				model.add(newChain);
			}

			newChain.addGroup(newG);


		} //modify caCod

	}



	/**
	 * Return the rmsd of the CAs between the input pro and this protein, at given positions.
	 * quite similar to transPdb but while that one transforms the whole ca2, this one only works on the res1 and res2 positions.
	 *
	 *  Modifies the coordinates in the second set of Atoms (pro).
	 *
	 * @return rmsd of CAs
	 */
	private static double calCaRmsd(Atom[] ca1, Atom[] pro, int resn, int[] res1, int[] res2) throws StructureException
	{



		Atom[] cod1 = getAtoms(ca1,  res1,resn,false);
		Atom[] cod2 = getAtoms(pro,  res2,resn,false);

		Matrix r;
		Atom t;

		if ( cod1.length == 0  || cod2.length == 0){
			logger.info("length of atoms  == 0!");
			return 99;
		}
		SVDSuperimposer svd = new SVDSuperimposer(cod1, cod2);

		r = svd.getRotation();
		t = svd.getTranslation();

		for (Atom a : cod2){

			Calc.rotate(a.getGroup(), r);
			Calc.shift(a.getGroup(),  t);
		}

		return SVDSuperimposer.getRMS(cod1, cod2);
	}

	/**
	 * Set the list of equivalent residues in the two proteins given a list of AFPs
	 *
	 * WARNING: changes the values for FocusRes1, focusRes2 and FocusResn in afpChain!
	 *
	 * @param afpChain the AFPChain to store resuts
	 * @param afpn nr of afp
	 * @param afpPositions
	 * @param listStart
	 * @return nr of eq residues
	 */

	public static int afp2Res(AFPChain afpChain, int afpn, int[] afpPositions , int listStart   )
	{
		int[] res1 = afpChain.getFocusRes1();
		int[] res2 = afpChain.getFocusRes2();
		int minLen = afpChain.getMinLen();

		int n = 0;

		List<AFP> afpSet =afpChain.getAfpSet();

		for(int i = listStart; i < listStart+afpn; i ++)      {
			int a = afpPositions[i];
			for(int j = 0; j < afpSet.get(a).getFragLen(); j ++)     {
				if(n >= minLen) {
					throw new RuntimeException("Error: too many residues in AFPChainer.afp2Res!");
				}
				res1[n] = afpSet.get(a).getP1() + j;
				res2[n] = afpSet.get(a).getP2() + j;
				n ++;
			}
		}

		afpChain.setFocusRes1(res1);
		afpChain.setFocusRes2(res2);
		afpChain.setFocusResn(n);

		if ( n == 0 ){
			logger.warn("n=0!!! + " + afpn + " " + listStart + " " + afpPositions.length);
		}
		return n;
	}

}
