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

import java.util.ArrayList;
import java.util.Arrays;
import java.util.Iterator;
import java.util.List;

import javax.vecmath.Matrix4d;

import org.biojava.nbio.structure.align.multiple.MultipleAlignment;
import org.biojava.nbio.structure.align.util.RotationAxis;
import org.biojava.nbio.structure.symmetry.internal.CESymmParameters.SymmetryType;

/**
 * Data Structure that stores all the symmetry axis that describe
 * the symmetry of a structure. Generalizes to all types of symmetry,
 * the classic ones (Cn, Dn) and any hierarchical or local symmetries.
 * <p>
 * Hierarchical symmetry can be visualized as a tree, where each level
 * has a fixed branching factor. Each level of the tree is associated
 * with a transformation operator, whose order determines the degree of
 * nodes at that level of the tree. Leaves of the tree implicitly
 * represent aligned repeats (indexed 0 to n-1), so care must be taken to
 * keep external references to the repeats (e.g. rows of a
 * {@link MultipleAlignment} in the same order implied by the tree.
 * <p>
 * Each node of the tree specifies an alignment between those repeats
 * below each of its children. It is also associated with a symmetry axis,
 * which is calculated based on the associated operator as well as any parent
 * operators.
 * It also stores the parts of the structure (symmetric units) involved
 * in each axis, in addition to the way to calculate them.
 * <p>
 * This is intended to provide a general axis support for the multiple
 * repeat alignment optimization and the axis display in Jmol. This
 * object is related to a MultipleAlignment object that defines the
 * symmetric units.
 *
 * @author Aleix Lafita
 * @since 4.2.0
 *
 */
public class SymmetryAxes {
	/*
	 * Implementation note: The tree is a nice explanation and a good image
	 * for developing algorithms, but it is not constructed explicitly.
	 * Instead, we just store one elementary axis for each level and reconstruct
	 * which operators apply to a particular leaf based on that leaf's index.
	 */

	/**
	 * Represents an axis of symmetry
	 * @author Spencer Bliven
	 *
	 */
	public static class Axis {
		private Matrix4d operator;
		private int order;
		private SymmetryType symmType;
		private int level;
		//private int indexInLevel;
		private int firstRepeat;
		private RotationAxis rotAxis;

		public Axis(Matrix4d operator, int order, SymmetryType type, int level, int firstRepeat) {
			if (order < 2) {
				throw new IllegalArgumentException("A symmetry axis should divide a structure in > 2 parts");
			}
			if(type != SymmetryType.OPEN && type != SymmetryType.CLOSED) {
				throw new IllegalArgumentException("Invalid symmetry type. Only OPEN and CLOSED are allowed");
			}

			this.operator = operator;
			this.order = order;
			this.symmType = type;
			setLevel(level);
			setFirstRepeat(firstRepeat);
			rotAxis = null;
		}
		/**
		 * Get the transformation operator for this axis as an homogeneous matrix
		 * @return the transformation operator
		 */
		public Matrix4d getOperator() {
			return operator;
		}
		public void setOperator(Matrix4d op) {
			this.operator = op;
		}
		/**
		 * Get the order of this axis (closed symm) or the number of repeats
		 * (open symm)
		 * @return the order
		 */
		public int getOrder() {
			return order;
		}
		/**
		 * @return the symmType (OPEN or CLOSED only)
		 */
		public SymmetryType getSymmType() {
			return symmType;
		}

		/**
		 * Get the transformation operator as a rotation axis. For open
		 * symmetry this will have a non-zero screw component.
		 * @return a RotationAxis for this Axis
		 */
		public RotationAxis getRotationAxis() {
			if( rotAxis == null) {
				rotAxis = new RotationAxis(operator);
			}
			return rotAxis;
		}
		/**
		 * @return The level of this axis within it's parent hierarchy, or -1 if unset
		 */
		public int getLevel() {
			return level;
		}
		/**
		 *
		 * @param level The level of this axis within it's parent hierarchy. Must be positive
		 */
		public void setLevel(int level) {
			if(level < 0) throw new IndexOutOfBoundsException("Level must be positive");
			this.level = level;
		}
//		/**
//		 * Each level can contain multiple equivalent axes. This index is
//		 * used to distinguish them.
//		 * @return the index of this axis relative to others at the same level
//		 */
//		public int getIndexInLevel() {
//			return indexInLevel;
//		}
//		/**
//		 *
//		 * @param indexInLevel the index of this axis relative to others at the same level
//		 */
//		public void setIndexInLevel(int indexInLevel) {
//			if( indexInLevel < 0 || getOrder() <= indexInLevel )
//				throw new IndexOutOfBoundsException("Invalid index for order "+getOrder());
//			this.indexInLevel = indexInLevel;
//		}
		/**
		 * Get the index of the first repeat used by this axis
		 * @return the firstRepeat
		 */
		public int getFirstRepeat() {
			return firstRepeat;
		}
		/**
		 * @param firstRepeat the index of the first repeat used by this axis
		 */
		public void setFirstRepeat(int firstRepeat) {
			this.firstRepeat = firstRepeat;
		}
	}

	/**
	 * List of all symmetry axis. They are sorted from higher to lower
	 * in the symmetry hierarchy, where higher means that they apply
	 * more globally and lower means that they apply to a local region
	 * of the higher axis division.
	 */
	private final List<Axis> axes;

	/**
	 * Constructor.
	 * Initializes variables only.
	 */
	public SymmetryAxes(){
		axes = new ArrayList<>();
	}

	/**
	 * Adds a new axis of symmetry to the bottom level of the tree
	 *
	 * @param axis the new axis of symmetry found
	 * @param order number of parts that this axis divides the structure in
	 * @param type indicates whether the axis has OPEN or CLOSED symmetry
	 */
	public void addAxis(Matrix4d axis, int order, SymmetryType type) {
		axes.add(new Axis(axis,order,type,axes.size(),0));
	}

	/**
	 * Return a list giving the number of times each axis must be applied
	 * to generate the given repeat.
	 * <P>
	 * For instance, for a D3 case <tt>getAxisCounts(4)</tt> would return [2,0],
	 * indicating that repeat 4 is generated by two applications of the 3-fold
	 * axis followed by 0 applications of the two-fold axis.
	 *
	 * @param repeat Index of the desired repeat
	 * @return array of the same length as axes giving the number of times
	 *  to apply each axis.
	 */
	private int[] getAxisCounts(int repeat) {
		int[] counts = new int[getNumLevels()];

		for(int i = counts.length-1; i >= 0; i--) {
			int d = axes.get(i).getOrder();
			counts[i] = repeat % d;
			repeat /= d;
		}
		assert repeat == 0 : "Invalid repeat index";
		return counts;
	}

//	/**
//	 * Inverse of {@link #getAxisCounts(int)}; Calculates the repeat for a
//	 * particular number of applications of each axis
//	 * @param counts Number of times to apply each axis
//	 * @return Repeat index
//	 */
//	private int getRepeatIndex(int[] counts) {
//		int repeat = 0;
//		for(int i = 0; i< counts.length; i++) {
//			repeat += counts[i]*axes.get(i).getOrder();
//		}
//		return repeat;
//	}
	/**
	 * Updates an axis of symmetry, after the superposition changed.
	 *
	 * @param index old axis index
	 * @param newAxis
	 */
	public void updateAxis(Integer index, Matrix4d newAxis){
		axes.get(index).setOperator(newAxis);
	}

	/**
	 * Return the operator for all elementary axes of symmetry of the structure, that is,
	 * the axes stored in the List as unique and from which all the symmetry
	 * axes are constructed.
	 *
	 * @return axes elementary axes of symmetry.
	 */
	public List<Matrix4d> getElementaryAxes(){
		List<Matrix4d> ops = new ArrayList<Matrix4d>(getNumLevels());
		for(Axis axis : axes) {
			ops.add(axis.getOperator());
		}
		return ops;
	}

	/**
	 * Return all elementary axes of symmetry of the structure, that is,
	 * the axes stored in the List as unique and from which all the symmetry
	 * axes are constructed.
	 *
	 * @return axes elementary axes of symmetry.
	 */
	public List<Axis> getElementaryAxesObjects() {
		return axes;
	}

	/**
	 * Get the indices of participating repeats in Cauchy two-line form.
	 * <p>
	 * Returns two lists of the same length.
	 * The first gives a list of all repeat indices which are aligned
	 * at the specified level of symmetry (e.g. 0 through the degree of this level).
	 * The second list gives the corresponding repeats after applying the
	 * operator once.
	 *
	 * @param level the axis index
	 * @return the double List of repeat relations, or null if the
	 * 			level is invalid
	 * @see #getRepeatsCyclicForm(int, int) for an equivalent specification with half the memory
	 */
	public List<List<Integer>> getRepeatRelation(int level){
		return getRepeatRelation(level,0);
	}

	public List<List<Integer>> getRepeatRelation(Axis axis){
		return getRepeatRelation(axis.getLevel(),axis.getFirstRepeat());
	}

	public List<List<Integer>> getRepeatRelation(int level, int firstRepeat) {
		Axis axis = axes.get(level);
		int m = getNumRepeats(level+1);//size of the children
		int d = axis.getOrder(); // degree of this node
		int n = m*d; // number of repeats included
		if(firstRepeat % n != 0)
			throw new IllegalArgumentException(String.format("Repeat %d cannot start a block at level %s of this tree",firstRepeat,level));
		if(axis.getSymmType() == SymmetryType.OPEN) {
			n -= m; // leave off last child for open symm
		}
		List<Integer> repeats = new ArrayList<>(n);
		List<Integer> equiv = new ArrayList<>(n);
		for(int i=0;i<n;i++) {
			repeats.add(i+firstRepeat);
			equiv.add( (i+m)%(m*d)+firstRepeat );
		}
		return Arrays.asList(repeats,equiv);

	}

	/**
	 * Get the indices of participating repeats in cyclic form.
	 * <p>
	 * Each inner list gives a set of equivalent repeats and should have length
	 * equal to the order of the axis' operator.
	 * @param level
	 * @param firstRepeat
	 * @return
	 */
	public List<List<Integer>> getRepeatsCyclicForm(int level, int firstRepeat) {
		Axis axis = axes.get(level);
		int m = getNumRepeats(level+1);//size of the children
		int d = axis.getOrder(); // degree of this node
		int n = m*d; // number of repeats included
		if(firstRepeat % n != 0) {
			throw new IllegalArgumentException(String.format("Repeat %d cannot start a block at level %s of this tree",firstRepeat,level));
		}
		if(axis.getSymmType() == SymmetryType.OPEN) {
			n -= m; // leave off last child for open symm
		}

		List<List<Integer>> repeats = new ArrayList<>(m);
		for(int i=0;i<m;i++) {
			List<Integer> cycle = new ArrayList<>(d);
			for(int j=0;j<d;j++) {
				cycle.add(firstRepeat+i+j*m);
			}
			repeats.add(cycle);
		}
		return repeats;
	}
	public List<List<Integer>> getRepeatsCyclicForm(Axis axis) {
		return getRepeatsCyclicForm(axis.getLevel(),axis.getFirstRepeat());
	}
	public List<List<Integer>> getRepeatsCyclicForm(int level) {
		return getRepeatsCyclicForm(level,0);
	}
	public String getRepeatsCyclicForm(Axis axis, List<?> repeats) {
		if(repeats.size() != getNumRepeats()) {
			throw new IllegalArgumentException("Mismatch in the number of repeats");
		}
		return getRepeatsCyclicForm(getRepeatsCyclicForm(axis), repeats);
	}
	public static String getRepeatsCyclicForm(List<List<Integer>> cycleForm, List<?> repeats) {
		StringBuilder str = new StringBuilder();
		for(List<Integer> cycle : cycleForm) {
			str.append("(");
			Iterator<Integer> cycleIt = cycle.iterator();
			str.append(repeats.get(cycleIt.next())); //should be at least one
			while(cycleIt.hasNext()) {
				str.append(";")
				.append(repeats.get( cycleIt.next() ));
			}
			str.append(")");
		}
		return str.toString();
	}

	/**
	 * Return the transformation that needs to be applied to a
	 * repeat in order to superimpose onto repeat 0.
	 *
	 * @param repeat the repeat index
	 * @return transformation matrix for the repeat
	 */
	public Matrix4d getRepeatTransform(int repeat){

		Matrix4d transform = new Matrix4d();
		transform.setIdentity();

		int[] counts = getAxisCounts(repeat);

		for(int t = counts.length-1; t>=0; t--) {
			if( counts[t] == 0 )
				continue;
			Matrix4d axis = new Matrix4d(axes.get(t).getOperator());
			for(int i=0;i<counts[t];i++) {
				transform.mul(axis);
			}
		}
		return transform;
	}

	/**
	 * Return the transformation that needs to be applied to
	 * repeat x in order to superimpose onto repeat y.
	 *
	 * @param x the first repeat index (transformed)
	 * @param y the second repeat index (fixed)
	 * @return transformation matrix for the repeat x
	 */
	public Matrix4d getRepeatTransform(int x, int y){

		Matrix4d transform = new Matrix4d();
		transform.setIdentity();

		int[] iCounts = getAxisCounts(x);
		int[] jCounts = getAxisCounts(y);

		int[] counts = new int[iCounts.length];
		for (int k = 0; k < iCounts.length; k++)
			counts[k] = iCounts[k] - jCounts[k];

		for(int t = counts.length-1; t>=0; t--) {
			if(counts[t] == 0)
				continue;
			if (counts[t] > 0) {
				Matrix4d axis = new Matrix4d(axes.get(t).getOperator());
				for(int i=0;i<counts[t];i++)
					transform.mul(axis);
			} else if (counts[t] < 0) {
				Matrix4d axis = new Matrix4d(axes.get(t).getOperator());
				axis.invert();
				for(int i=0;i<counts[t];i++)
					transform.mul(axis);
			}
		}
		return transform;
	}

	/**
	 * Return all symmetry axes of of the structure: the set of axes that
	 * describe all parts of the structure. This combines the elementary
	 * axes to generate all possible axes. The axes are returned in the repeat
	 * degrees.
	 * @return axes all symmetry axes of the structure.
	 */
	public List<Axis> getSymmetryAxes(){

		List<Axis> symmAxes = new ArrayList<>();

		Matrix4d prior = new Matrix4d();
		prior.setIdentity();

		getSymmetryAxes(symmAxes,prior,0,0);


		return symmAxes;
	}
	/**
	 * Recursive helper
	 * @param symmAxes output list
	 * @param prior transformation aligning the first repeat of this axis with the first overall
	 * @param level current level
	 */
	private void getSymmetryAxes(List<Axis> symmAxes, Matrix4d prior, int level, int firstRepeat) {
		if(level >= getNumLevels() ) {
			return;
		}

		Axis elem = axes.get(level);
		Matrix4d elemOp = elem.getOperator();

		// Current axis:
		// elementary maps B -> A
		// prior maps I -> A and J -> B
		// want J -> I = J -> B -> A <- I= inv(prior) * elementary * prior
		Matrix4d currAxisOp = new Matrix4d(prior);
		currAxisOp.invert();
		currAxisOp.mul(elemOp);
		currAxisOp.mul(prior);
		Axis currAxis = new Axis(currAxisOp,elem.getOrder(),elem.getSymmType(),level,firstRepeat);
		symmAxes.add(currAxis);

		//Remember that all degrees are at least 2
		getSymmetryAxes(symmAxes,prior,level+1,firstRepeat);
		//New prior is elementary^d*prior
		Matrix4d newPrior = new Matrix4d(elemOp);
		newPrior.mul(prior);
		int childSize = getNumRepeats(level+1);
		getSymmetryAxes(symmAxes,newPrior,level+1,firstRepeat+childSize);
		for(int d=2;d<elem.getOrder();d++) {
			newPrior.mul(elemOp,newPrior);
			getSymmetryAxes(symmAxes,newPrior,level+1,firstRepeat+childSize*d);
		}
	}


//	public Matrix4d getSymmetryAxis(int level, int axisNum) {
//		if(level == 0) {
//			if( axisNum != 0 )
//				throw new IndexOutOfBoundsException("Axis number out of bounds");
//			return axes.get(0);
//		} else {
//			if( axisNum >= degrees.get(level-1) )
//				throw new IndexOutOfBoundsException("Axis number out of bounds");
//			// Convert axisNum into a count of
//
//	}
	/**
	 * Get the number of repeats. This is equal to the product of all degrees.
	 * @return Number of repeats (leaves of the tree).
	 */
	public int getNumRepeats() {
		return getNumRepeats(0);
	}

	/**
	 * Get the number of leaves from a node at the specified level. This is
	 * equal to the product of all degrees at or below the level.
	 * @param level level of the tree to cut at
	 * @return Number of repeats (leaves of the tree).
	 */
	private int getNumRepeats(int level) {
		int size = 1;
		// Return 1 for illegally high level
		if(level < getNumLevels()) {
			for(Axis axis : axes.subList(level, getNumLevels())) {
				size *= axis.getOrder();
			}
		}
		return size;
	}

	/**
	 * Get the first repeat index of each axis of a specified level.
	 * @param level level of the tree to cut at
	 * @return List of first Repeats of each index, sorted in ascending order
	 */
	public List<Integer> getFirstRepeats(int level) {
		List<Integer> firstRepeats = new ArrayList<>();
		int m = getNumRepeats(level+1); //size of the level
		int d = axes.get(level).getOrder(); //degree of this level
		int n = m*d; // number of repeats included in each axis
		for (int firstRepeat = 0; firstRepeat < getNumRepeats(); firstRepeat+=n)
			firstRepeats.add(firstRepeat);
		return firstRepeats;
	}

	public Axis getElementaryAxis(int level) {
		return axes.get(level);
	}

	public int getNumLevels() {
		return axes.size();
	}

}
