package prefuse.data.util;

import java.util.Iterator;

import prefuse.Constants;
import prefuse.data.Edge;
import prefuse.data.Node;
import prefuse.data.Tuple;
import prefuse.util.collections.Queue;

/**
 * Provides a distance-limited breadth first traversal over nodes, edges,
 * or both, using any number of traversal "roots".
 *  
 * @author <a href="http://jheer.org">jeffrey heer</a>
 */
public class BreadthFirstIterator implements Iterator {

    protected Queue m_queue = new Queue();
    protected int   m_depth;
    protected int   m_traversal;
    protected boolean m_includeNodes;
    protected boolean m_includeEdges;
    
    /**
     * Create an uninitialized BreadthFirstIterator. Use the
     * {@link #init(Object, int, int)} method to initialize the iterator.
     */
    public BreadthFirstIterator() {
        // do nothing, requires init call
    }
    
    /**
     * Create a new BreadthFirstIterator starting from the given source node.
     * @param n the source node from which to begin the traversal
     * @param depth the maximum graph distance to traverse
     * @param traversal the traversal type, one of
     * {@link prefuse.Constants#NODE_TRAVERSAL},
     * {@link prefuse.Constants#EDGE_TRAVERSAL}, or
     * {@link prefuse.Constants#NODE_AND_EDGE_TRAVERSAL}
     */
    public BreadthFirstIterator(Node n, int depth, int traversal) {
        init(new Node[] {n}, depth, traversal);
    }
    
    /**
     * Create a new BreadthFirstIterator starting from the given source nodes.
     * @param it an Iterator over the source nodes from which to begin the
     * traversal
     * @param depth the maximum graph distance to traverse
     * @param traversal the traversal type, one of
     * {@link prefuse.Constants#NODE_TRAVERSAL},
     * {@link prefuse.Constants#EDGE_TRAVERSAL}, or
     * {@link prefuse.Constants#NODE_AND_EDGE_TRAVERSAL}
     */
    public BreadthFirstIterator(Iterator it, int depth, int traversal) {
        init(it, depth, traversal);
    }
    
    /**
     * Initialize (or re-initialize) this iterator.
     * @param o Either a source node or iterator over source nodes
     * @param depth the maximum graph distance to traverse
     * @param traversal the traversal type, one of
     * {@link prefuse.Constants#NODE_TRAVERSAL},
     * {@link prefuse.Constants#EDGE_TRAVERSAL}, or
     * {@link prefuse.Constants#NODE_AND_EDGE_TRAVERSAL}
     */
    public void init(Object o, int depth, int traversal) {
        // initialize the member variables
        m_queue.clear();
        m_depth = depth;
        if ( traversal < 0 || traversal >= Constants.TRAVERSAL_COUNT )
            throw new IllegalArgumentException(
                    "Unrecognized traversal type: "+traversal);
        m_traversal = traversal;
        m_includeNodes = (traversal == Constants.NODE_TRAVERSAL || 
                traversal == Constants.NODE_AND_EDGE_TRAVERSAL);
        m_includeEdges = (traversal == Constants.EDGE_TRAVERSAL ||
                traversal == Constants.NODE_AND_EDGE_TRAVERSAL);
        
        // seed the queue
        // TODO: clean this up? (use generalized iterator?)
        if ( m_includeNodes ) {
            if ( o instanceof Node ) {
                m_queue.add(o, 0);
            } else {
                Iterator tuples = (Iterator)o;
                while ( tuples.hasNext() )
                    m_queue.add(tuples.next(), 0);
            }
        } else {
            if ( o instanceof Node ) {
                Node n = (Node)o;
                m_queue.visit(n, 0);
                Iterator edges = getEdges(n);
                while ( edges.hasNext() ) {
                    Edge e = (Edge)edges.next();
                    Node nn = e.getAdjacentNode(n);
                    m_queue.visit(nn, 1);
                    if ( m_queue.getDepth(e) < 0 )
                        m_queue.add(e, 1);
                }
            } else {
                Iterator tuples = (Iterator)o;
                while ( tuples.hasNext() ) {
                    // TODO: graceful error handling when non-node in set?
                    Node n = (Node)tuples.next();
                    m_queue.visit(n, 0);
                    Iterator edges = getEdges(n);
                    while ( edges.hasNext() ) {
                        Edge e = (Edge)edges.next();
                        Node nn = e.getAdjacentNode(n);
                        m_queue.visit(nn, 1);
                        if ( m_queue.getDepth(e) < 0 )
                            m_queue.add(e, 1);
                    }
                }
            }
        }
    }
    
    // ------------------------------------------------------------------------
    
    /**
     * @see java.util.Iterator#remove()
     */
    public void remove() {
        throw new UnsupportedOperationException();
    }

    /**
     * @see java.util.Iterator#hasNext()
     */
    public boolean hasNext() {
        return !m_queue.isEmpty();
    }

    /**
     * Determines which edges are traversed for a given node.
     * @param n a node
     * @return an iterator over edges incident on the node
     */
    protected Iterator getEdges(Node n) {
        return n.edges(); // TODO: add support for all edges, in links only, out links only
    }
    
    /**
     * Get the traversal depth at which a particular tuple was encountered.
     * @param t the tuple to lookup
     * @return the traversal depth of the tuple, or -1 if the tuple has not
     * been visited by the traversal.
     */
    public int getDepth(Tuple t) {
        return m_queue.getDepth(t);
    }
    
    /**
     * @see java.util.Iterator#next()
     */
    public Object next() {        
        Tuple t = (Tuple)m_queue.removeFirst();

        switch ( m_traversal ) {
        
        case Constants.NODE_TRAVERSAL:
        case Constants.NODE_AND_EDGE_TRAVERSAL:
            for ( ; true; t = (Tuple)m_queue.removeFirst() ) {
                if ( t instanceof Edge ) {
                    return t;
                } else {
                    Node n = (Node)t;
                    int d = m_queue.getDepth(n);
                    
                    if ( d < m_depth ) {
                        int dd = d+1;
                        Iterator edges = getEdges(n);
                        while ( edges.hasNext() ) {
                            Edge e = (Edge)edges.next();
                            Node v = e.getAdjacentNode(n);
                        
                            if ( m_includeEdges && m_queue.getDepth(e) < 0 )
                                m_queue.add(e, dd);
                            if ( m_queue.getDepth(v) < 0 )
                                m_queue.add(v, dd);
                        }
                    }
                    else if ( m_includeEdges && d == m_depth )
                    {
                        Iterator edges = getEdges(n);
                        while ( edges.hasNext() ) {
                            Edge e = (Edge)edges.next();
                            Node v = e.getAdjacentNode(n);
                            int dv = m_queue.getDepth(v);
                            if ( dv > 0 && m_queue.getDepth(e) < 0 ) {
                                m_queue.add(e, Math.min(d,dv));
                            }
                        }
                    }
                    return n;
                }
            }
                
        case Constants.EDGE_TRAVERSAL:
            Edge e = (Edge)t;
            Node u = e.getSourceNode();
            Node v = e.getTargetNode();
            int du = m_queue.getDepth(u);
            int dv = m_queue.getDepth(v);

            if ( du != dv ) {
                Node n = (dv > du ? v : u);
                int  d = Math.max(du, dv);
            
                if ( d < m_depth ) {
                    int dd = d+1;
                    Iterator edges = getEdges(n);
                    while ( edges.hasNext() ) {
                        Edge ee = (Edge)edges.next();
                        if ( m_queue.getDepth(ee) >= 0 )
                            continue; // already visited
            
                        Node nn = ee.getAdjacentNode(n);
                        m_queue.visit(nn, dd);
                        m_queue.add(ee, dd);
                    }
                }
            }
            return e;
        
        default:
            throw new IllegalStateException();
        }
    }

} // end of class BreadthFirstIterator