/*
 *                    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.axis;

import org.biojava.nbio.structure.geometry.CalcPoint;
import org.biojava.nbio.structure.geometry.MomentsOfInertia;
import org.biojava.nbio.structure.symmetry.core.Helix;
import org.biojava.nbio.structure.symmetry.core.HelixLayers;
import org.biojava.nbio.structure.symmetry.core.QuatSymmetryResults;
import org.biojava.nbio.structure.symmetry.core.QuatSymmetrySubunits;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import javax.vecmath.*;

import java.util.*;

public class HelixAxisAligner extends AxisAligner {

	private static final Logger logger = LoggerFactory
			.getLogger(HelixAxisAligner.class);


	private static final Vector3d Y_AXIS = new Vector3d(0,1,0);
	private static final Vector3d Z_AXIS = new Vector3d(0,0,1);

	private QuatSymmetrySubunits subunits = null;
	private HelixLayers helixLayers = null;

	private Matrix4d transformationMatrix = new Matrix4d();
	private Matrix4d reverseTransformationMatrix = new Matrix4d();
	private Vector3d referenceVector = new Vector3d();
	private Vector3d principalRotationVector = new Vector3d();
	private Vector3d[] principalAxesOfInertia = null;
	private List<List<Integer>> alignedOrbits = null;

	private Vector3d minBoundary = new Vector3d();
	private Vector3d maxBoundary = new Vector3d();
	private double xzRadiusMax = Double.MIN_VALUE;

	boolean modified = true;

	public HelixAxisAligner(QuatSymmetryResults results) {
		this.subunits = new QuatSymmetrySubunits(results.getSubunitClusters());
		this.helixLayers = results.getHelixLayers();
		if (subunits == null) {
			throw new IllegalArgumentException("HelixAxisTransformation: Subunits are null");
		} else if (helixLayers == null) {
			throw new IllegalArgumentException("HelixAxisTransformation: HelixLayers is null");
		} else if (subunits.getSubunitCount() == 0) {
			throw new IllegalArgumentException("HelixAxisTransformation: Subunits is empty");
		} else if (helixLayers.size() == 0) {
			throw new IllegalArgumentException("HelixAxisTransformation: HelixLayers is empty");
		}
	}


	/* (non-Javadoc)
	 * @see org.biojava.nbio.structure.quaternary.core.AxisAligner#getTransformation()
	 */
	@Override
	public String getSymmetry() {
		run();
		return "H";
	}

	/* (non-Javadoc)
	 * @see org.biojava.nbio.structure.quaternary.core.AxisAligner#getTransformation()
	 */
	@Override
	public Matrix4d getTransformation() {
		run();
		return transformationMatrix;
	}

	/* (non-Javadoc)
	 * @see org.biojava.nbio.structure.quaternary.core.AxisAligner#getRotationMatrix()
	 */
	@Override
	public Matrix3d getRotationMatrix() {
		run();
		Matrix3d m = new Matrix3d();
		transformationMatrix.getRotationScale(m);
		return m;
	}

	/* (non-Javadoc)
	 * @see org.biojava.nbio.structure.quaternary.core.AxisAligner#getReverseTransformation()
	 */
	@Override
	public Matrix4d getReverseTransformation() {
		run();
		return reverseTransformationMatrix;
	}

	/* (non-Javadoc)
	 * @see org.biojava.nbio.structure.quaternary.core.AxisAligner#getPrincipalRotationAxis()
	 */
	@Override
	public Vector3d getPrincipalRotationAxis() {
		run();
		return principalRotationVector;
	}

	/* (non-Javadoc)
	 * @see org.biojava.nbio.structure.quaternary.core.AxisAligner#getRotationReferenceAxis()
	 */
	@Override
	public Vector3d getRotationReferenceAxis() {
		run();
		return referenceVector;
	}

	/* (non-Javadoc)
	 * @see org.biojava.nbio.structure.quaternary.core.AxisAligner#getPrincipalAxesOfInertia()
	 */
	@Override
	public Vector3d[] getPrincipalAxesOfInertia() {
		run();
		return principalAxesOfInertia;
	}

	/* (non-Javadoc)
	 * @see org.biojava.nbio.structure.quaternary.core.AxisAligner#getDimension()
	 */
	@Override
	public Vector3d getDimension() {
		run();
		Vector3d dimension = new Vector3d();
		dimension.sub(maxBoundary, minBoundary);
		dimension.scale(0.5);
		return dimension;
	}

	/* (non-Javadoc)
	 * @see org.biojava.nbio.structure.quaternary.core.AxisAligner#getXYRadius()
	 */
	@Override
	public double getRadius() {
		run();
		return xzRadiusMax;
	}

	/* (non-Javadoc)
	 * @see org.biojava.nbio.structure.quaternary.core.AxisAligner#getGeometicCenterTransformation()
	 */
	@Override
	public Matrix4d getGeometicCenterTransformation() {
		run();

		Matrix4d geometricCentered = new Matrix4d(reverseTransformationMatrix);
		geometricCentered.setTranslation(new Vector3d(getGeometricCenter()));

		return geometricCentered;
	}

	/* (non-Javadoc)
	 * @see org.biojava.nbio.structure.quaternary.core.AxisAligner#getGeometricCenter()
	 */
	@Override
	public Point3d getGeometricCenter() {
		run();

		Point3d geometricCenter = new Point3d();
		Vector3d translation = new Vector3d();
//		reverseTransformationMatrix.get(translation);

		// TODO does this apply to the helic case?
		// calculate adjustment around z-axis and transform adjustment to
		//  original coordinate frame with the reverse transformation

//		Vector3d corr = new Vector3d(0,minBoundary.y+getDimension().y, 0);
//		reverseTransformationMatrix.transform(corr);
//		geometricCenter.set(corr);

		reverseTransformationMatrix.transform(translation);
		geometricCenter.add(translation);
		return geometricCenter;
	}

	/* (non-Javadoc)
	 * @see org.biojava.nbio.structure.quaternary.core.AxisAligner#getCentroid()
	 */
	@Override
	public Point3d getCentroid() {
		return new Point3d(subunits.getCentroid());
	}

	/* (non-Javadoc)
	 * @see org.biojava.nbio.structure.quaternary.core.AxisAligner#getSubunits()
	 */
	@Override
	public QuatSymmetrySubunits getSubunits() {
		return subunits;
	}

	public HelixLayers getHelixLayers() {
		run();
		return helixLayers;
	}

	/* (non-Javadoc)
	 * @see org.biojava.nbio.structure.quaternary.core.AxisAligner#getOrbits()
	 */
	@Override
	public List<List<Integer>> getOrbits() {
		run();
		return alignedOrbits;
	}

	/**
	 * @return
	 */

	private void run () {
		if (modified) {
			calcPrincipalRotationVector();
			calcPrincipalAxes();
			calcReferenceVector();
			calcTransformation();
			// orient helix along Y axis by rotating 90 degrees around X-axis
			transformationMatrix = reorientHelix(0);

			calcReverseTransformation();
			calcBoundaries();
			calcAlignedOrbits();
			calcCenterOfRotation();

			// orient helix along Y axis by rotating 90 degrees around X-axis
//			transformationMatrix = reorientHelix(0);
//			calcReverseTransformation();

			modified = false;
		}
	}

	public Point3d calcCenterOfRotation() {
		List<Integer> line = getLongestLayerLine();

		// can't determine center of rotation if there are only 2 points
		// TODO does this ever happen??
		if (line.size() < 3) {
			return subunits.getCentroid();
		}

		Point3d centerOfRotation = new Point3d();
		List<Point3d> centers = subunits.getOriginalCenters();

		// calculate helix mid points for each set of 3 adjacent subunits
		for (int i = 0; i < line.size()-2; i++) {
			Point3d p1 = new Point3d(centers.get(line.get(i)));
			Point3d p2 = new Point3d(centers.get(line.get(i+1)));
			Point3d p3 = new Point3d(centers.get(line.get(i+2)));
			transformationMatrix.transform(p1);
			transformationMatrix.transform(p2);
			transformationMatrix.transform(p3);
			centerOfRotation.add(getMidPoint(p1, p2, p3));
		}

		// average over all midpoints to find best center of rotation
		centerOfRotation.scale(1/(line.size()-2));
		// since helix is aligned along the y-axis, with an origin at y = 0, place the center of rotation there
		centerOfRotation.y = 0;
		// transform center of rotation to the original coordinate frame
		reverseTransformationMatrix.transform(centerOfRotation);
//		System.out.println("center of rotation: " + centerOfRotation);
		return centerOfRotation;
	}

	private List<Integer> getLongestLayerLine() {
		int len = 0;
		int index = 0;

		Helix helix = helixLayers.getByLargestContacts();
		List<List<Integer>> layerLines = helix.getLayerLines();
		for (int i = 0; i < layerLines.size(); i++) {
			if (layerLines.get(i).size() > len) {
				len = layerLines.get(i).size();
				index = i;
			}
		}
		return layerLines.get(index);
	}

	/**
	 * Return a midpoint of a helix, calculated from three positions
	 * of three adjacent subunit centers.
	 * @param p1 center of first subunit
	 * @param p2 center of second subunit
	 * @param p3 center of third subunit
	 * @return midpoint of helix
	 */
	private Point3d getMidPoint(Point3d p1, Point3d p2, Point3d p3) {
		Vector3d v1 = new Vector3d();
		v1.sub(p1, p2);
		Vector3d v2 = new Vector3d();
		v2.sub(p3, p2);
		Vector3d v3 = new Vector3d();
		v3.add(v1);
		v3.add(v2);
		v3.normalize();

		// calculat the total distance between to subunits
		double dTotal = v1.length();
		// calculate the rise along the y-axis. The helix axis is aligned with y-axis,
		// therfore, the rise between subunits is the y-distance
		double rise = p2.y - p1.y;
		// use phythagorean theoremm to calculate chord length between two subunit centers
		double chord = Math.sqrt(dTotal*dTotal - rise*rise);
//		System.out.println("Chord d: " + dTotal + " rise: " + rise + "chord: " + chord);
		double angle = helixLayers.getByLargestContacts().getAxisAngle().angle;

		// using the axis angle and the chord length, we can calculate the radius of the helix
		// http://en.wikipedia.org/wiki/Chord_%28geometry%29
		double radius = chord/Math.sin(angle/2)/2; // can this go to zero?
//		System.out.println("Radius: " + radius);

		// project the radius onto the vector that points toward the helix axis
		v3.scale(radius);
		v3.add(p2);
//		System.out.println("Angle: " + Math.toDegrees(helixLayers.getByLowestAngle().getAxisAngle().angle));
		Point3d cor = new Point3d(v3);
		return cor;
	}

	private Matrix4d reorientHelix(int index) {
		Matrix4d matrix = new Matrix4d();
		matrix.setIdentity();
		matrix.setRotation(new AxisAngle4d(1,0,0,Math.PI/2*(index+1)));
		matrix.mul(transformationMatrix);
	return matrix;
	}

	/**
	 * Returns a list of orbits (an orbit is a cyclic permutation of subunit indices that are related
	 * by a rotation around the principal rotation axis) ordered from the +z direction to the -z direction (z-depth).
	 * Within an orbit, subunit indices are ordered with a clockwise rotation around the z-axis.
	 * @return list of orbits ordered by z-depth
	 */
	private void calcAlignedOrbits() {
		Map<Double, List<Integer>> depthMap = new TreeMap<Double, List<Integer>>();
		double[] depth = getSubunitZDepth();
		alignedOrbits = calcOrbits();

		// calculate the mean depth of orbit along z-axis
		for (List<Integer> orbit: alignedOrbits) {
			// calculate the mean depth along z-axis for each orbit
			double meanDepth = 0;
			for (int subunit: orbit) {
				meanDepth += depth[subunit];
			}
			meanDepth /= orbit.size();

			if (depthMap.get(meanDepth) != null) {
				meanDepth += 0.01;
			}
//			System.out.println("calcAlignedOrbits: " + meanDepth + " orbit: " + orbit);
			depthMap.put(meanDepth, orbit);
		}

		// now fill orbits back into list ordered by depth
		alignedOrbits.clear();
		for (List<Integer> orbit: depthMap.values()) {
			// order subunit in a clockwise rotation around the z-axis
			/// starting at the 12 O-clock position (+y position)
			// TODO how should this be aligned??
	//		alignWithReferenceAxis(orbit);
			alignedOrbits.add(orbit);
		}
	}


	private void calcTransformation() {
		calcTransformationBySymmetryAxes();
		// make sure this value is zero. On Linux this value is 0.
		transformationMatrix.setElement(3, 3, 1.0);
	}

	private void calcReverseTransformation() {
		reverseTransformationMatrix.invert(transformationMatrix);
	}

	private void calcTransformationBySymmetryAxes() {
		Vector3d[] axisVectors = new Vector3d[2];
		axisVectors[0] = new Vector3d(principalRotationVector);
		axisVectors[1] = new Vector3d(referenceVector);

		//  y,z axis centered at the centroid of the subunits
		Vector3d[] referenceVectors = new Vector3d[2];
		referenceVectors[0] = new Vector3d(Z_AXIS);
		referenceVectors[1] = new Vector3d(Y_AXIS);

		transformationMatrix = alignAxes(axisVectors, referenceVectors);

		// combine with translation
		Matrix4d combined = new Matrix4d();
		combined.setIdentity();
		Vector3d trans = new Vector3d(subunits.getCentroid());
		trans.negate();
		combined.setTranslation(trans);
		transformationMatrix.mul(combined);

		// for helical geometry, set a canonical view for the Z direction
		calcZDirection();
	}

	/**
	 * Returns a transformation matrix that rotates refPoints to match
	 * coordPoints
	 * @param refPoints the points to be aligned
	 * @param referenceVectors
	 * @return
	 */
	private Matrix4d alignAxes(Vector3d[] axisVectors, Vector3d[] referenceVectors) {
		Matrix4d m1 = new Matrix4d();
		AxisAngle4d a = new AxisAngle4d();
		Vector3d axis = new Vector3d();

		// calculate rotation matrix to rotate refPoints[0] into coordPoints[0]
		Vector3d v1 = new Vector3d(axisVectors[0]);
		Vector3d v2 = new Vector3d(referenceVectors[0]);
		double dot = v1.dot(v2);
		if (Math.abs(dot) < 0.999) {
			axis.cross(v1,v2);
			axis.normalize();
			a.set(axis, v1.angle(v2));
			m1.set(a);
			// make sure matrix element m33 is 1.0. It's 0 on Linux.
			m1.setElement(3,  3, 1.0);
		} else if (dot > 0) {
			// parallel axis, nothing to do -> identity matrix
			m1.setIdentity();
		} else if (dot < 0) {
			// anti-parallel axis, flip around x-axis
			m1.set(flipX());
		}

		// apply transformation matrix to all refPoints
		m1.transform(axisVectors[0]);
		m1.transform(axisVectors[1]);

		// calculate rotation matrix to rotate refPoints[1] into coordPoints[1]
		v1 = new Vector3d(axisVectors[1]);
		v2 = new Vector3d(referenceVectors[1]);
		Matrix4d m2 = new Matrix4d();
		dot = v1.dot(v2);
		if (Math.abs(dot) < 0.999) {
			axis.cross(v1,v2);
			axis.normalize();
			a.set(axis, v1.angle(v2));
			m2.set(a);
			// make sure matrix element m33 is 1.0. It's 0 on Linux.
			m2.setElement(3,  3, 1.0);
		} else if (dot > 0) {
			// parallel axis, nothing to do -> identity matrix
			m2.setIdentity();
		} else if (dot < 0) {
			// anti-parallel axis, flip around z-axis
			m2.set(flipZ());
		}

		// apply transformation matrix to all refPoints
		m2.transform(axisVectors[0]);
		m2.transform(axisVectors[1]);

		// combine the two rotation matrices
		m2.mul(m1);

		// the RMSD should be close to zero
		Point3d[] axes = new Point3d[2];
		axes[0] = new Point3d(axisVectors[0]);
		axes[1] = new Point3d(axisVectors[1]);
		Point3d[] ref = new Point3d[2];
		ref[0] = new Point3d(referenceVectors[0]);
		ref[1] = new Point3d(referenceVectors[1]);
		if (CalcPoint.rmsd(axes, ref) > 0.1) {
			logger.warn("AxisTransformation: axes alignment is off. RMSD: "
					+ CalcPoint.rmsd(axes, ref));
		}

		return m2;
	}

	private void calcPrincipalAxes() {
		MomentsOfInertia moi = new MomentsOfInertia();

		for (Point3d[] list: subunits.getTraces()) {
			for (Point3d p: list) {
				moi.addPoint(p, 1.0);
			}
		}
		principalAxesOfInertia = moi.getPrincipalAxes();
	}

	/**
	 * Calculates the min and max boundaries of the structure after it has been
	 * transformed into its canonical orientation.
	 */
	private void calcBoundaries() {
		minBoundary.x = Double.MAX_VALUE;
		maxBoundary.x = Double.MIN_VALUE;
		minBoundary.y = Double.MAX_VALUE;
		maxBoundary.x = Double.MIN_VALUE;
		minBoundary.z = Double.MAX_VALUE;
		maxBoundary.z = Double.MIN_VALUE;
		xzRadiusMax = Double.MIN_VALUE;

		Point3d probe = new Point3d();

		for (Point3d[] list: subunits.getTraces()) {
			for (Point3d p: list) {
				probe.set(p);
				transformationMatrix.transform(probe);

				minBoundary.x = Math.min(minBoundary.x, probe.x);
				maxBoundary.x = Math.max(maxBoundary.x, probe.x);
				minBoundary.y = Math.min(minBoundary.y, probe.y);
				maxBoundary.y = Math.max(maxBoundary.y, probe.y);
				minBoundary.z = Math.min(minBoundary.z, probe.z);
				maxBoundary.z = Math.max(maxBoundary.z, probe.z);
				xzRadiusMax = Math.max(xzRadiusMax, Math.sqrt(probe.x*probe.x + probe.z * probe.z));
			}
		}
//		System.out.println("MinBoundary: " + minBoundary);
//		System.out.println("MaxBoundary: " + maxBoundary);
//		System.out.println("zxRadius: " + xzRadiusMax);
	}

	/*
	 * Modifies the rotation part of the transformation axis for
	 * a Cn symmetric complex, so that the narrower end faces the
	 * viewer, and the wider end faces away from the viewer. Example: 3LSV
	 */
	private void calcZDirection() {
		calcBoundaries();

		// if the longer part of the structure faces towards the back (-z direction),
		// rotate around y-axis so the longer part faces the viewer (+z direction)
		if (Math.abs(minBoundary.z) > Math.abs(maxBoundary.z)) {
			Matrix4d rot = flipY();
			rot.mul(transformationMatrix);
			transformationMatrix.set(rot);
		}
	}

	private double[] getSubunitZDepth() {
		int n = subunits.getSubunitCount();
		double[] depth = new double[n];
		Point3d probe = new Point3d();

		// transform subunit centers into z-aligned position and calculate
		// z-coordinates (depth) along the z-axis.
		for (int i = 0; i < n; i++) {
			Point3d p= subunits.getCenters().get(i);
			probe.set(p);
			transformationMatrix.transform(probe);
			depth[i] = probe.z;
		}
		return depth;
	}

	/**
	 * Returns a list of list of subunit ids that form an "orbit", i.e. they
	 * are transformed into each other during a rotation around the principal symmetry axis (z-axis)
	 * @return
	 */
	private List<List<Integer>> calcOrbits() {
		int n = subunits.getSubunitCount();

		List<List<Integer>> orbits = new ArrayList<List<Integer>>();
		for (int i = 0; i < n; i++) {
			orbits.add(Collections.singletonList(i));
		}

		return orbits;
	}

	/**
	 * Returns a vector along the principal rotation axis for the
	 * alignment of structures along the z-axis
	 * @return principal rotation vector
	 */
	private void calcPrincipalRotationVector() {
//		AxisAngle4d axisAngle = helixLayers.getByLowestAngle().getAxisAngle();
		AxisAngle4d axisAngle = helixLayers.getByLargestContacts().getAxisAngle();
		principalRotationVector = new Vector3d(axisAngle.x, axisAngle.y, axisAngle.z);
	}

	/**
	 * Returns a vector perpendicular to the principal rotation vector
	 * for the alignment of structures in the xy-plane
	 * @return reference vector
	 */
	private void calcReferenceVector() {
		referenceVector = getReferenceAxisCylic();

		if (referenceVector == null) {
			logger.warn("no reference vector found. Using y-axis.");
			referenceVector = new Vector3d(Y_AXIS);
		}
		// make sure reference vector is perpendicular principal roation vector
		referenceVector = orthogonalize(principalRotationVector, referenceVector);
	}

	private Vector3d orthogonalize(Vector3d vector1, Vector3d vector2) {
		double dot = vector1.dot(vector2);
		Vector3d ref = new Vector3d(vector2);
//		System.out.println("p.r: " + dot);
//		System.out.println("Orig refVector: " + referenceVector);
		if (dot < 0) {
			vector2.negate();
		}
		if (Math.abs(dot) < 0.00001) {
			logger.info("HelixAxisAligner: reference axis parallel");
		}
		vector2.cross(vector1, vector2);
//		System.out.println("Intermed. refVector: " + vector2);
		vector2.normalize();
//		referenceVector.cross(referenceVector, principalRotationVector);
		vector2.cross(vector1, vector2);
		vector2.normalize();
		if (ref.dot(vector2) < 0) {
			vector2.negate();
		}
//		System.out.println("Mod. refVector: " + vector2);
		return vector2;
	}
	/**
	 * Returns the default reference vector for the alignment of Cn structures
	 * @return
	 */
	private Vector3d getReferenceAxisCylic() {
		// get principal axis vector that is perpendicular to the principal
		// rotation vector
		Vector3d vmin = null;
		double dotMin = 1.0;
		for (Vector3d v: principalAxesOfInertia) {
			if (Math.abs(principalRotationVector.dot(v)) < dotMin) {
				dotMin = Math.abs(principalRotationVector.dot(v));
				vmin = new Vector3d(v);
			}
		}
		if (principalRotationVector.dot(vmin) < 0) {
			vmin.negate();
		}

		return vmin;
	}

	private static Matrix4d flipX() {
		Matrix4d rot = new Matrix4d();
		rot.m00 = 1;
		rot.m11 = -1;
		rot.m22 = -1;
		rot.m33 = 1;
		return rot;
	}

	private static Matrix4d flipY() {
		Matrix4d rot = new Matrix4d();
		rot.m00 = -1;
		rot.m11 = 1;
		rot.m22 = -1;
		rot.m33 = 1;
		return rot;
	}

	private static Matrix4d flipZ() {
		Matrix4d rot = new Matrix4d();
		rot.m00 = -1;
		rot.m11 = -1;
		rot.m22 = 1;
		rot.m33 = 1;
		return rot;
	}

}
