/*
 *                    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/
 *
 */

package org.biojava.nbio.structure.symmetry.core;

import org.biojava.nbio.structure.symmetry.geometry.DistanceBox;
import org.biojava.nbio.structure.symmetry.geometry.MomentsOfInertia;
import org.biojava.nbio.structure.symmetry.geometry.SphereSampler;
import org.biojava.nbio.structure.symmetry.geometry.SuperPosition;

import javax.vecmath.AxisAngle4d;
import javax.vecmath.Matrix4d;
import javax.vecmath.Point3d;
import javax.vecmath.Vector3d;
import java.util.ArrayList;
import java.util.HashSet;
import java.util.List;
import java.util.Set;


/**
 *
 * @author Peter
 */
public class RotationSolver implements QuatSymmetrySolver {
    private Subunits subunits = null;
    private QuatSymmetryParameters parameters = null;

    private double distanceThreshold = 0.0f;
    private DistanceBox<Integer> box = null;
    private Vector3d centroid = new Vector3d();
    private Matrix4d centroidInverse = new Matrix4d();
    private Point3d[] originalCoords = null;
    private Point3d[] transformedCoords = null;
    private Set<List<Integer>> hashCodes = new HashSet<List<Integer>>();

    private RotationGroup rotations = new RotationGroup();

    public RotationSolver(Subunits subunits, QuatSymmetryParameters parameters) {
    	if (subunits.getSubunitCount()== 2) {
    		throw new IllegalArgumentException("RotationSolver cannot be applied to subunits with 2 centers");
    	}
        this.subunits = subunits;
        this.parameters = parameters;
    }

	public RotationGroup getSymmetryOperations() {
		if (rotations.getOrder() == 0) {
            solve();
            completeRotationGroup();
            rotations.complete();
        }
               
        return rotations;
    }

    private void solve() {	
        initialize();
        
        int maxSymOps = subunits.getSubunitCount();
        
        boolean isSpherical = isSpherical();
        
        // for cases with icosahedral symmetry n cannot be higher than 60, should check for spherical symmetry here
        // isSpherical check added 08-04-11
        if (maxSymOps % 60 == 0 && isSpherical) {
            maxSymOps = 60;
         }

        AxisAngle4d sphereAngle = new AxisAngle4d();
        Matrix4d transformation = new Matrix4d();

        int n = subunits.getSubunitCount();
        
        int sphereCount = SphereSampler.getSphereCount();
        
        List<Double> angles = getAngles();

       for (int i = 0; i < sphereCount; i++) {
            SphereSampler.getAxisAngle(i, sphereAngle);
            
            for (double angle : angles) {
                // apply rotation
                sphereAngle.angle = angle;
                transformation.set(sphereAngle);
                // Make sure matrix element m33 is 1.0. It's not on Linux.
                transformation.setElement(3, 3, 1.0);
                for (int j = 0; j < n; j++) {
                    transformedCoords[j].set(originalCoords[j]);
                    transformation.transform(transformedCoords[j]);
                }

                // get permutation of subunits and check validity/uniqueness             
                List<Integer> permutation = getPermutation();
  //              System.out.println("Rotation Solver: permutation: " + i + ": " + permutation);
                
                boolean isValidPermuation = isValidPermutation(permutation);
                if (! isValidPermuation) {
                    continue;
                }
               
                boolean newPermutation = evaluatePermutation(permutation);
                if (newPermutation) {
                	completeRotationGroup();
                }
                
                // check if all symmetry operations have been found.          
                if (rotations.getOrder() >= maxSymOps) {           	
                	return;
                }
            }
        }
    }
    
    private void completeRotationGroup() {
    	PermutationGroup g = new PermutationGroup();
    	for (int i = 0; i < rotations.getOrder(); i++) {
    		Rotation s = rotations.getRotation(i);
    		g.addPermutation(s.getPermutation());
    	}
    	g.completeGroup();
    	
    	// the group is complete, nothing to do
    	if (g.getOrder() == rotations.getOrder()) {
    		return;
    	}
    	   	
    	// try to complete the group
    	for (int i = 0; i < g.getOrder(); i++) {
    		List<Integer> permutation = g.getPermutation(i);
    		
    		boolean isValidPermutation = isValidPermutation(permutation);	
    		if (isValidPermutation) {
    			
    			  // perform permutation of subunits
                evaluatePermutation(permutation);
    		}
    	}
    }

	private boolean evaluatePermutation(List<Integer> permutation) {
		// permutate subunits
		for (int j = 0, n = subunits.getSubunitCount(); j < n; j++) {
		    transformedCoords[j].set(originalCoords[permutation.get(j)]);
		}

		int fold = PermutationGroup.getOrder(permutation);
		
		// get optimal transformation and axisangle by superimposing subunits
		AxisAngle4d axisAngle = new AxisAngle4d();
		Matrix4d transformation = SuperPosition.superposeAtOrigin(transformedCoords, originalCoords, axisAngle);
		double subunitRmsd 		= SuperPosition.rmsd(transformedCoords, originalCoords);
		
		if (subunitRmsd < parameters.getRmsdThreshold()) {
			combineWithTranslation(transformation);
			
			// evaluate superposition of CA traces
			QuatSymmetryScores scores = QuatSuperpositionScorer.calcScores(subunits, transformation, permutation);
			if (scores.getRmsd() < 0.0 || scores.getRmsd() > parameters.getRmsdThreshold()) {
				return false;
			}
			
			scores.setRmsdCenters(subunitRmsd);
			Rotation symmetryOperation = createSymmetryOperation(permutation, transformation, axisAngle, fold, scores);
			rotations.addRotation(symmetryOperation);
			return true;
		}
		return false;
	}

    private List<Double> getAngles() {
        int n = subunits.getSubunitCount();
        // for spherical symmetric cases, n cannot be higher than 60
        if (n % 60 == 0 && isSpherical()) {
           n = 60;
        }
        List<Integer> folds = subunits.getFolds();
        List<Double> angles = new ArrayList<Double>(folds.size()-1);

        // note this loop starts at 1, we do ignore 1-fold symmetry, which is the first entry
        for (int fold: folds) {
        	if (fold > 0 && fold <= n) {
        		angles.add(2* Math.PI/fold);
        	}
        }
        return angles;
    }
    
    private boolean isSpherical() {
    	MomentsOfInertia m = subunits.getMomentsOfInertia();
    	return m.getSymmetryClass(0.05) == MomentsOfInertia.SymmetryClass.SYMMETRIC;
    }

    private boolean isValidPermutation(List<Integer> permutation) {
    	 if (permutation.size() == 0) {
             return false;
         }
    	 
    	 // if this permutation is a duplicate, return false
    	if (hashCodes.contains(permutation)) {
    		return false;
    	}
       
        // check if permutation is allowed
        if (! isAllowedPermutation(permutation)) {
        	return false;
        }
        
        // get fold and make sure there is only one E (fold=1) permutation
        int fold = PermutationGroup.getOrder(permutation);
        if (rotations.getOrder() > 1 && fold == 1) {
            return false;
        }
        
        if (fold == 0 || subunits.getSubunitCount() % fold != 0) {
        	return false;
        }
        
        // if this permutation is a duplicate, returns false
        return hashCodes.add(permutation);
    }

    private boolean isAllowedPermutation(List<Integer> permutation) {
    	List<Integer> seqClusterId = subunits.getSequenceClusterIds();
    	for (int i = 0; i < permutation.size(); i++) {
    		int j = permutation.get(i);
    		if (seqClusterId.get(i) != seqClusterId.get(j)) {
    			return false;
    		}
    	}
    	return true;
    }
    /**
     * Adds translational component to rotation matrix
     * @param rotTrans
     * @param rotation
     * @return
     */
    private void combineWithTranslation(Matrix4d rotation) {
        rotation.setTranslation(centroid);
        rotation.mul(rotation, centroidInverse);
    }

    private Rotation createSymmetryOperation(List<Integer> permutation, Matrix4d transformation, AxisAngle4d axisAngle, int fold, QuatSymmetryScores scores) {
        Rotation s = new Rotation();
        s.setPermutation(new ArrayList<Integer>(permutation));
        s.setTransformation(new Matrix4d(transformation));
        s.setAxisAngle(new AxisAngle4d(axisAngle));
        s.setFold(fold);
        s.setScores(scores);

        return s;
    }

    private void setupDistanceBox() {
        distanceThreshold = calcDistanceThreshold();
        box = new DistanceBox<Integer>(distanceThreshold);

        for (int i = 0; i < originalCoords.length; i++) {
            box.addPoint(originalCoords[i], i);
        }
    }

    private double calcDistanceThreshold() {
        double threshold = Double.MAX_VALUE;
        int n = subunits.getSubunitCount();
        List<Point3d> centers = subunits.getCenters();
        
        for (int i = 0; i < n - 1; i++) {
            Point3d pi = centers.get(i);
            for (int j = i + 1; j < n; j++) {
                Point3d pj = centers.get(j);
                threshold = Math.min(threshold, pi.distanceSquared(pj));
            }
        }
        double distanceThreshold = Math.sqrt(threshold);

        distanceThreshold = Math.max(distanceThreshold, parameters.getRmsdThreshold());
        
        return distanceThreshold;
    }

    private List<Integer> getPermutation() {
        List<Integer> permutation = new ArrayList<Integer>(transformedCoords.length);
        double sum = 0.0f;

        for (Point3d t: transformedCoords) {
            List<Integer> neighbors = box.getNeighborsWithCache(t);
            int closest = -1;
            double minDist = Double.MAX_VALUE;

           for (int j : neighbors) {
            	double dist = t.distanceSquared(originalCoords[j]);
                if (dist < minDist) {
                    closest = j;
                    minDist = dist;
                } 
            }
            
            sum += minDist;
            if (closest == -1) {
         	   break;
            }
            permutation.add(closest);
        }
        double rmsd = Math.sqrt(sum / transformedCoords.length);

        if (rmsd > distanceThreshold || permutation.size() != transformedCoords.length) {
            permutation.clear();
            return permutation;
        }

        // check uniqueness of indices
        Set<Integer> set = new HashSet<Integer>(permutation);
        
        // if size mismatch, clear permutation (its invalid)
        if (set.size() != originalCoords.length) {
            permutation.clear();
        }

        return permutation;
    }

    private void initialize() {        
        // translation to centered coordinate system
        centroid = new Vector3d(subunits.getCentroid());
        // translation back to original coordinate system
        Vector3d reverse = new Vector3d(centroid);
        reverse.negate();
        centroidInverse.set(reverse);
        // Make sure matrix element m33 is 1.0. An old version vecmath did not set this element.
        centroidInverse.setElement(3, 3,  1.0);

        List<Point3d> centers = subunits.getCenters();
        int n = subunits.getSubunitCount();

        originalCoords = new Point3d[n];
        transformedCoords = new Point3d[n];

        for (int i = 0; i < n; i++) {
            originalCoords[i] = centers.get(i);
            transformedCoords[i] = new Point3d();
        }

        setupDistanceBox();
    }
}
