/*
 *                    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 javax.vecmath.AxisAngle4d;
import javax.vecmath.Matrix4d;
import javax.vecmath.Point3d;
import javax.vecmath.Quat4d;
import javax.vecmath.Vector3d;

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

import java.util.ArrayList;
import java.util.HashMap;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Objects;
import java.util.Set;


/**
 *
 * @author Peter
 */
public class RotationSolver implements QuatSymmetrySolver {
        private QuatSymmetrySubunits 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;
        // Cache whether a permutation is invalid (null) vs has been added to rotations
        private Map<List<Integer>,Rotation> evaluatedPermutations = new HashMap<>();

        private RotationGroup rotations = new RotationGroup();

        public RotationSolver(QuatSymmetrySubunits 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;
        }

        @Override
        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++) {
                        // Sampled orientation
                        //TODO The SphereSampler samples 4D orientation space. We really
                        // only need to sample 3D unit vectors, since we use limited
                        // angles. -SB
                        SphereSampler.getAxisAngle(i, sphereAngle);

                        // Each valid rotation angle
                        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);

                                // check if novel
                                if ( evaluatedPermutations.containsKey(permutation)) {
                                        continue; //either invalid or already added
                                }

                                Rotation newPermutation = isValidPermutation(permutation);
                                if (newPermutation != null) {
                                        completeRotationGroup(newPermutation);
                                }

                                // check if all symmetry operations have been found.
                                if (rotations.getOrder() >= maxSymOps) {
                                        return;
                                }
                        }
                }
        }

        /**
         * Combine current rotations to make all possible permutations.
         * If these are all valid, add them to the rotations
         * @param additionalRots Additional rotations we are considering adding to this.rotations
         * @return whether the rotations were valid and added
         */
        private boolean completeRotationGroup(Rotation... additionalRots) {
                PermutationGroup g = new PermutationGroup();
                for (Rotation s : rotations) {
                        g.addPermutation(s.getPermutation());
                }
                for( Rotation s : additionalRots) {
                        g.addPermutation(s.getPermutation());
                        // inputs should not have been added already
                        assert evaluatedPermutations.get(s.getPermutation()) == null;
                }
                g.completeGroup();

                // the group is complete, nothing to do
                if (g.getOrder() == rotations.getOrder()+additionalRots.length) {
                        for( Rotation s : additionalRots) {
                                addRotation(s);
                        }
                        return true;
                }

                // try to complete the group
                List<Rotation> newRots = new ArrayList<>(g.getOrder());
                // First, quick check for whether they're allowed
                for (List<Integer> permutation : g) {
                        if( evaluatedPermutations.containsKey(permutation)) {
                                Rotation rot = evaluatedPermutations.get(permutation);
                                if( rot == null ) {
                                        return false;
                                }
                        } else {
                                if( ! isAllowedPermutation(permutation)) {
                                        return false;
                                }
                        }
                }
                // Slower check including the superpositions
                for (List<Integer> permutation : g) {
                        Rotation rot;
                        if( evaluatedPermutations.containsKey(permutation)) {
                                rot = evaluatedPermutations.get(permutation);
                        } else {
                                rot = isValidPermutation(permutation);
                        }

                        if( rot == null ) {
                                // if any induced rotation is invalid, abort
                                return false;
                        }
                        if(!evaluatedPermutations.containsKey(permutation)){ //novel
                                newRots.add(rot);
                        }
                }
                // Add rotations
                for( Rotation rot : newRots) {
                        addRotation(rot);
                }
                return true;
        }

        private void addRotation(Rotation rot) {
                evaluatedPermutations.put(rot.getPermutation(),rot);
                rotations.addRotation(rot);
        }

        /**
         * Superimpose subunits based on the given permutation. Then check whether
         * the superposition passes RMSD thresholds and create a Rotation to
         * represent it if so.
         * @param permutation A list specifying which subunits should be aligned by the current transformation
         * @return A Rotation representing the permutation, or null if the superposition did not meet thresholds.
         */
        private Rotation superimposePermutation(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 subunit superposition
                // TODO implement this piece of code using at origin superposition
                Quat4d quat = UnitQuaternions.relativeOrientation(
                                originalCoords, transformedCoords);
                AxisAngle4d axisAngle = new AxisAngle4d();
                Matrix4d transformation = new Matrix4d();

                transformation.set(quat);
                axisAngle.set(quat);

                Vector3d axis = new Vector3d(axisAngle.x, axisAngle.y, axisAngle.z);
                if (axis.lengthSquared() < 1.0E-6) {
                        axisAngle.x = 0;
                        axisAngle.y = 0;
                        axisAngle.z = 1;
                        axisAngle.angle = 0;
                } else {
                        axis.normalize();
                        axisAngle.x = axis.x;
                        axisAngle.y = axis.y;
                        axisAngle.z = axis.z;
                }

                CalcPoint.transform(transformation, transformedCoords);
                double subunitRmsd = CalcPoint.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 null;
                        }

                        scores.setRmsdCenters(subunitRmsd);
                        Rotation symmetryOperation = createSymmetryOperation(permutation, transformation, axisAngle, fold, scores);
                        return symmetryOperation;
                }
                return null;
        }

        /**
         * Get valid rotation angles given the number of subunits
         * @return The rotation angle corresponding to each fold of {@link Subunits#getFolds()}
         */
        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<>(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;
        }

        /**
         * Checks if a particular permutation is allowed and superimposes well.
         * Caches results.
         * @param permutation
         * @return null if invalid, or a rotation if valid
         */
        private Rotation isValidPermutation(List<Integer> permutation) {
                if (permutation.size() == 0) {
                        return null;
                }

                // cached value
                if (evaluatedPermutations.containsKey(permutation)) {
                        return evaluatedPermutations.get(permutation);
                }

                // check if permutation is allowed
                if (! isAllowedPermutation(permutation)) {
                        evaluatedPermutations.put(permutation, null);
                        return null;
                }

                // check if superimposes
                Rotation rot = superimposePermutation(permutation);
                return rot;
        }

        /**
         * The permutation must map all subunits onto an equivalent subunit
         * and no subunit onto itself
         * @param permutation
         * @return
         */
        private boolean isAllowedPermutation(List<Integer> permutation) {
                List<Integer> seqClusterId = subunits.getClusterIds();
                int selfaligned = 0;
                for (int i = 0; i < permutation.size(); i++) {
                        int j = permutation.get(i);
                        if (!Objects.equals(seqClusterId.get(i), seqClusterId.get(j))) {
                                return false;
                        }
                        if(i == j ) {
                                selfaligned++;
                        }
                }
                // either identity (all self aligned) or all unique
                return selfaligned == 0 || selfaligned == permutation.size();
        }

        /**
         * Adds translational component to rotation matrix
         * @param rotation
         * @return
         */
        private void combineWithTranslation(Matrix4d rotation) {
                rotation.setTranslation(centroid);
                rotation.mul(rotation, centroidInverse);
        }

        private static 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<>(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;
        }

        /**
         * Compare this.transformedCoords with the original coords. For each
         * subunit, return the transformed subunit with the closest position.
         * @return A list mapping each subunit to the closest transformed subunit
         */
        private List<Integer> getPermutation() {
                List<Integer> permutation = new ArrayList<>(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<>(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();
        }
}