/*
 *                    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/
 *
 * Created on 2013-05-23
 *
 */
package org.biojava.nbio.structure.symmetry.core;

import org.biojava.nbio.structure.Calc;
import org.biojava.nbio.structure.Structure;
import org.biojava.nbio.structure.cluster.*;
import org.biojava.nbio.structure.contact.BoundingBox;
import org.biojava.nbio.structure.contact.Grid;
import org.jgrapht.graph.SimpleGraph;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import org.jgrapht.alg.clique.CliqueMinimalSeparatorDecomposition;
import org.jgrapht.graph.DefaultEdge;
import org.jgrapht.graph.AsSubgraph;
import org.jgrapht.Graph;

import javax.vecmath.Point3d;
import java.util.*;
import java.util.stream.Collectors;

/**
 * Detects the symmetry (global, pseudo, internal and local) of protein
 * structures.
 * <p>
 * The {@link SubunitClustererParameters} determine the subunit definition and
 * clustering, while the {@link QuatSymmetryParameters} determine the calculated
 * symmetry results (point group and axes).
 *
 * @author Peter Rose
 * @author Aleix Lafita
 *
 */
public class QuatSymmetryDetector {

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


	/**
	 * Maximal distance between Calpha atoms of residues of different subunits
	 * to establish a residue contact.
	 */
	private static final double CONTACT_GRAPH_DISTANCE_CUTOFF = 8;
	/**
	 * The minimal number of residue contacts between subunits to consider
	 * them connected and add an edge to the contact graph.
	 */
	private static final int CONTACT_GRAPH_MIN_CONTACTS = 5;

	/** Prevent instantiation **/
	private QuatSymmetryDetector() {
	}

	/**
	 * Calculate GLOBAL symmetry results. This means that all {@link Subunit}
	 * are included in the symmetry.
	 *
	 * @param structure
	 *            protein chains will be extracted as {@link Subunit}
	 * @param symmParams
	 *            quaternary symmetry parameters
	 * @param clusterParams
	 *            subunit clustering parameters
	 * @return GLOBAL quaternary structure symmetry results
	 */
	public static QuatSymmetryResults calcGlobalSymmetry(Structure structure,
			QuatSymmetryParameters symmParams,
			SubunitClustererParameters clusterParams) {
		Stoichiometry composition = SubunitClusterer.cluster(structure, clusterParams);
		return calcGlobalSymmetry(composition, symmParams);
	}

	/**
	 * Calculate GLOBAL symmetry results. This means that all {@link Subunit}
	 * are included in the symmetry.
	 *
	 * @param subunits
	 *            list of {@link Subunit}
	 * @param symmParams
	 *            quaternary symmetry parameters
	 * @param clusterParams
	 *            subunit clustering parameters
	 * @return GLOBAL quaternary structure symmetry results
	 */
	public static QuatSymmetryResults calcGlobalSymmetry(
			List<Subunit> subunits, QuatSymmetryParameters symmParams,
			SubunitClustererParameters clusterParams) {
		Stoichiometry composition = SubunitClusterer.cluster(subunits,
				clusterParams);
		return calcGlobalSymmetry(composition, symmParams);
	}

	/**
	 * Calculate GLOBAL symmetry results. This means that all {@link Subunit}
	 * are included in the symmetry.
	 *
	 * @param composition
	 *            {@link Stoichiometry} object that contains clustering results
	 * @param symmParams
	 *            quaternary symmetry parameters
	 * @return GLOBAL quaternary structure symmetry results
	 */
	public static QuatSymmetryResults calcGlobalSymmetry(
			Stoichiometry composition, QuatSymmetryParameters symmParams) {
		return calcQuatSymmetry(composition, symmParams);
	}


	/**
	 * Returns a List of LOCAL symmetry results. This means that a subset of the
	 * {@link SubunitCluster} is left out of the symmetry calculation. Each
	 * element of the List is one possible LOCAL symmetry result.
	 * <p>
	 * Determine local symmetry if global structure is: (1) asymmetric, C1; (2)
	 * heteromeric (belongs to more than 1 subunit cluster); (3) more than 2
	 * subunits (heteromers with just 2 chains cannot have local symmetry)
	 *
	 * @param structure
	 *            protein chains will be extracted as {@link Subunit}
	 * @param symmParams
	 *            quaternary symmetry parameters
	 * @param clusterParams
	 *            subunit clustering parameters
	 * @return List of LOCAL quaternary structure symmetry results. Empty if
	 *         none.
	 */
	public static List<QuatSymmetryResults> calcLocalSymmetries(
			Structure structure, QuatSymmetryParameters symmParams,
			SubunitClustererParameters clusterParams) {

		Stoichiometry composition = SubunitClusterer.cluster(structure, clusterParams);
		return calcLocalSymmetries(composition, symmParams);
	}

	/**
	 * Returns a List of LOCAL symmetry results. This means that a subset of the
	 * {@link SubunitCluster} is left out of the symmetry calculation. Each
	 * element of the List is one possible LOCAL symmetry result.
	 * <p>
	 * Determine local symmetry if global structure is: (1) asymmetric, C1; (2)
	 * heteromeric (belongs to more than 1 subunit cluster); (3) more than 2
	 * subunits (heteromers with just 2 chains cannot have local symmetry)
	 *
	 * @param subunits
	 *            list of {@link Subunit}
	 * @param symmParams
	 *            quaternary symmetry parameters
	 * @param clusterParams
	 *            subunit clustering parameters
	 * @return List of LOCAL quaternary structure symmetry results. Empty if
	 *         none.
	 */
	public static List<QuatSymmetryResults> calcLocalSymmetries(
			List<Subunit> subunits, QuatSymmetryParameters symmParams,
			SubunitClustererParameters clusterParams) {

		Stoichiometry composition = SubunitClusterer.cluster(subunits, clusterParams);
		return calcLocalSymmetries(composition, symmParams);
	}

	/**
	 * Returns a List of LOCAL symmetry results. This means that a subset of the
	 * {@link SubunitCluster} is left out of the symmetry calculation. Each
	 * element of the List is one possible LOCAL symmetry result.
	 * <p>
	 * Determine local symmetry if global structure is: (1) asymmetric, C1; (2)
	 * heteromeric (belongs to more than 1 subunit cluster); (3) more than 2
	 * subunits (heteromers with just 2 chains cannot have local symmetry)
	 *
	 * @param globalComposition
	 *            {@link Stoichiometry} object that contains global clustering results
	 * @param symmParams
	 *            quaternary symmetry parameters
	 * @return List of LOCAL quaternary structure symmetry results. Empty if
	 *         none.
	 */

	public static List<QuatSymmetryResults> calcLocalSymmetries(Stoichiometry globalComposition, QuatSymmetryParameters symmParams) {

		Set<Set<Integer>> knownCombinations = new HashSet<>();
		List<SubunitCluster> clusters = globalComposition.getClusters();
		//more than one subunit per cluster required for symmetry
		List<SubunitCluster> nontrivialClusters =
				clusters.stream().
					filter(cluster -> (cluster.size()>1)).
					collect(Collectors.toList());

		QuatSymmetrySubunits consideredSubunits = new QuatSymmetrySubunits(nontrivialClusters);
		if (consideredSubunits.getSubunitCount() < 2)
			return new ArrayList<>();

		Graph<Integer, DefaultEdge> graph = initContactGraph(nontrivialClusters);
		Stoichiometry nontrivialComposition = new Stoichiometry(nontrivialClusters,false);

		List<Integer> allSubunitIds = new ArrayList<>(graph.vertexSet());
		Collections.sort(allSubunitIds);
		List<Integer> allSubunitClusterIds = consideredSubunits.getClusterIds();

		// since clusters are rearranged and trimmed, we need a reference to the original data
		// to maintain consistent IDs of clusters and subunits across all solutions
		Map<Integer, List<Integer>> clusterIdToSubunitIds =
				allSubunitIds.stream().
					collect(Collectors.
						groupingBy(allSubunitClusterIds::get, Collectors.toList()));

		List<QuatSymmetryResults> redundantSymmetries = new ArrayList<>();
		// first, find symmetries for single clusters and their groups
		// grouping is done based on symmetries found (i.e., no exhaustive permutation search is performed)
		if (clusters.size()>1) {

			List<QuatSymmetryResults> clusterSymmetries =
					calcLocalSymmetriesCluster(nontrivialComposition, clusterIdToSubunitIds,symmParams, knownCombinations);
			redundantSymmetries.addAll(clusterSymmetries);
		}
		//find symmetries for groups based on connectivity of subunits
		// disregarding initial clustering
		List<QuatSymmetryResults> graphSymmetries = calcLocalSymmetriesGraph(nontrivialComposition,
																			allSubunitClusterIds,
																			clusterIdToSubunitIds,
																			symmParams,
																			knownCombinations,
																			graph);

		redundantSymmetries.addAll(graphSymmetries);

		// find symmetries which are not superseded by any other symmetry
		// e.g., we have global stoichiometry of A3B3C,
		// the local symmetries found are C3 with stoichiometries A3, B3, A3B3;
		// then output only A3B3.

		List<QuatSymmetryResults> outputSymmetries =
				redundantSymmetries.stream().
					filter(a -> redundantSymmetries.stream().
						noneMatch(b -> a!=b && a.isSupersededBy(b))).
						collect(Collectors.toList());

		if(symmParams.isLocalLimitsExceeded(knownCombinations)) {
			logger.warn("Exceeded calculation limits for local symmetry detection. The results may be incomplete.");
		}

		return outputSymmetries;
	}


	private static  Graph<Integer, DefaultEdge> initContactGraph(List<SubunitCluster> clusters){

		Graph<Integer, DefaultEdge> graph = new SimpleGraph<>(DefaultEdge.class);

		// extract Ca coords from every subunit of every cluster.
		// all subunit coords are used for contact evaluation,
		// not only the aligned equivalent residues
		List <Point3d[]> clusterSubunitCoords =
				clusters.stream().
					flatMap(c -> c.getSubunits().stream()).
						map(r -> Calc.atomsToPoints(r.getRepresentativeAtoms())).
						collect(Collectors.toList());

		for (int i = 0; i < clusterSubunitCoords.size(); i++) {
			graph.addVertex(i);
		}

		// pre-compute bounding boxes
		List<BoundingBox> boundingBoxes = new ArrayList<>();
		clusterSubunitCoords.forEach(c -> boundingBoxes.add(new BoundingBox(c)));

		for (int i = 0; i < clusterSubunitCoords.size() - 1; i++) {
			Point3d[] coords1 = clusterSubunitCoords.get(i);
			BoundingBox bb1 = boundingBoxes.get(i);

			for (int j = i + 1; j < clusterSubunitCoords.size(); j++) {
				Point3d[] coords2 = clusterSubunitCoords.get(j);
				BoundingBox bb2 = boundingBoxes.get(j);
				Grid grid = new Grid(CONTACT_GRAPH_DISTANCE_CUTOFF);
				grid.addCoords(coords1, bb1, coords2, bb2);

				if (grid.getIndicesContacts().size() >= CONTACT_GRAPH_MIN_CONTACTS) {
					graph.addEdge(i, j);
				}
			}
		}
		return graph;
	}

	private static List<QuatSymmetryResults> calcLocalSymmetriesCluster(Stoichiometry nontrivialComposition,
	                                                                    Map<Integer, List<Integer>> clusterIdToSubunitIds,
	                                                                    QuatSymmetryParameters symmParams,
	                                                                    Set<Set<Integer>> knownCombinations) {

		List<QuatSymmetryResults> clusterSymmetries = new ArrayList<>();

		// find solutions for single clusters first
		for (int i=0;i<nontrivialComposition.numberOfComponents();i++) {
			QuatSymmetryResults localResult =
					calcQuatSymmetry(nontrivialComposition.getComponent(i),symmParams);

			if(localResult!=null && !localResult.getSymmetry().equals("C1")) {
				localResult.setLocal(true);
				clusterSymmetries.add(localResult);
				Set<Integer> knownResult = new HashSet<>(clusterIdToSubunitIds.get(i));
				// since symmetry is found,
				// do not try graph decomposition of this set of subunits later
				knownCombinations.add(knownResult);
			}
		}

		// group clusters by symmetries found, in case they all share axes and have the same number of subunits
		Map<String, Map<Integer,List<QuatSymmetryResults>>> groupedSymmetries =
				clusterSymmetries.stream().
					collect(Collectors.
						groupingBy(QuatSymmetryResults::getSymmetry,Collectors.
							groupingBy(QuatSymmetryResults::getSubunitCount,Collectors.toList())));

		for (Map<Integer,List<QuatSymmetryResults>> symmetriesByGroup: groupedSymmetries.values()) {
			for (List<QuatSymmetryResults> symmetriesBySubunits: symmetriesByGroup.values()) {

				Stoichiometry groupComposition =
						symmetriesBySubunits.stream().
							map(QuatSymmetryResults::getStoichiometry).
								reduce(Stoichiometry::combineWith).get();

				if (groupComposition.numberOfComponents() < 2) {
					continue;
				}
				//check if grouped clusters also have symmetry
				QuatSymmetryResults localResult = calcQuatSymmetry(groupComposition,symmParams);

				if(localResult!=null && !localResult.getSymmetry().equals("C1")) {
					localResult.setLocal(true);
					clusterSymmetries.add(localResult);
					// find subunit ids in this cluster list
					Set<Integer> knownResult = new HashSet<>();
					for (SubunitCluster cluster: groupComposition.getClusters()) {
						int i = nontrivialComposition.getClusters().indexOf(cluster);
						knownResult.addAll(clusterIdToSubunitIds.get(i));
					}
					// since symmetry is found,
					// do not try graph decomposition of this set of subunits later
					knownCombinations.add(knownResult);
				}
			}
		}
		return clusterSymmetries;
	}


	private static List<QuatSymmetryResults> calcLocalSymmetriesGraph(final Stoichiometry globalComposition,
	                                                                  final List<Integer> allSubunitClusterIds,
	                                                                  final Map<Integer, List<Integer>> clusterIdToSubunitIds,
	                                                                  QuatSymmetryParameters symmParams,
	                                                                  Set<Set<Integer>> knownCombinations,
	                                                                  Graph<Integer, DefaultEdge> graph) {

		List<QuatSymmetryResults> localSymmetries = new ArrayList<>();

		// do not go any deeper into recursion if over the time/combinations limit
		if(symmParams.isLocalLimitsExceeded(knownCombinations)) {
			return localSymmetries;
		}
		// extract components of a (sub-)graph
		CliqueMinimalSeparatorDecomposition<Integer, DefaultEdge> cmsd =
				new CliqueMinimalSeparatorDecomposition<>(graph);

		// only consider components with more than 1 vertex (subunit)
		Set<Set<Integer>> graphComponents =
				cmsd.getAtoms().stream().
					filter(component -> component.size()>1).
					collect(Collectors.toSet());

		//do not go into what has already been explored
		graphComponents.removeAll(knownCombinations);

		for (Set<Integer> graphComponent: graphComponents) {
			knownCombinations.add(graphComponent);

			List<Integer> usedSubunitIds = new ArrayList<>(graphComponent);
			Collections.sort(usedSubunitIds);
			// get clusters which contain only subunits in the current component
			Stoichiometry localStoichiometry =
					trimSubunitClusters(globalComposition, allSubunitClusterIds, clusterIdToSubunitIds, usedSubunitIds);

			if (localStoichiometry.numberOfComponents()==0) {
				continue;
			}

			//NB: usedSubunitIds might have changed when trimming clusters
			// if a subunit belongs to a cluster with no other subunits,
			// it is removed inside trimSubunitClusters
			Set<Integer> usedSubunitIdsSet = new HashSet<>(usedSubunitIds);
			if(!graphComponent.equals(usedSubunitIdsSet)) {
				if(knownCombinations.contains(usedSubunitIdsSet)) {
					continue;
				} else {
					knownCombinations.add(usedSubunitIdsSet);
				}
			}

			QuatSymmetryResults localResult = calcQuatSymmetry(localStoichiometry,symmParams);
			if(localResult!=null && !localResult.getSymmetry().equals("C1")) {
				localResult.setLocal(true);
				localSymmetries.add(localResult);
				continue;
			}

			if (usedSubunitIds.size() < 3) {
				// cannot decompose this component any further
				continue;
			}

			for (Integer removeSubunitId: usedSubunitIds) {
				// try removing subunits one by one and decompose the sub-graph recursively
				Set<Integer> prunedGraphVertices = new HashSet<>(usedSubunitIds);
				prunedGraphVertices.remove(removeSubunitId);
				if (knownCombinations.contains(prunedGraphVertices)) {
					continue;
				}
				knownCombinations.add(prunedGraphVertices);

				Graph<Integer, DefaultEdge> subGraph = new AsSubgraph<>(graph,prunedGraphVertices);

				List<QuatSymmetryResults> localSubSymmetries = calcLocalSymmetriesGraph(globalComposition,
																						allSubunitClusterIds,
																						clusterIdToSubunitIds,
																						symmParams,
																						knownCombinations,
																						subGraph);
				localSymmetries.addAll(localSubSymmetries);
			}

		}

		return localSymmetries;
	}

	private static Stoichiometry trimSubunitClusters(Stoichiometry globalComposition,
	                                                        List<Integer> allSubunitClusterIds,
	                                                        Map<Integer, List<Integer>> clusterIdToSubunitIds,
	                                                        List<Integer> usedSubunitIds) {
		List<SubunitCluster> globalClusters = globalComposition.getClusters();
		List<SubunitCluster> localClusters = new ArrayList<>();

		TreeSet<Integer> usedClusterIds =
				usedSubunitIds.stream().
					map(allSubunitClusterIds::get).
					distinct().
					collect(Collectors.toCollection(TreeSet::new));

		// for each used cluster, remove unused subunits
		for(Integer usedClusterId:usedClusterIds) {
			SubunitCluster originalCluster = globalClusters.get(usedClusterId);
			List<Integer> allSubunitIdsInCluster = clusterIdToSubunitIds.get(usedClusterId);

			//subunit numbering is global for the entire graph
			// make it zero-based for the inside of a cluster
			int minSUValue = Collections.min(allSubunitIdsInCluster);
			List<Integer> usedSubunitIdsInCluster = new ArrayList<>(allSubunitIdsInCluster);
			usedSubunitIdsInCluster.retainAll(usedSubunitIds);

			List<Integer> subunitsToRetain =
					usedSubunitIdsInCluster.stream().
						map(i -> i-minSUValue).
						collect(Collectors.toList());

			if (subunitsToRetain.size()>1) {
				SubunitCluster filteredCluster = new SubunitCluster(originalCluster, subunitsToRetain);
				localClusters.add(filteredCluster);
			} else {
				// if the cluster ends up having only 1 subunit, remove it from further processing
				usedSubunitIds.removeAll(usedSubunitIdsInCluster);
			}
		}
		return new Stoichiometry(localClusters,false);
	}


	private static QuatSymmetryResults calcQuatSymmetry(Stoichiometry composition, QuatSymmetryParameters parameters) {

		QuatSymmetrySubunits subunits = new QuatSymmetrySubunits(composition.getClusters());

		if (subunits.getSubunitCount() == 0)
			return null;

		RotationGroup rotationGroup = null;
		SymmetryPerceptionMethod method = null;
		if (subunits.getFolds().size() == 1) {
			// no symmetry possible, create empty ("C1") rotation group
			method = SymmetryPerceptionMethod.NO_ROTATION;
			rotationGroup = new RotationGroup();
			rotationGroup.setC1(subunits.getSubunitCount());
		} else if (subunits.getSubunitCount() == 2
				&& subunits.getFolds().contains(2)) {
			method = SymmetryPerceptionMethod.C2_ROTATION;
			QuatSymmetrySolver solver = new C2RotationSolver(subunits,
					parameters);
			rotationGroup = solver.getSymmetryOperations();
		} else {
			method = SymmetryPerceptionMethod.ROTATION;
			QuatSymmetrySolver solver = new RotationSolver(subunits, parameters);
			rotationGroup = solver.getSymmetryOperations();
		}

		QuatSymmetryResults results = new QuatSymmetryResults(composition,
				rotationGroup, method);

		String symmetry = results.getSymmetry();

		// Check structures with Cn symmetry (n = 1, ...) for helical symmetry
		if (symmetry.startsWith("C")) {
			HelixSolver hc = new HelixSolver(subunits,
					rotationGroup.getOrder(), parameters);
			HelixLayers helixLayers = hc.getSymmetryOperations();

			if (helixLayers.size() > 0) {
				// if the complex has no symmetry (c1) or has Cn (n>=2) symmetry
				// and the
				// helical symmetry has a lower RMSD than the cyclic symmetry,
				// set helical symmetry
				// If the RMSD for helical and cyclic symmetry is similar, a
				// slight preference is
				// given to the helical symmetry by the helixRmsdThreshold
				// parameter.
				double cRmsd = rotationGroup.getScores().getRmsd();
				double hRmsd = helixLayers.getScores().getRmsd();
				// System.out.println("cRMSD: " + cRmsd + " hRMSD: " + hRmsd);
				double deltaRmsd = hRmsd - cRmsd;
				if (symmetry.equals("C1")
						|| (!symmetry.equals("C1") && deltaRmsd <= parameters
								.getHelixRmsdThreshold())) {
					method = SymmetryPerceptionMethod.ROTO_TRANSLATION;
					results = new QuatSymmetryResults(composition, helixLayers,
							method);
				}
			}
		}

		return results;
	}
}
