/*
 *                    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.Structure;
import org.biojava.nbio.structure.symmetry.utils.CombinationGenerator;
import org.jgrapht.UndirectedGraph;
import org.jgrapht.alg.ConnectivityInspector;
import org.jgrapht.graph.DefaultEdge;
import org.jgrapht.graph.UndirectedSubgraph;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import javax.vecmath.Point3d;

import java.math.BigInteger;
import java.util.*;

/**
 * Detects global and local quaternary protein structure symmetry in a structure.
 *
 * The QuatSymmetryParameter settings affect the calculated results.
 *
 * @author Peter Rose
 *
 */
public class QuatSymmetryDetector {

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

        private Structure structure = null;
        private QuatSymmetryParameters parameters = null;

        private List<QuatSymmetryResults> globalSymmetry = new ArrayList<QuatSymmetryResults>();
        private List<List<QuatSymmetryResults>> localSymmetries = new ArrayList<List<QuatSymmetryResults>>();
        private int proteinChainCount = 0;
        private boolean complete = false;

        public QuatSymmetryDetector(Structure structure, QuatSymmetryParameters parameters) {
                this.structure = structure;
                this.parameters = parameters;
        }

        /**
         * Returns true if structure contains protein subunits. The other methods
         * in this class will only return data if protein subunits are present.
         * Always use this method first, before retrieving global or local symmetry
         * results.
         *
         * @return true if protein subunits are present
         */
        public boolean hasProteinSubunits() {
                run();
                return proteinChainCount > 0;
        }

        /**
         * Returns list of global quaternary structure symmetry results
         *
         * @return list of global quaternary structure symmetry results
         */
        public List<QuatSymmetryResults> getGlobalSymmetry() {
                run();
                return globalSymmetry;
        }

        /**
         * Returns a list of lists of local quaternary structure symmetry results
         *
         * @return list of lists of local quaternary structure symmetry results
         */
        public List<List<QuatSymmetryResults>> getLocalSymmetries() {
                run();
                return localSymmetries;
        }

        private void run() {
                if (complete) {
                        return;
                }
                complete = true;
                //Cluster chains by sequence
                ClusterProteinChains clusterer = new ClusterProteinChains(structure, parameters);
                proteinChainCount = clusterer.getProteinChainCount();

                if (! hasProteinSubunits()) {
                        return;
                }

                int nucleicAcidChainCount = clusterer.getNucleicAcidChainCount();

                // sort seq. identity thresholds from smallest to largest. This reduces the total number of calculations necessary.
                double[] thresholds = parameters.getSequenceIdentityThresholds().clone();
                Arrays.sort(thresholds);

                for (int index = 0; index < thresholds.length; index++) {
                        // Map structure info to the sequence cluster
                        ChainClusterer chainClusterer = new ChainClusterer(clusterer.getSequenceAlignmentClusters(thresholds[index]));

                        // determine global symmetry
                        Subunits globalSubunits = createGlobalSubunits(chainClusterer, nucleicAcidChainCount);
                        QuatSymmetryResults gSymmetry = calcQuatSymmetry(globalSubunits);
                        gSymmetry.setSequenceIdentityThreshold(thresholds[index]);
                        globalSymmetry.add(gSymmetry);
//                      SymmetryDeviation sd = new SymmetryDeviation(globalSubunits, gSymmetry.getRotationGroup());


                        // determine local symmetry if global structure is
                        // (1) asymmetric (C1)
                        // (2) heteromeric (belongs to more than 1 sequence cluster)
                        // (3) more than 2 chains (heteromers with just 2 chains cannot have local symmetry)

                        // TODO example 2PT7: global C2, but local C6 symm., should that be included here ...?
                        // i.e., include all heteromers here, for example if higher symmetry is possible by stoichiometry? A6B2 -> local A6  can have higher symmetry
                        if (parameters.isLocalSymmetry() && globalSubunits.getSubunitCount() <= parameters.getMaximumLocalSubunits()) {
                                if (gSymmetry.getSymmetry().equals("C1") && proteinChainCount > 2) {
                                        List<QuatSymmetryResults> lSymmetry = new ArrayList<QuatSymmetryResults>();

                                        long start = System.nanoTime();

                                        for (Subunits subunits: createLocalSubunits(chainClusterer)) {
                                                QuatSymmetryResults result = calcQuatSymmetry(subunits);
                                                addToLocalSymmetry(result, lSymmetry);

                                                double time = (System.nanoTime()- start)/1000000000;
                                                if (time > parameters.getLocalTimeLimit()) {
                                                        logger.warn("Exceeded time limit for local symmetry calculations: " + time +
                                                                        " seconds. Quat symmetry results may be incomplete");
                                                        break;
                                                }
                                        }
                                        localSymmetries.add(lSymmetry);
                                }
                        }

                        if (! gSymmetry.getSubunits().isPseudoStoichiometric()) {
                                break;
                        }
                }


                trimGlobalSymmetryResults();
                trimLocalSymmetryResults();
                setPseudoSymmetry();
                setPreferredResults();
        }


        /**
         * trims asymmetric global symmetry results that are C1 and have pseudoStoichiometry
         */
        private void trimGlobalSymmetryResults() {
                for (Iterator<QuatSymmetryResults> iter = globalSymmetry.iterator(); iter.hasNext();) {
                        QuatSymmetryResults result = iter.next();
                        if (result.getSymmetry().equals("C1") && result.getSubunits().isPseudoStoichiometric()) {
                                iter.remove();
                        }
                }
        }

        /**
         * trims local symmetry results if any global symmetry is found. This only happens in special cases.
         */
        private void trimLocalSymmetryResults() {
                boolean hasGlobalSymmetry = false;
                for (QuatSymmetryResults result: globalSymmetry) {
                        if (! result.getSymmetry().equals("C1")) {
                                hasGlobalSymmetry = true;
                                break;
                        }
                }

                if (hasGlobalSymmetry) {
                        localSymmetries.clear();
                }
        }

        /**
         * Sets pseudosymmetry flag for results that have pseudosymmetry
         * @param thresholds sequence identity thresholds
         */
        private void setPseudoSymmetry() {
                setPseudoSymmetry(globalSymmetry);
                for (List<QuatSymmetryResults> localSymmetry: localSymmetries) {
                        setPseudoSymmetry(localSymmetry);
                }
        }

        /**
         * Sets pseudosymmetry based on the analysis of all symmetry results
         */
        private void setPseudoSymmetry(List<QuatSymmetryResults> results) {
                // find maximum symmetry order for cases of non-pseudostoichiometry
                int maxOrder = 0;
                for (QuatSymmetryResults result: results) {
                        if (result.getRotationGroup() != null && !result.getSubunits().isPseudoStoichiometric()) {
                                if (result.getRotationGroup().getOrder() > maxOrder) {
                                        maxOrder = result.getRotationGroup().getOrder();
                                }
                        }
                }

                // if the order of symmetry for a pseudstoichiometric case is higher
                // than the order for a non-pseudostoichiometry, set the pseudoSymmetry flag to true (i.e. 4HHB)
                for (QuatSymmetryResults result: results) {
                        if (result.getRotationGroup() != null) {
                                if (result.getRotationGroup().getOrder() > maxOrder) {
                                        result.getSubunits().setPseudoSymmetric(true);
                                }
                        }
                }
        }

        /**
         * Sets preferred results flag for symmetry result that should be shown by default in visualization programs
         * @param thresholds sequence identity thresholds
         */
        private void setPreferredResults() {
                int[] score = new int[globalSymmetry.size()];

                // check global symmetry
                for (int i = 0; i < globalSymmetry.size(); i++) {
                        QuatSymmetryResults result = globalSymmetry.get(i);
                        if (! result.getSymmetry().equals("C1")) {
                                score[i] += 2;
                        }
                        if (! result.getSubunits().isPseudoStoichiometric()) {
                                score[i]++;
                        }
                }

                int bestGlobal = 0;
                int bestScore = 0;
                for (int i = 0; i < score.length; i++) {
                        if (score[i] > bestScore) {
                                bestScore = score[i];
                                bestGlobal = i;
                        }
                }
                if (bestScore >= 2) {
                        QuatSymmetryResults g = globalSymmetry.get(bestGlobal);
                        g.setPreferredResult(true);
                        return;
                }

                // check local symmetry
                score = new int[localSymmetries.size()];

                for (int i = 0; i < localSymmetries.size(); i++) {
                        List<QuatSymmetryResults> results = localSymmetries.get(i);
                        for (QuatSymmetryResults result: results) {
                                if (! result.getSymmetry().equals("C1")) {
                                        score[i] += 2;
                                }
                                if (! result.getSubunits().isPseudoStoichiometric()) {
                                        score[i]++;
                                }
                        }
                }

                int bestLocal = 0;
                bestScore = 0;
                for (int i = 0; i < score.length; i++) {
                        if (score[i] > bestScore) {
                                bestScore = score[i];
                                bestLocal = i;
                        }
                }
                if (bestScore > 0) {
                        List<QuatSymmetryResults> results = localSymmetries.get(bestLocal);
                        for (QuatSymmetryResults result: results) {
                                result.setPreferredResult(true);
                        }
                } else {
                        QuatSymmetryResults g = globalSymmetry.get(bestGlobal);
                        g.setPreferredResult(true);
                }
        }

        private void addToLocalSymmetry(QuatSymmetryResults testResults, List<QuatSymmetryResults> localSymmetry) {
                if (testResults.getSymmetry().equals("C1")) {
                        return;
                }

                for (QuatSymmetryResults results: localSymmetry) {
                        if (results.getSubunits().overlaps(testResults.getSubunits())) {
                                return;
                        }
                }
                testResults.setLocal(true);
                localSymmetry.add(testResults);
        }

        private Subunits createGlobalSubunits(ChainClusterer chainClusterer, int nucleicAcidChainCount) {
                Subunits subunits = new Subunits(chainClusterer.getCalphaCoordinates(),
                                chainClusterer.getSequenceClusterIds(),
                                chainClusterer.getPseudoStoichiometry(),
                                chainClusterer.getMinSequenceIdentity(),
                                chainClusterer.getMaxSequenceIdentity(),
                                chainClusterer.getFolds(),
                                chainClusterer.getChainIds(),
                                chainClusterer.getModelNumbers());
                subunits.setNucleicAcidChainCount(nucleicAcidChainCount);
                return subunits;
        }

        private List<Subunits> createLocalSubunits(ChainClusterer chainClusterer) {
                List<Subunits> subunits = new ArrayList<Subunits>();
                List<List<Integer>> subClusters = decomposeClusters(chainClusterer.getCalphaCoordinates(), chainClusterer.getSequenceClusterIds());
                for (List<Integer> subCluster: subClusters) {
                        subunits.add(createLocalSubunit(subCluster, chainClusterer));
                }
                return subunits;
        }

        private Subunits createLocalSubunit(List<Integer> subCluster, ChainClusterer chainClusterer) {
              List<Point3d[]> subCalphaCoordinates = new ArrayList<Point3d[]>(subCluster.size());
              List<Integer> subSequenceIds = new ArrayList<Integer>(subCluster.size());
              List<Boolean> subPseudoStoichiometry = new ArrayList<Boolean>(subCluster.size());
              List<Double> subMinSequenceIdentity = new ArrayList<Double>(subCluster.size());
              List<Double> subMaxSequenceIdentity = new ArrayList<Double>(subCluster.size());
              List<String> subChainIds = new ArrayList<String>(subCluster.size());
              List<Integer> subModelNumbers = new ArrayList<Integer>(subCluster.size());

              for (int index: subCluster) {
                  subCalphaCoordinates.add(chainClusterer.getCalphaCoordinates().get(index));
                  subSequenceIds.add(chainClusterer.getSequenceClusterIds().get(index));
                  subPseudoStoichiometry.add(chainClusterer.getPseudoStoichiometry().get(index));
                  subMinSequenceIdentity.add(chainClusterer.getMinSequenceIdentity().get(index));
                  subMaxSequenceIdentity.add(chainClusterer.getMaxSequenceIdentity().get(index));
                  subChainIds.add(chainClusterer.getChainIds().get(index));
                  subModelNumbers.add(chainClusterer.getModelNumbers().get(index));
              }

              standardizeSequenceIds(subSequenceIds);

              Integer[] array = subSequenceIds.toArray(new Integer[subSequenceIds.size()]);
              List<Integer> subFolds = getFolds(array, subSequenceIds.size());
              Subunits subunits = new Subunits(subCalphaCoordinates,
                                        subSequenceIds,
                                        subPseudoStoichiometry,
                                        subMinSequenceIdentity,
                                        subMaxSequenceIdentity,
                                subFolds,
                                        subChainIds,
                                        subModelNumbers);
                        return subunits;
        }

        /**
         * Resets list of arbitrary sequence ids into integer order: 0, 1, ...
         * @param subSequenceIds
         */
        private static void standardizeSequenceIds(List<Integer> subSequenceIds) {
                int count = 0;
              int current = subSequenceIds.get(0);
              for (int i = 0; i < subSequenceIds.size(); i++) {
                  if (subSequenceIds.get(i) > current) {
                          current = subSequenceIds.get(i);
                          count++;
                  }
                  subSequenceIds.set(i, count);
              }
        }

        private List<List<Integer>> decomposeClusters(List<Point3d[]> caCoords, List<Integer> clusterIds) {
                List<List<Integer>> subClusters = new ArrayList<List<Integer>>();

                int last = getLastMultiSubunit(clusterIds);
                List<Point3d[]> subList = caCoords;
                if (last < caCoords.size()) {
                        subList = caCoords.subList(0, last);
                } else {
                        last = caCoords.size();
                }

                SubunitGraph subunitGraph = new SubunitGraph(subList);
                UndirectedGraph<Integer, DefaultEdge> graph = subunitGraph.getProteinGraph();
                logger.debug("Graph: {}", graph.toString());

                for (int i = last; i > 1; i--) {
                        CombinationGenerator generator = new CombinationGenerator(last, i);
                        int[] indices = null;
                        Integer[] subCluster = new Integer[i];

                        // avoid combinatorial explosion, i.e. for 1FNT
                        BigInteger maxCombinations = BigInteger.valueOf(parameters.getMaximumLocalCombinations());
                        logger.debug("Number of combinations: {}", generator.getTotal());
                    if (generator.getTotal().compareTo(maxCombinations) > 0) {
                        logger.warn("Number of combinations exceeds limit for biounit with {} subunits in groups of {} subunits. Will not check local symmetry for them", last, i);
                        continue;
                    }

                        while (generator.hasNext()) {
                                indices = generator.getNext();


                                // only consider sub clusters that can have rotational symmetry based on the number of subunits
                                // TODO this however may exclude a few cases of helical symmetry. Checking all combinations for
                                // helical symmetry would be prohibitively slow.
                                for (int j = 0; j < indices.length; j++) {
                                        subCluster[j] = clusterIds.get(indices[j]);
                                }
                                List<Integer> folds = getFolds(subCluster, last);

                                if (folds.size() < 2) {
                                        continue;
                                }

                                List<Integer> subSet = new ArrayList<Integer>(indices.length);
                                for (int index: indices) {
                                        subSet.add(index);
                                }

                                // check if this subset of subunits interact with each other
                                UndirectedGraph<Integer, DefaultEdge> subGraph = new UndirectedSubgraph<Integer, DefaultEdge>(graph, new HashSet<Integer>(subSet), null);
                                if (isConnectedGraph(subGraph)) {
                                        subClusters.add(subSet);
                                        if (subClusters.size() > parameters.getMaximumLocalResults()) {
                                                return subClusters;
                                        }
                                }
                        }
                }

                return subClusters;
        }

        private static int getLastMultiSubunit(List<Integer> clusterIds) {
                for (int i = 0, n = clusterIds.size(); i < n; i++) {
                        if (i < n-2) {
                                if (clusterIds.get(i)!=clusterIds.get(i+1) &&
                                                clusterIds.get(i+1) != clusterIds.get(i+2)) {
                                        return i+1;
                                }
                        }
                        if (i == n-2) {
                                if (clusterIds.get(i)!=clusterIds.get(i+1)) {
                                        return i+1;
                                }
                        }
                }
                return clusterIds.size();
        }

        private static boolean isConnectedGraph(UndirectedGraph<Integer, DefaultEdge> graph) {
                ConnectivityInspector<Integer, DefaultEdge> inspector = new ConnectivityInspector<Integer, DefaultEdge>(graph);
                return inspector.isGraphConnected();
        }

        private static List<Integer> getFolds(Integer[] subCluster, int size) {
                List<Integer> denominators = new ArrayList<Integer>();
                int[] counts = new int[size];
                for (int element: subCluster) {
                        counts[element]++;
                }

                for (int d = 1; d <= subCluster.length; d++) {
                        int count = 0;
                        for (int i = 0; i < size; i++) {
                                if (counts[i] > 0 && (counts[i] % d == 0)) {
                                        count += counts[i];
                                }
                        }
                        if (count == subCluster.length) {
                                denominators.add(d);
                        }
                }

                Collections.sort(denominators);
                return denominators;
        }

        private QuatSymmetryResults calcQuatSymmetry(Subunits subunits){
                return calcQuatSymmetry(subunits, parameters);
        }

        public static QuatSymmetryResults calcQuatSymmetry(Subunits subunits,
                        QuatSymmetryParameters parameters) {
                if (subunits.getSubunitCount() == 0) {
                        return null;
                }

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

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

                // asymmetric structures cannot be pseudosymmetric
                String symmetry = results.getSymmetry();
                if (symmetry.equals("C1")) {
                        subunits.setPseudoSymmetric(false);
                }

                // 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 = "rottranslation";
                                        results = new QuatSymmetryResults(subunits, helixLayers, method);
                                }
                        }
                }

                return results;
        }
}