/* 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.
 * http://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.bio.structure.align.fatcat.calc;



import org.biojava.bio.structure.Atom;
import org.biojava.bio.structure.Calc;
import org.biojava.bio.structure.Group;
import org.biojava.bio.structure.SVDSuperimposer;
import org.biojava.bio.structure.StructureException;
import org.biojava.bio.structure.StructureTools;
import org.biojava.bio.structure.jama.Matrix;



public class StructureAlignmentOptimizer
{

   //private static final boolean showAlig = false;
   
   int pro1Len;
   int pro2Len;
   int maxLen;
   Atom[] cod1;
   Atom[] cod2;

   int[][] equSet;
   int equLen;
   int equLen0;
   double[][]sij;

   int maxKeepStep;
   int keepStep;

   double  Dc;   //the criteria for structural equivalent residues, eg. 3.0 (CE), 6.0(ProSup)
   double  rmsdCut;//the criteria for stoping optimization
   double  increase;
   double  stopLenPer;
   double  stopRmsdPer;
   double  stopRmsd;

   double  gapIni;
   double  gapExt;

   double rmsd;

   private static final boolean debug = FatCatAligner.debug;

   /**
     optimize the structural alignment by update the equivalent residues
      and then run dynamic programming
      input: len1 the length of structure 1; c1: the structure information of 1
             len2 the length of structure 2; c2: the structure information of 2
             iniLen and iniSet is the length and list of initial equivalent residues
           

    */

   public StructureAlignmentOptimizer(int b1, int end1, Atom[] c1, int b2, int end2, Atom[] c2,
         int iniLen, int[][] iniSet) throws StructureException{
      
      //System.err.println("optimizing range:" + b1 + "-" + end1 + "("+ (end1-b1) + ") b2:  " + b2 + "-" + end2+ "("+ (end2-b2) + ") iniLen " + iniLen);
      //System.out.println("ca1: " + c1.length + " ca2: " + c2.length);

      int len1 = end1-b1;
      int len2 = end2-b2;
           
      //System.err.println("len1: " + len1 + " len2:" + len2);
      
      pro1Len = len1;
      pro2Len = len2;
            
      cod1 = new Atom[len1];
      cod2 = new Atom[len2];
      
      for(int i = 0; i < len1; i ++)      {           
         Atom a = c1[i+b1];
         //cod1[i] = (Atom)a.clone();
         Group parent = (Group)a.getGroup().clone();
         //cod1[i].setParent(parent);
         cod1[i] = parent.getAtom(StructureTools.caAtomName);
         //cod1[i] = c1[i];
      }
      for(int i = 0; i < len2; i ++)      {
         Atom a = c2[i+b2];
         //cod2[i]= (Atom)a.clone();
         Group parent = (Group)a.getGroup().clone();
         //cod2[i].setParent(parent);
         cod2[i] = parent.getAtom(StructureTools.caAtomName);
         //cod2[i] = c2[i];
      }


      //initial equivalent sets
      maxLen = (len1 < len2)?len1:len2;

      equSet = new int[2][maxLen];
      for(int i = 0; i < iniLen; i ++)    {

         equSet[0][i] = iniSet[0][i];
         equSet[1][i] = iniSet[1][i];
         if(iniSet[0][i] > len1 || iniSet[1][i] > len2)  {
            throw new RuntimeException(String.format("StructureAlignmentOptimizer: focus exceeds the protein 1 or 2 length!"));            
         }
      }

      equLen    = iniLen;
      equLen0   = equLen;

      setParameters();

      sij =  new double[pro1Len][pro2Len];

//      if (showAlig) 
//         showCurrentAlignment(iniLen, equSet, "initial alignment");

   }

//   private void showCurrentAlignment(int len, int[][]set, String title){
//      BiojavaJmol jmol = new BiojavaJmol();
//      jmol.setTitle(title);
//
//      Chain c1 = new ChainImpl();
//      c1.setName("A");
//
//      Chain c2 = new ChainImpl();
//      c2.setName("B");
//      for(int i = 0; i < len; i ++)    {
//         Atom a1 = cod1[set[0][i]];
//         Atom a2 = cod2[set[1][i]];
//
//         Group g1 = a1.getParent();
//         Group g2 = a2.getParent();
//
//         try {
//            Group n1 = new AminoAcidImpl();
//            n1.setPDBCode(g1.getPDBCode());
//            n1.setPDBName(g1.getPDBName());
//            n1.addAtom(a1);
//
//            Group n2 = new AminoAcidImpl();
//            n2.setPDBCode(g2.getPDBCode());
//            n2.setPDBName(g2.getPDBName());
//            n2.addAtom(a2);
//
//
//            c1.addGroup(n1);
//            c2.addGroup(n2);
//         } catch (Exception e){
//            //
//         }
//
//      }
//
//      Structure s = new StructureImpl();
//      s.setPDBCode(title);
//      List<Chain> model1 = new ArrayList<Chain>();
//      model1.add(c1);
//      List<Chain> model2 = new ArrayList<Chain>();
//      model2.add(c2);
//      s.addModel(model1);
//      s.addModel(model2);
//      s.setNmr(true);
//      jmol.setStructure(s);
//      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 show both models again.
//      jmol.evalString("model 0;");
//
//   }


   /** run the optimization
    *
    * @param maxi maximum nr. of iterations
    * @throws StructureException
    */
   public void runOptimization(int maxi) throws StructureException{
      superimposeBySet();
      if ( debug)
         System.err.println("   initial rmsd " + rmsd);

//      if (showAlig) 
//         showCurrentAlignment(equLen, equSet, "after initial superimposeBySet Len:" +equLen + " rmsd:" +rmsd);

      maxKeepStep = 4;
      keepStep = 0;

      optimize(maxi);
   }

   /**
    * refer CE, similarity = Dc - dij, Dc is increased by 0.5 each cycle,
  optimization continues until either
 i)alignment length is less than 95% of alignment length before optimization
 ii)rmsd is less than 110% of rmsd at the cycle when condition i) was first satisfied
    */
   private void setParameters()
   {
      Dc = 3.0; //Dc = 2.0
      increase = 0.5;
      stopLenPer = 0.95;
      stopRmsdPer = 1.1;
      stopRmsd = -1.0;
      rmsdCut = 3.0;
      gapIni = 5.0;
      gapExt = 0.5;
   }


   /**
 //superimpose two structures according to the equivalent residues
    */
   private void superimposeBySet ()
   throws StructureException
   {

      //extract the coordinations of equivalent residues
      Atom[] tmp1 = new Atom[equLen];
      Atom[] tmp2 = new Atom[equLen];
      int     i,  r1, r2;

     
      for(i = 0; i < equLen; i ++)    {
         r1 = equSet[0][i];
         r2 = equSet[1][i];

         tmp1[i] =       cod1[ r1 ];
         tmp2[i] = (Atom)cod2[ r2 ].clone(); // have to be cloned!
         //tmp2[i] = cod2[ r2 ];
         

         /*try {
            System.out.println("before superimpos: " + equSet[0][i]+"-"+ equSet[1][i]+ " dist:" + Calc.getDistance(tmp1[i], cod2[equSet[1][i]]));
         } catch (Exception e){
            e.printStackTrace();
         }*/
      }

      //superimpose the equivalent residues
      SVDSuperimposer svd = new SVDSuperimposer(tmp1, tmp2);

      Matrix    m = svd.getRotation();
      Atom      t = svd.getTranslation();

      for (Atom a: tmp2) {
         
         Calc.rotate(a,m);
         Calc.shift(a,t);
         
      }
      
      // weird, why does it take the RMSD before the rotation?
      // the rmsd is only for the subset contained in the tmp arrays.
      rmsd = SVDSuperimposer.getRMS(tmp1,tmp2);

      //System.err.println("rmsd after superimpose by set: " + rmsd);

      //transform structure 2 according to the superimposition of the equivalent residues
      for (i = 0 ; i< cod2.length; i++) {
         Atom a = cod2[i];
         Calc.rotate(a,m);
         Calc.shift(a,t);
         
      }

  
      
//      for(i = 0; i < equLen; i ++)    {
//         try {
//            System.err.println("after superimpos: " + equSet[0][i]+"-"+ equSet[1][i]+ " dist:" + Calc.getDistance(tmp1[i], cod2[equSet[1][i]]));
//         } catch (Exception e){
//            e.printStackTrace();
//         }
//      }
      
       

   }


   private void optimize(int maxi) throws StructureException
   {
      long optStart = System.currentTimeMillis();
      if ( debug)
         System.out.println("Optimizing up to " + maxi + " iterations.. ");
      boolean ifstop = true;;
      int     i, alnLen;
      alnLen = 0;

      int[][]     alnList =  new int[2][maxLen];
      for(i = 0; i < maxi; i ++)      {

         //if ( debug){
         //   System.out.println("optimize iteration: " + i);
         //}

         calMatrix();

         FCAlignHelper aln = new FCAlignHelper(sij,pro1Len,pro2Len,gapIni, gapExt);

         //ALIGN0 *aln = new ALIGN0(sij, pro1Len, pro2Len, gapIni, gapExt);
         alnLen = aln.getAlignPos(alnList);
         if(alnLen < 3)  ifstop = true; //very rare, mark by Y.Y on 5/1/03
         else    ifstop = defineEquPos(alnLen, alnList);

         if(ifstop)      break;
         Dc += increase;

//         if (showAlig) 
//            if ( i == 0 )
//               showCurrentAlignment(alnLen, alnList,  "optimizing alignment - after " + i + " iterations alnLen:" + alnLen + " rmsd " + rmsd);     
      }

      if  (debug){
         if(i < maxi)    {
            System.out.println(String.format("   optimize converged at %d iterations\n", i));
         }
         else    System.out.println("   optimize stop without convergence\n");
         System.out.println("optimization time: " + (System.currentTimeMillis() - optStart) + " ms.");
      }

//      if (showAlig) 
//         showCurrentAlignment(alnLen, alnList,  "optimizing alignment - after " + i + " iterations alnLen:" + alnLen + " rmsd " + rmsd);     
   }

   //--------------------------------------------------------------------------------------------------------
   //the definition of matrix between residues:
   //              Sij = Dc^2 - Dij^2 if Dij <= Dc
   //                    0            else
   //--------------------------------------------------------------------------------------------------------
   private void calMatrix() throws StructureException
   {
      int     i, j;
      double  dis;
      for(i = 0; i < pro1Len; i ++)   {
         for(j = 0; j < pro2Len; j ++)   {
            dis = Calc.getDistance(cod1[i],cod2[j]);

            if(dis < Dc) {
               sij[i][j] = Dc - dis;               
            }
            else  {
               sij[i][j] = 0;
            }
         }
      }
   }

   /**the equivalent residues: residues where Dij &lt;= Dc and i,j is an aligned pair
 use the previous superimposing
    */

   private boolean defineEquPos(int alnLen, int[][] alnList)
   throws StructureException
   {
      int     i, r1, r2;
      int     equLenOld = equLen;
      int[][]    equSetOld = new int[2][equLenOld];
      for(i = 0; i < equLen; i ++)    {
         equSetOld[0][i] = equSet[0][i];
         equSetOld[1][i] = equSet[1][i];
      }
      double  rmsdOld = rmsd;
      double  dis;
      equLen = 0;
      //if (debug)
      //   System.out.println(String.format(" OPT: Dc %f, equLenOld %d, rmsdOld %f, alnLen %d", Dc, equLenOld, rmsdOld, alnLen));
      for(i = 0; i < alnLen; i ++)    {
         r1 = alnList[0][i];
         r2 = alnList[1][i];
         dis = Calc.getDistance(cod1[r1],cod2[r2]);         
         if(dis <= Dc)   {
            //System.out.println(r1 + "-"  + r2 + " d:" + dis); 
            equSet[0][equLen] = r1;
            equSet[1][equLen] = r2;
            equLen ++;
         }
      }

      superimposeBySet();

      //if (debug)
      //   System.out.println(String.format(" OPT: new equLen %d rmsd %f", equLen, rmsd));

      boolean     ifstop = false;

//      if (debug) {
//         System.out.print(" OPT: rmsd diff: " + Math.abs(rmsd - rmsdOld) + " equLens: " + equLenOld + ":"+ equLen);
//         if ( Math.abs(rmsd - rmsdOld) < 1e-10)
//            System.out.println(" NO DIFF!");
//         else 
//            System.out.println(" DIFF!");
//      }

      if((Math.abs(rmsd - rmsdOld) < 1e-10 ) && (equLenOld == equLen)) keepStep ++;
      else    keepStep = 0;

      if(keepStep > maxKeepStep)      {
         ifstop = true; //converge
      } //allowing up to maxKeepStep instead of 1 is essential for some special cases
      else if(stopRmsd < 0) {
         ifstop = false; //condition 1, continue
      }
      else if((rmsd <= stopRmsd * stopRmsdPer) || (rmsd < rmsdCut)) {
         ifstop = false; //condition 2, continue
      } //rmsdCut is adopted or not? to be tuned
      else    {
         ifstop = true; //get worse
      }


      if((stopRmsd <0) && (equLen >= stopLenPer * equLen0))    {
         //System.err.println("stopRmsd: " + stopRmsd + " setting to rmsd:" + rmsd);
         stopRmsd = rmsd; //condition 1
      }

      return ifstop;
   }


   public double optimizeResult(int[] optLen, int optLenpos, int[][] list)
   {
      optLen[optLenpos] = equLen;

      for(int i = 0; i < equLen; i ++)        {
         list[0][i] = equSet[0][i];
         list[1][i] = equSet[1][i];
      }
      return rmsd;
   }

}
