/*
 *                    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.geometry.CalcPoint;
import org.biojava.nbio.structure.geometry.UnitQuaternions;
import org.biojava.nbio.structure.symmetry.utils.PermutationGenerator;

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


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

	private Point3d[] originalCoords = null;
	private Point3d[] transformedCoords = null;
	private RotationGroup rotations = new RotationGroup();
	private Vector3d centroid = new Vector3d();
	private Matrix4d centroidInverse = new Matrix4d();
	private Set<List<Integer>> hashCodes = new HashSet<>();

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

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

	private void solve() {
		initialize();
		int n = subunits.getSubunitCount();
		PermutationGenerator g = new PermutationGenerator(n);

		// loop over all permutations
		while (g.hasMore()) {
			int[] perm = g.getNext();
			List<Integer> permutation = new ArrayList<>(perm.length);
			for (int j = 0; j < n; j++) {
				permutation.add(perm[j]);
			}

			if (! isValidPermutation(permutation)) {
				continue;
			}

			boolean newPermutation = evaluatePermutation(permutation);
			if (newPermutation) {
				completeRotationGroup();
			}

			if (rotations.getOrder() >= subunits.getSubunitCount()) {
				return;
			}
		}
	}

	/**
	 * 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 completeRotationGroup() {
		PermutationGroup g = new PermutationGroup();
		for (int i = 0; i < rotations.getOrder(); i++) {
			Rotation s = rotations.getRotation(i);
			g.addPermutation(s.getPermutation());
		}
		g.completeGroup();

//   	System.out.println("Completing rotation group from: " +symmetryOperations.getSymmetryOperationCount() + " to " + g.getPermutationCount());

		// the group is complete, nothing to do
		if (g.getOrder() == rotations.getOrder()) {
			return;
		}

//  	System.out.println("complete group: " +  rotations.getOrder() +"/" + g.getOrder());
		// try to complete the group
		for (int i = 0; i < g.getOrder(); i++) {
			List<Integer> permutation = g.getPermutation(i);
			if (isValidPermutation(permutation)) {
				  // perform permutation of subunits
				evaluatePermutation(permutation);
			}
		}
	}

	private boolean isValidPermutation(List<Integer> permutation) {
		if (permutation.isEmpty()) {
			return false;
		}

		// if this permutation is a duplicate, return false
		if (hashCodes.contains(permutation)) {
			return false;
		}

		// check if permutation is pseudosymmetric
		if (! isAllowedPermuation(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 isAllowedPermuation(List<Integer> permutation) {
		List<Integer> seqClusterId = subunits.getClusterIds();
		for (int i = 0; i < permutation.size(); i++) {
			int j = permutation.get(i);
			if (!Objects.equals(seqClusterId.get(i), seqClusterId.get(j))) {
				return false;
			}
		}
		return true;
	}

	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);

		// 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()) {
			// transform to original coordinate system
			combineWithTranslation(transformation);
			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 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();
		}
	}
}
