/*
 *                  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 4, 2005
 *
 */
package org.biojava.nbio.structure;

import javax.vecmath.Matrix4d;

import org.biojava.nbio.structure.jama.Matrix;
import org.biojava.nbio.structure.jama.SingularValueDecomposition;


/** A class that calculates the superimposition between two sets of atoms
 * inspired by the biopython SVDSuperimposer class...
 *
 *
 * example usage:
 * <pre>
                        // get some arbitrary amino acids from somewhere
                        String filename   =  "/Users/ap3/WORK/PDB/5pti.pdb" ;

                        PDBFileReader pdbreader = new PDBFileReader();
                        Structure struc = pdbreader.getStructure(filename);
                        Group g1 = (Group)struc.getChain(0).getGroup(21).clone();
                        Group g2 = (Group)struc.getChain(0).getGroup(53).clone();

                        if ( g1.getPDBName().equals("GLY")){
                                if ( g1 instanceof AminoAcid){
                                        Atom cb = Calc.createVirtualCBAtom((AminoAcid)g1);
                                        g1.addAtom(cb);
                                }
                        }

                        if ( g2.getPDBName().equals("GLY")){
                                if ( g2 instanceof AminoAcid){
                                        Atom cb = Calc.createVirtualCBAtom((AminoAcid)g2);
                                        g2.addAtom(cb);
                                }
                        }

                        Structure struc2 = new StructureImpl((Group)g2.clone());

                        System.out.println(g1);
                        System.out.println(g2);


                        Atom[] atoms1 = new Atom[3];
                        Atom[] atoms2 = new Atom[3];

                        atoms1[0] = g1.getAtom("N");
                        atoms1[1] = g1.getAtom("CA");
                        atoms1[2] = g1.getAtom("CB");


                        atoms2[0] = g2.getAtom("N");
                        atoms2[1] = g2.getAtom("CA");
                        atoms2[2] = g2.getAtom("CB");


                        SVDSuperimposer svds = new SVDSuperimposer(atoms1,atoms2);


                        Matrix rotMatrix = svds.getRotation();
                        Atom tranMatrix = svds.getTranslation();


                        // now we have all the info to perform the rotations ...

                        Calc.rotate(struc2,rotMatrix);

                        //          shift structure 2 onto structure one ...
                        Calc.shift(struc2,tranMatrix);

                        //
                        // write the whole thing to a file to view in a viewer

                        String outputfile = "/Users/ap3/WORK/PDB/rotated.pdb";

                        FileOutputStream out= new FileOutputStream(outputfile);
                        PrintStream p =  new PrintStream( out );

                        Structure newstruc = new StructureImpl();

                        Chain c1 = new ChainImpl();
                        c1.setName("A");
                        c1.addGroup(g1);
                        newstruc.addChain(c1);

                        Chain c2 = struc2.getChain(0);
                        c2.setName("B");
                        newstruc.addChain(c2);

                        // show where the group was originally ...
                        Chain c3 = new ChainImpl();
                        c3.setName("C");
                        //c3.addGroup(g1);
                        c3.addGroup(g2);

                        newstruc.addChain(c3);
                        p.println(newstruc.toPDB());

                        p.close();

                        System.out.println("wrote to file " + outputfile);

                </pre>
 *
 *
 * @author Andreas Prlic
 * @since 1.5
 * @version %I% %G%

 */
public class SVDSuperimposer {

        Matrix rot;
        Matrix tran;

        Matrix centroidA;
        Matrix centroidB;

        /** Create a SVDSuperimposer object and calculate a SVD superimposition of two sets of atoms.
         *
         * @param atomSet1 Atom array 1
         * @param atomSet2 Atom array 2
         * @throws StructureException
         */
        public SVDSuperimposer(Atom[] atomSet1,Atom[]atomSet2)
        throws StructureException{

                if ( atomSet1.length != atomSet2.length ){
                        throw new StructureException("The two atom sets are not of same length!");
                }

                Atom cena = Calc.getCentroid(atomSet1);
                Atom cenb = Calc.getCentroid(atomSet2);

                double[][] centAcoords = new double[][]{{cena.getX(),cena.getY(),cena.getZ()}};
                centroidA = new Matrix(centAcoords);

                double[][] centBcoords = new double[][]{{cenb.getX(),cenb.getY(),cenb.getZ()}};
                centroidB = new Matrix(centBcoords);

                //      center at centroid

                Atom[] ats1 = Calc.centerAtoms(atomSet1,cena);
                Atom[] ats2 = Calc.centerAtoms(atomSet2,cenb);

                double[][] coordSet1 = new double[ats1.length][3];
                double[][] coordSet2 = new double[ats2.length][3];

                // copy the atoms into the internal coords;
                for (int i =0 ; i< ats1.length;i++) {
                        coordSet1[i] = ats1[i].getCoords();
                        coordSet2[i] = ats2[i].getCoords();
                }

                calculate(coordSet1,coordSet2);


        }


        /** Do the actual calculation.
         *
         * @param coordSet1 coordinates for atom array 1
         * @param coordSet2 coordiantes for atom array 2
         */
        private void calculate(double[][] coordSet1, double[][]coordSet2){
                // now this is the bridge to the Jama package:
                Matrix a = new Matrix(coordSet1);
                Matrix b = new Matrix(coordSet2);


//      # correlation matrix

                Matrix b_trans = b.transpose();
                Matrix corr = b_trans.times(a);


                SingularValueDecomposition svd = corr.svd();

                Matrix u = svd.getU();
                // v is alreaady transposed ! difference to numermic python ...
                Matrix vt =svd.getV();

                Matrix vt_orig = (Matrix) vt.clone();
                Matrix u_transp = u.transpose();

                Matrix rot_nottrans = vt.times(u_transp);
                rot = rot_nottrans.transpose();

                // check if we have found a reflection

                double det = rot.det();

                 if (det<0) {
                        vt = vt_orig.transpose();
                        vt.set(2,0,(0 - vt.get(2,0)));
                        vt.set(2,1,(0 - vt.get(2,1)));
                        vt.set(2,2,(0 - vt.get(2,2)));

                        Matrix nv_transp = vt.transpose();
                        rot_nottrans = nv_transp.times(u_transp);
                        rot = rot_nottrans.transpose();

                }

                Matrix cb_tmp = centroidB.times(rot);
                tran = centroidA.minus(cb_tmp);


        }

        /** Calculate the RMS (root mean square) deviation of two sets of atoms.
         *
         * Atom sets must be pre-rotated.
         *
         * @param atomSet1 atom array 1
         * @param atomSet2 atom array 2
         * @return the RMS of two atom sets
         * @throws StructureException
         */
        public static double getRMS(Atom[] atomSet1, Atom[] atomSet2) throws StructureException {
                if ( atomSet1.length != atomSet2.length ){
                        throw new StructureException("The two atom sets are not of same length!");
                }

                double sum = 0.0;
                for ( int i =0 ; i < atomSet1.length;i++){
                        double d = Calc.getDistance(atomSet1[i],atomSet2[i]);
                        sum += (d*d);

                }

                double avd = ( sum/ atomSet1.length);
                //System.out.println("av dist" + avd);
                return Math.sqrt(avd);

        }

        /**
         * Calculate the TM-Score for the superposition.
         *
         * <em>Normalizes by the <strong>minimum</strong>-length structure (that is, {@code min\{len1,len2\}}).</em>
         *
         * Atom sets must be pre-rotated.
         *
         * <p>Citation:<br/>
         * <i>Zhang Y and Skolnick J (2004). "Scoring function for automated assessment
         * of protein structure template quality". Proteins 57: 702 - 710.</i>
         *
         * @param atomSet1 atom array 1
         * @param atomSet2 atom array 2
         * @param len1 The full length of the protein supplying atomSet1
         * @param len2 The full length of the protein supplying atomSet2
         * @return The TM-Score
         * @throws StructureException
         */
        public static double getTMScore(Atom[] atomSet1, Atom[] atomSet2, int len1, int len2) throws StructureException {
                if ( atomSet1.length != atomSet2.length ){
                        throw new StructureException("The two atom sets are not of same length!");
                }
                if ( atomSet1.length > len1 ){
                        throw new StructureException("len1 must be greater or equal to the alignment length!");
                }
                if ( atomSet2.length > len2 ){
                        throw new StructureException("len2 must be greater or equal to the alignment length!");
                }

                int Lmin = Math.min(len1,len2);
                int Laln = atomSet1.length;

                double d0 = 1.24 * Math.cbrt(Lmin - 15.) - 1.8;
                double d0sq = d0*d0;

                double sum = 0;
                for(int i=0;i<Laln;i++) {
                        double d = Calc.getDistance(atomSet1[i],atomSet2[i]);
                        sum+= 1./(1+d*d/d0sq);
                }

                return sum/Lmin;
        }

        /**
         * Calculate the TM-Score for the superposition.
         *
         * <em>Normalizes by the <strong>maximum</strong>-length structure (that is, {@code max\{len1,len2\}}) rather than the minimum.</em>
         *
         * Atom sets must be pre-rotated.
         *
         * <p>Citation:<br/>
         * <i>Zhang Y and Skolnick J (2004). "Scoring function for automated assessment
         * of protein structure template quality". Proteins 57: 702 - 710.</i>
         *
         * @param atomSet1 atom array 1
         * @param atomSet2 atom array 2
         * @param len1 The full length of the protein supplying atomSet1
         * @param len2 The full length of the protein supplying atomSet2
         * @return The TM-Score
         * @throws StructureException
         * @see {@link #getTMScore(Atom[], Atom[], int, int)}, which normalizes by the minimum length
         */
        public static double getTMScoreAlternate(Atom[] atomSet1, Atom[] atomSet2, int len1, int len2) throws StructureException {
                if ( atomSet1.length != atomSet2.length ){
                        throw new StructureException("The two atom sets are not of same length!");
                }
                if ( atomSet1.length > len1 ){
                        throw new StructureException("len1 must be greater or equal to the alignment length!");
                }
                if ( atomSet2.length > len2 ){
                        throw new StructureException("len2 must be greater or equal to the alignment length!");
                }

                int Lmax = Math.max(len1,len2);
                int Laln = atomSet1.length;

                double d0 = 1.24 * Math.cbrt(Lmax - 15.) - 1.8;
                double d0sq = d0*d0;

                double sum = 0;
                for(int i=0;i<Laln;i++) {
                        double d = Calc.getDistance(atomSet1[i],atomSet2[i]);
                        sum+= 1./(1+d*d/d0sq);
                }

                return sum/Lmax;
        }

        /**  Get the Rotation matrix that is required to superimpose the two atom sets.
         *
         * @return a rotation matrix.
         */
        public Matrix getRotation(){
                return rot;
        }

        /** Get the shift vector.
         *
         * @return the shift vector
         */
        public Atom getTranslation(){

                Atom a = new AtomImpl();
                a.setX(tran.get(0,0));
                a.setY(tran.get(0,1));
                a.setZ(tran.get(0,2));
                return a;
        }

        public Matrix4d getTransformation() {
                return Calc.getTransformation(rot, tran);
        }

}