/*
 *                    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 Dec 1, 2009
 * Author: Andreas Prlic 
 *
 */

package org.biojava.nbio.structure.test.align.fatcat;

import junit.framework.TestCase;
import org.biojava.nbio.structure.*;
import org.biojava.nbio.structure.align.StructureAlignment;
import org.biojava.nbio.structure.align.fatcat.FatCatFlexible;
import org.biojava.nbio.structure.align.fatcat.FatCatRigid;
import org.biojava.nbio.structure.align.fatcat.calc.FatCatParameters;
import org.biojava.nbio.structure.align.model.AFPChain;
import org.biojava.nbio.structure.align.util.AtomCache;
import org.biojava.nbio.structure.align.xml.AFPChainXMLConverter;
import org.biojava.nbio.structure.align.xml.AFPChainXMLParser;
import org.biojava.nbio.structure.jama.Matrix;

import java.io.IOException;
import java.util.ArrayList;
import java.util.List;



public class TestFlexibleRotationMatrices extends TestCase{

  


   public void testFlexibleRotationMatrices(){

      String name1 = "1a21.A";
      String name2 = "1hwg.C";

      compare(name1,name2, false);
      compare(name1,name2, true);

      String name3 = "5pti.A";
      String name4 = "1znf.A";
      compare(name3,name4, false);
      compare(name3,name4, true);
   }

   private void compare(String name1, String name2, boolean doRigid){

	   AtomCache cache = new AtomCache();

      try {
         Atom[] ca1orig = cache.getAtoms(name1);
         Atom[] ca2orig = cache.getAtoms(name2);

         Atom[] ca1 = StructureTools.cloneAtomArray(ca1orig);
         Atom[] ca2 = StructureTools.cloneAtomArray(ca2orig);

         Atom[] ca3 = StructureTools.cloneAtomArray(ca2);

         AFPChain afpChain = getAlignment(name1, name2, ca1, ca2, doRigid);
         afpChain.setCalculationTime(-1);
         String xml = AFPChainXMLConverter.toXML(afpChain,ca1,ca2);
         //System.out.println(xml);
         AFPChain newChain = AFPChainXMLParser.fromXML (xml, ca1, ca3);
         
         Matrix[] maxs1 = afpChain.getBlockRotationMatrix();
         //Atom[] shifts1 = afpChain.getBlockShiftVector();
         double[] blockRmsd = afpChain.getBlockRmsd();

         assertTrue( afpChain.getBlockNum() == newChain.getBlockNum());

         // make sure the XML conversion worked OK.
         for ( int i = 0 ; i < newChain.getBlockNum();i++) {
        	
            assertTrue(compareMatrices(maxs1[i],newChain.getBlockRotationMatrix()[i]));
            //assertTrue(compareVectors(shifts1[i],newChain.getBlockShiftVector()[i]));
            assertTrue(compareRmsd(blockRmsd[i], newChain.getBlockRmsd()[i]) );
         }

         assertTrue(afpChain.getOptLength() == newChain.getOptLength());

         // get the aligned blocks and check RMSD
         int[] blockSize =afpChain.getBlockSize();

         int[][][] blocks1 = afpChain.getOptAln();
         int[][][] blocks2 = newChain.getOptAln();

         for ( int x = 0 ; x < blocks1.length && x < afpChain.getBlockNum() ; x++){
            for ( int y = 0 ; y < blocks1[x].length ; y++){
               for ( int z = 0 ; z < blocks1[x][y].length && z < blockSize[x] ; z++){
                  //System.out.println(x + " " + y + " " + z);
                  assertEquals("The values in the optAln arrays don't match! " + 
                        x + " " + z + " " + blocks1[x][y][z]+ " " + 
                        blocks2[x][y][z],blocks1[x][y][z], blocks2[x][y][z]);
               }
            }


            Atom[] ca1new = StructureTools.cloneAtomArray(ca1orig);
            Atom[] ca2new = StructureTools.cloneAtomArray(ca2orig);

            compareBlock(x,afpChain, newChain,ca1new,ca2new );

         }

      } catch (IOException e){
         fail(e.getMessage());
      } catch (StructureException e){
         fail(e.getMessage());
      }


   }

   private boolean compareRmsd(double rmsdOrig, double rmsdNew) {
      //System.out.println("orig: " + rmsdOrig + " " + rmsdNew);
      String rmsdString1 = String.format("%5.2f",rmsdOrig);
      String rmsdString2 = String.format("%5.2f",rmsdNew);

      return rmsdString1.equals(rmsdString2);


   }

   private void compareBlock(int blockNr, AFPChain afpChain, AFPChain newChain,
         Atom[] ca1, Atom[] ca2) throws StructureException {
      


      Atom[] ca1Copy = StructureTools.cloneAtomArray(ca1);
      Atom[] ca2Copy = StructureTools.cloneAtomArray(ca2);
      Atom[] ca2Copy2 = StructureTools.cloneAtomArray(ca2);

     // int[][][] blocks1 = afpChain.getOptAln();
      int[][][] blocks2 = newChain.getOptAln();

     // Matrix[] maxs1 = afpChain.getBlockRotationMatrix();
     // Atom[] shifts1 = afpChain.getBlockShiftVector();

      Matrix[] maxs2 = newChain.getBlockRotationMatrix();
      Atom[] shifts2 = newChain.getBlockShiftVector();

      // get the eqr atoms of block X:
      int[] optLen =afpChain.getOptLen();
      List<Atom> eqrPos1 = new ArrayList<Atom>();
      List<Atom> eqrPos2 = new ArrayList<Atom>();
      List<Atom> eqrPos2copy = new ArrayList<Atom>();
      for ( int z = 0 ; z < blocks2[blockNr][0].length && z < optLen[blockNr] ; z++){
         int pos1 = blocks2[blockNr][0][z];
         int pos2 = blocks2[blockNr][1][z];

         Atom c1 = ca1Copy[pos1];
         Atom c2 = ca2Copy[pos2];
         Atom c3 = ca2Copy2[pos2];

         eqrPos1.add(c1);
         eqrPos2.add(c2);
         eqrPos2copy.add(c3);
      }

      assertTrue("The nr of Atoms in block " + blockNr + " does not match the expected nr. Expected:" + afpChain.getOptLen()[blockNr] + " but found: " + eqrPos2.size() , eqrPos2.size() == afpChain.getOptLen()[blockNr]);



      // THIS IS ROTATING the coordinates according to what is in the file.

      Atom[] blockSet1 = eqrPos1.toArray(new Atom[eqrPos1.size()]);
      Atom[] blockSet2 = eqrPos2.toArray(new Atom[eqrPos2.size()]);
      Atom[] blockSet2copy = eqrPos2copy.toArray(new Atom[eqrPos2copy.size()]);


      //System.out.println(shift );

      // rotate group 2...
      for ( Atom a : blockSet2){
         for ( int i =0 ; i<= blockNr;i++ ) {
            Matrix max   = maxs2[  i];
            Atom   shift = shifts2[i];
            Calc.rotate(a, max);
            Calc.shift( a, shift);
         }
      }
      // calc RMSD


      double rmsdFile = SVDSuperimposer.getRMS(blockSet1, blockSet2);		

      // this is the value from the file. it never seems to match precisely, probably is calculated from initial block.
      // we can't reproduce the initial block, since we don;t serialize it.
      //double rmsdOrig =afpChain.getBlockRmsd()[blockNr];


      // THIS IS CALCULATING THE "correct" rotation matrix, that should be in the file

      SVDSuperimposer svd = new SVDSuperimposer(blockSet1, blockSet2copy);
      //double rmsdForce = SVDSuperimposer.getRMS(atomSet1, atomSet2);
      Matrix m = svd.getRotation();
      Atom   s  = svd.getTranslation();

      Matrix max   = maxs2[blockNr];
      Atom   shift = shifts2[blockNr];
    
      compareMatrices(max, m);

      
      if ( blockNr == 0) {
    	  compareVectors(shift, s);
      } else {
    	  System.err.println("Not testing shift vectors for blocks > 1. There is still a problem...");
    	  
      }

      for ( Atom a : ca2Copy2){
         Calc.rotate(a, m);
         Calc.shift( a, s);
      }
      double rmsd3 = SVDSuperimposer.getRMS(blockSet1,blockSet2copy);

      assertTrue("The RMSD values don;t match after rotation / shift for block " + blockNr + "! should be: " + rmsd3 + " but found: " +rmsdFile, compareRmsd(rmsd3, rmsdFile));



      //this fails: is fatcat is using the rmsd before optimization?
      //assertTrue("The RMSD values don;t match after rotation / shift for block " + blockNr + "! should be: " + rmsdOrig + " but found: " +rmsdNew, compareRmsd(rmsdOrig, rmsdNew));
      // get the RMSD between the aligned blocks...


   }

   private boolean compareVectors(Atom atom1, Atom atom2) throws StructureException {

	   //System.out.println(Math.abs(atom1.getX()-  atom2.getX()));
	   assertTrue("The X coordinates are too far apart!", Math.abs(atom1.getX()-  atom2.getX()) < 0.01);
	   assertTrue("The Y coordinates are too far apart!", Math.abs(atom1.getY()-  atom2.getY()) < 0.01);
	   assertTrue("The Z coordinates are too far apart!", Math.abs(atom1.getZ()-  atom2.getZ()) < 0.01);
      //return (atom1.getX() == atom2.getX() && atom1.getY() == atom2.getY() && atom1.getZ() == atom2.getZ());
	   return true;


   }

   private boolean compareMatrices(Matrix matrix1, Matrix matrix2) {

      String m1 = matrix1.toString();
      String m2 = matrix2.toString();

      return (m1.equals(m2));


   }

   private AFPChain getAlignment (String name1, String name2, Atom[] ca1, Atom[] ca2 , boolean doRigid) throws StructureException,IOException{
      FatCatParameters params = new FatCatParameters();

      StructureAlignment fatCat ;

      if ( doRigid)
         fatCat = new FatCatRigid();
      else 
         fatCat = new FatCatFlexible();

      AFPChain afpChain = fatCat.align(ca1,ca2,params);

      afpChain.setName1(name1);
      afpChain.setName2(name2);

      // flexible original results:
      //String fatcat = afpChain.toFatcat(ca1,ca2);
      //System.out.println(result1);


      //String xml = AFPChainXMLConverter.toXML(afpChain,ca1,ca2);


      return afpChain;
   }


}
