/*
 *                    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 July 1, 2010
 * Author: Mark Chapman
 */

package org.biojava.nbio.alignment;

import org.biojava.nbio.core.alignment.SimpleProfile;
import org.biojava.nbio.alignment.template.GuideTreeNode;
import org.biojava.nbio.alignment.template.PairwiseSequenceScorer;
import org.biojava.nbio.core.alignment.template.Profile;
import org.biojava.nbio.core.alignment.template.ProfilePair;
import org.biojava.nbio.core.sequence.AccessionID;
import org.biojava.nbio.core.sequence.template.Compound;
import org.biojava.nbio.core.sequence.template.Sequence;
import org.biojava.nbio.phylo.ForesterWrapper;
import org.biojava.nbio.phylo.TreeConstructor;
import org.biojava.nbio.phylo.TreeConstructorType;
import org.forester.evoinference.matrix.distance.BasicSymmetricalDistanceMatrix;
import org.forester.phylogeny.Phylogeny;
import org.forester.phylogeny.PhylogenyNode;

import javax.swing.tree.TreeNode;

import java.util.*;
import java.util.concurrent.Future;

/**
 * Implements a data structure for a guide tree used during progressive multiple sequence alignment.  Leaf
 * {@link Node}s correspond to single {@link Sequence}s.  Internal {@link Node}s correspond to multiple sequence
 * alignments.  The root {@link Node} corresponds to the full multiple sequence alignment.
 *
 * @author Mark Chapman
 * @param <S> each {@link Sequence} in the tree is of type S
 * @param <C> each element of a {@link Sequence} is a {@link Compound} of type C
 */
public class GuideTree<S extends Sequence<C>, C extends Compound> implements Iterable<GuideTreeNode<S, C>> {

        private List<S> sequences;
        private List<PairwiseSequenceScorer<S, C>> scorers;
        private BasicSymmetricalDistanceMatrix distances;
        private String newick;
        private Node root;

        /**
         * Creates a guide tree for use during progressive multiple sequence alignment.
         *
         * @param sequences the {@link List} of {@link Sequence}s to align
         * @param scorers list of sequence pair scorers, one for each pair of sequences given
         */
        public GuideTree(List<S> sequences, List<PairwiseSequenceScorer<S, C>> scorers) {
                this.sequences = Collections.unmodifiableList(sequences);
                this.scorers = Collections.unmodifiableList(scorers);
                distances = new BasicSymmetricalDistanceMatrix(sequences.size());
                for (int i = 0, n = 0; i < sequences.size(); i++) {
                        AccessionID id = sequences.get(i).getAccession();
                        String str = (id == null) ? Integer.toString(i + 1) : id.getID();
                        distances.setIdentifier(i, str);
                        for (int j = i+1; j < sequences.size(); j++) {
                                double dist = scorers.get(n++).getDistance();
                                distances.setValue(i, j, dist);
                        }
                }
                BasicSymmetricalDistanceMatrix distclone = ForesterWrapper.cloneDM(distances);
                Phylogeny phylogeny = TreeConstructor.distanceTree(distclone, TreeConstructorType.NJ);
                newick = phylogeny.toString();
                root = new Node(phylogeny.getRoot(), null);
        }

        /**
         * Returns a sequence pair score for all {@link Sequence} pairs in the given {@link List}.
         *
         * @return list of sequence pair scores
         */
        public double[] getAllPairsScores() {
                double[] scores = new double[scorers.size()];
                int n = 0;
                for (PairwiseSequenceScorer<S, C> scorer : scorers) {
                        scores[n++] = scorer.getScore();
                }
                return scores;
        }

        /**
         * Returns the distance matrix used to construct this guide tree.  The scores have been normalized.
         *
         * @return the distance matrix used to construct this guide tree
         */
        public double[][] getDistanceMatrix() {
                double[][] matrix = new double[distances.getSize()][distances.getSize()];
                for (int i = 0; i < matrix.length; i++) {
                        for (int j = i+1; j < matrix.length; j++) {
                                matrix[i][j] = matrix[j][i] = distances.getValue(i, j);
                        }
                }
                return matrix;
        }

        /**
         * Returns the root {@link Node} which corresponds to the full multiple sequence alignment.
         *
         * @return the root node
         */
        public Node getRoot() {
                return root;
        }

        /**
         * Returns the similarity matrix used to construct this guide tree.  The scores have not been normalized.
         *
         * @return the similarity matrix used to construct this guide tree
         */
        public double[][] getScoreMatrix() {
                double[][] matrix = new double[sequences.size()][sequences.size()];
                for (int i = 0, n = 0; i < matrix.length; i++) {
                        matrix[i][i] = scorers.get(i).getMaxScore();
                        for (int j = i+1; j < matrix.length; j++) {
                                matrix[i][j] = matrix[j][i] = scorers.get(n++).getScore();
                        }
                }
                return matrix;
        }

        /**
         * Returns the {@link Sequence}s which make up the leaves of this tree.
         *
         * @return the sequences which make up the leaves of this tree
         */
        public List<S> getSequences() {
                return sequences;
        }

        // method for Iterable

        /**
         * Returns a post-order {@link Iterator} that traverses the tree from leaves to root.
         */
        @Override
        public Iterator<GuideTreeNode<S, C>> iterator() {
                return new PostOrderIterator();
        }

        // method from Object

        @Override
        public String toString() {
                return newick;
        }

        /**
         * Implements a data structure for the node in a guide tree used during progressive multiple sequence alignment.
         */
        public class Node implements GuideTreeNode<S, C> {

                private GuideTreeNode<S, C> parent, child1, child2;
                private double distance;
                private String name;
                private boolean isLeaf, isVisited;
                private Profile<S, C> profile;
                private Future<ProfilePair<S, C>> profileFuture;

                private Node(PhylogenyNode node, Node parent) {
                        this.parent = parent;
                        distance = node.getDistanceToParent();
                        name = node.getName();
                        if(isLeaf = node.isExternal()) {
                                profile = new SimpleProfile<S, C>(sequences.get(distances.getIndex(name)));
                        } else {
                                child1 = new Node(node.getChildNode1(), this);
                                child2 = new Node(node.getChildNode2(), this);
                        }
                }

                // methods for GuideTreeNode

                @Override
                public GuideTreeNode<S, C> getChild1() {
                        return child1;
                }

                @Override
                public GuideTreeNode<S, C> getChild2() {
                        return child2;
                }

                @Override
                public double getDistanceToParent() {
                        return distance;
                }

                @Override
                public String getName() {
                        return name;
                }

                @Override
                public Profile<S, C> getProfile() {
                        return profile;
                }

                @Override
                public Future<ProfilePair<S, C>> getProfileFuture() {
                        return profileFuture;
                }

                @Override
                public void setProfile(Profile<S, C> profile) {
                        this.profile = profile;
                        profileFuture = null;
                }

                @Override
                public void setProfileFuture(Future<ProfilePair<S, C>> profileFuture) {
                        this.profileFuture = profileFuture;
                        profile = null;
                }

                // methods for TreeNode

                @Override
                public Enumeration<GuideTreeNode<S, C>> children() {
                        Vector<GuideTreeNode<S, C>> children = new Vector<GuideTreeNode<S, C>>();
                        children.add(getChild1());
                        children.add(getChild2());
                        return children.elements();
                }

                @Override
                public boolean getAllowsChildren() {
                        return !isLeaf();
                }

                @Override
                public GuideTreeNode<S, C> getChildAt(int childIndex) {
                        if (childIndex == 1) {
                                return getChild1();
                        } else if (childIndex == 2) {
                                return getChild2();
                        }
                        throw new IndexOutOfBoundsException();
                }

                @Override
                public int getChildCount() {
                        return 2;
                }

                @Override
                public int getIndex(TreeNode child) {
                        return getChildAt(1) == child ? 1 : (getChildAt(2) == child ? 2 : -1);
                }

                @Override
                public GuideTreeNode<S, C> getParent() {
                        return parent;
                }

                @Override
                public boolean isLeaf() {
                        return isLeaf;
                }

                // helper methods for iterator

                private void clearVisited() {
                        isVisited = false;
                        if (!isLeaf()) {
                                ((Node) getChild1()).clearVisited();
                                ((Node) getChild2()).clearVisited();
                        }
                }

                private boolean isVisited() {
                        return isVisited;
                }

                private void visit() {
                        isVisited = true;
                }

        }

        // helper class that defines the default post-order (leaves to root) traversal
        private class PostOrderIterator implements Iterator<GuideTreeNode<S, C>> {

                private Stack<Node> nodes;

                private PostOrderIterator() {
                        getRoot().clearVisited();
                        nodes = new Stack<Node>();
                        nodes.push(getRoot());
                }

                // methods for Iterator

                @Override
                public boolean hasNext() {
                        return !nodes.isEmpty();
                }

                @Override
                public GuideTreeNode<S, C> next() {
                        if(!hasNext()){
                                throw new NoSuchElementException();
                        }

                        while (hasNext()) {
                                Node next = nodes.peek(), child1 = (Node) next.getChild1(), child2 = (Node) next.getChild2();
                                if (child1 != null && !child1.isVisited()) {
                                        nodes.push(child1);
                                } else if (child2 != null && !child2.isVisited()) {
                                        nodes.push(child2);
                                } else {
                                        next.visit();
                                        return nodes.pop();
                                }
                        }
                        return null;
                }

                @Override
                public void remove() {
                        throw new UnsupportedOperationException();
                }

        }

}