/* Copyright (c) Jython Developers */
package org.python.modules;

import java.util.ArrayList;
import java.util.List;

import org.python.core.ArgParser;
import org.python.core.ClassDictInit;
import org.python.core.Py;
import org.python.core.PyException;
import org.python.core.PyInteger;
import org.python.core.PyIterator;
import org.python.core.PyNone;
import org.python.core.PyObject;
import org.python.core.PyString;
import org.python.core.PyTuple;
import org.python.core.PyXRange;

/**
 * Functional tools for creating and using iterators. Java implementation of the CPython module
 * itertools.
 * 
 * @since 2.5
 */
public class itertools implements ClassDictInit {

    public static PyString __doc__ = new PyString(
            "Functional tools for creating and using iterators.\n\nInfinite iterators:\n"
                    + "count([n]) --> n, n+1, n+2, ...\ncycle(p) --> p0, p1, ... plast, p0, p1, ...\n"
                    + "repeat(elem [,n]) --> elem, elem, elem, ... endlessly or up to n times\n\n"
                    + "Iterators terminating on the shortest input sequence:"
                    + "\nizip(p, q, ...) --> (p[0], q[0]), (p[1], q[1]), ... \n"
                    + "ifilter(pred, seq) --> elements of seq where pred(elem) is True\n"
                    + "ifilterfalse(pred, seq) --> elements of seq where pred(elem) is False\n"
                    + "islice(seq, [start,] stop [, step]) --> elements from\n       seq[start:stop:step]\n"
                    + "imap(fun, p, q, ...) --> fun(p0, q0), fun(p1, q1), ...\n"
                    + "starmap(fun, seq) --> fun(*seq[0]), fun(*seq[1]), ...\n"
                    + "chain(p, q, ...) --> p0, p1, ... plast, q0, q1, ... \n"
                    + "takewhile(pred, seq) --> seq[0], seq[1], until pred fails\n"
                    + "dropwhile(pred, seq) --> seq[n],seq[n+1], starting when pred fails\n"
                    + "groupby(iterable[, keyfunc]) -> create an iterator which returns\n"
                    + "(key, sub-iterator) grouped by each value of key(value)."
                    + "tee(iterable, n=2) --> tuple of n independent iterators.");

    /**
     * Iterator base class used by most methods.
     */
    static abstract class ItertoolsIterator extends PyIterator {

        /**
         * Returns the next element from an iterator. If it raises/throws StopIteration just store
         * the Exception and return null according to PyIterator practice.
         */
        protected PyObject nextElement(PyObject pyIter) {
            PyObject element = null;
            try {
                element = pyIter.__iternext__();//next();
            } catch (PyException pyEx) {
                if (pyEx.match(Py.StopIteration)) {
                    // store exception - will be used by PyIterator.next()
                    stopException = pyEx;
                } else {
                    throw pyEx;
                }
            }
            return element;
        }
    }

    public static void classDictInit(PyObject dict) {
    }
    

    public static PyString __doc__count = new PyString(
            "count([firstval]) --> count object\n\nReturn a count object whose .next() "
                    + "method returns consecutive\nintegers starting from zero or, if specified, from firstval.");

    /**
     * Creates an iterator that returns consecutive integers starting at <code>init</code>.
     */
    public static PyIterator count(final int init) {
        return new PyIterator() {
            int counter = init;

            public PyObject __iternext__() {
                return new PyInteger(counter++);
            }
            
            public PyString __repr__() {
                return (PyString)(Py.newString("count(%d)").__mod__(Py.newInteger(counter)));
            }

        };
    }

    /**
     * Creates an iterator that returns consecutive integers starting at 0.
     */
    public static PyIterator count() {
        return itertools.count(0);
    }

    public static PyString __doc__cycle = new PyString(
            "cycle(iterable) --> cycle object\n\nReturn elements from the iterable "
                    + "until itis exhausted.\nThen repeat the sequence indefinitely.");

    /**
     * Returns an iterator that iterates over an iterable, saving the values for each iteration.
     * When the iterable is exhausted continues to iterate over the saved values indefinitely.
     */
    public static PyIterator cycle(final PyObject sequence) {
        return new ItertoolsIterator() {
            List<PyObject> saved = new ArrayList<PyObject>();
            int counter = 0;
            PyObject iter = sequence.__iter__();

            boolean save = true;

            public PyObject __iternext__() {
                if (save) {
                    PyObject obj = nextElement(iter);
                    if (obj != null) {
                        saved.add(obj);
                        return obj;
                    } else {
                        save = false;
                    }
                }
                if (saved.size() == 0) {
                    return null;
                }
                
                // pick element from saved List
                if (counter >= saved.size()) {
                    // start over again
                    counter = 0;
                }
                return saved.get(counter++);
            }

        };
    }
    
    

    public static PyString __doc__chain = new PyString(
            "chain(*iterables) --> chain object\n\nReturn a chain object "
                    + "whose .next() method returns elements from the\nfirst iterable until it is exhausted, then elements"
                    + " from the next\niterable, until all of the iterables are exhausted.");

    /**
     * Creates an iterator that iterates over a <i>chain</i> of iterables.
     */
    public static PyIterator chain(final PyObject[] iterables) {
        final PyObject[] iterators = new PyObject[iterables.length];
        for (int i = 0; i < iterables.length; i++) {
            iterators[i] = iterables[i].__iter__();
        }

        return new ItertoolsIterator() {
            int iteratorIndex = 0;

            public PyObject __iternext__() {
                PyObject next = null;
                for (; iteratorIndex < iterators.length; iteratorIndex++) {
                    next = nextElement(iterators[iteratorIndex]);
                    if (next != null) {
                        break;
                    }
                }
                return next;
            }

        };
    }

    public static PyString __doc__repeat = new PyString(
            "'repeat(element [,times]) -> create an iterator which returns the element\n"
                    + "for the specified number of times.  If not specified, returns the element\nendlessly.");

    /**
     * Creates an iterator that returns the same object the number of times given by
     * <code>times</code>.
     */
    public static PyIterator repeat(final PyObject object, final int times) {
        return new PyIterator() {
            int counter = times;

            public PyObject __iternext__() {
                if (counter > 0) {
                    counter--;
                    return object;
                }
                return null;
            }
            
            public int __len__() {
               return times; 
            }
            
            public PyString __repr__() {
                return (PyString)(Py.newString("repeat(%r, %d)").
                        __mod__(new PyTuple(object, Py.newInteger(counter))));
            }
        };
    }

    /**
     * Creates an iterator that returns the same object over and over again.
     */
    public static PyIterator repeat(final PyObject object) {
        return new PyIterator() {
            public PyObject __iternext__() {
                return object;
            }
                        
            public PyString __repr__() {
                return (PyString)(Py.newString("repeat(%r)").
                        __mod__(new PyTuple(object)));
            }
        };
    }

    public static PyString __doc__imap = new PyString(
            "'map(func, *iterables) --> imap object\n\nMake an iterator that computes the "
                    + "function using arguments from\neach of the iterables.\tLike map() except that it returns\n"
                    + "an iterator instead of a list and that it stops when the shortest\niterable is exhausted "
                    + "instead of filling in None for shorter\niterables.");

    /**
     * Works as <code>__builtin__.map()</code> but returns an iterator instead of a list. (Code in
     * this method is based on __builtin__.map()).
     */
    public static PyIterator imap(PyObject[] argstar) {
        final int n = argstar.length - 1;
        if (n < 1) {
            throw Py.TypeError("imap requires at least two arguments");
        }

        final PyObject callable = argstar[0];
        final PyObject[] iters = new PyObject[n];

        for (int j = 0; j < n; j++) {
            iters[j] = Py.iter(argstar[j + 1], "argument " + (j + 1)
                    + " to imap() must support iteration");
        }

        return new PyIterator() {
            PyObject[] args = new PyObject[n];

            PyObject element = null;

            public PyObject __iternext__() {

                for (int i = 0; i < n; i++) {
                    if ((element = iters[i].__iternext__()) != null) {
                        // collect the arguments for the callable
                        args[i] = element;
                    } else {
                        // break iteration
                        return null;
                    }
                }
                if (callable == Py.None) {
                    // if None is supplied as callable we just return what's in
                    // the iterable(s)
                    if (n == 1) {
                        return args[0];
                    } else {
                        return new PyTuple(args.clone());
                    }
                } else {
                    return callable.__call__(args);
                }
            }
        };
    }

    public static PyString __doc__islice = new PyString(
            "islice(iterable, [start,] stop [, step]) --> islice object\n"
                    + "\nReturn an iterator whose next() method returns selected values from an\n"
                    + "iterable.  If start is specified, will skip all preceding elements;\notherwise, start defaults to zero."
                    + "Step defaults to one.  If\nspecified as another value, step determines how manyvalues are \n"
                    + "skipped between successive calls.  Works like a slice() on a list\nbut returns an iterator.");

    
    private static int py2int(PyObject obj, int defaultValue, String msg) {
        if (obj instanceof PyNone) {
            return defaultValue;
        } else {
            int value = defaultValue;
            try {
                value = Py.py2int(obj);
            }
            catch (PyException pyEx) {
                if (pyEx.match(Py.TypeError)) {
                    throw Py.ValueError(msg);
                } else {
                    throw pyEx;
                }
            }
            return value;
        }
    }
    /**
     * Creates an iterator that returns selected values from an iterable.
     * 
     * @param startObj
     *            the index of where in the iterable to start returning values
     * @param stopObj
     *            the index of where in the iterable to stop returning values
     * @param stepObj
     *            the number of steps to take beween each call to <code>next()</code>
     */
    public static PyIterator islice(final PyObject iterable, PyObject startObj,
            PyObject stopObj, PyObject stepObj) {
        final int stop = py2int(stopObj, 0, "Stop argument must be a non-negative integer or None");
        final int start = py2int(startObj, 0, "Start argument must be a non-negative integer or None");
        final int step = py2int(stepObj, 1, "Step argument must be a non-negative integer or None");
        final boolean stopNone = stopObj instanceof PyNone;

        if (start < 0 || step < 0 || stop < 0) {
            throw Py.ValueError("Indices for islice() must be non-negative integers");
        }

        if (step == 0) {
            throw Py.ValueError("Step must be one or larger for islice()");
        }

        return new ItertoolsIterator() {
            int counter = start;

            int lastCount = 0;

            PyObject iter = iterable.__iter__();

            public PyObject __iternext__() {
                PyObject result = null;

                if (counter >= stop && !stopNone) {
                    return null;
                }

                while (lastCount <= counter) {
                    result = nextElement(iter);
                    lastCount++;
                }
                counter += step;
                return result;
            }

        };

    }

    /**
     * @see #islice(PyObject, PyObject, PyObject, PyObject) startObj defaults to 0 and stepObj to 1
     */
    public static PyIterator islice(PyObject iterable, PyObject stopObj) {
        return islice(iterable, new PyInteger(0), stopObj, new PyInteger(1));
    }

    /**
     * @see #islice(PyObject, PyObject, PyObject, PyObject) stepObj defaults to 1
     */
    public static PyIterator islice(PyObject iterable, PyObject start,
            PyObject stopObj) {
        return islice(iterable, start, stopObj, new PyInteger(1));
    }

    /**
     * Iterator base class for iterators returned by <code>ifilter</code> and
     * <code>ifilterfalse</code>.
     */
    static class FilterIterator extends ItertoolsIterator {
        private PyObject predicate;

        private PyObject iterator;

        private boolean filterTrue;

        FilterIterator(PyObject predicate, PyObject iterable, boolean filterTrue) {
            if (predicate instanceof PyNone) {
                this.predicate = null;
            } else {
                this.predicate = predicate;
            }
            this.iterator = iterable.__iter__();
            this.filterTrue = filterTrue;
        }

        public PyObject __iternext__() {

            while (true) {
                PyObject element = nextElement(iterator);
                if (element != null) {
                    // the boolean value of calling predicate with the element
                    // or if predicate is null/None of the element itself
                    boolean booleanValue = predicate != null ? predicate
                            .__call__(element).__nonzero__() : element
                            .__nonzero__();
                    if (booleanValue == filterTrue) {
                        // if the boolean value is the same as filterTrue return
                        // the element
                        // for ifilter filterTrue is always true, for
                        // ifilterfalse always false
                        return element;
                    }
                } else {
                    return null;
                }
            }
        }
    }

    public static PyString __doc__ifilter = new PyString(
            "ifilter(function or None, sequence) --> ifilter object\n\n"
                    + "Return those items of sequence for which function(item) is true.\nIf function is None, "
                    + "return the items that are true.");

    /**
     * Creates an iterator that returns the items of the iterable for which
     * <code>predicate(item)</code> is <code>true</code>. If <code>predicate</code> is null
     * (None) return the items that are true.
     */
    public static PyIterator ifilter(PyObject predicate, PyObject iterable) {
        return new FilterIterator(predicate, iterable, true);
    }

    public static PyString __doc__ifilterfalse = new PyString(
            "'ifilterfalse(function or None, sequence) --> ifilterfalse object\n\n"
                    + "Return those items of sequence for which function(item) is false.\nIf function is None, "
                    + "return the items that are false.'");

    /**
     * Creates an iterator that returns the items of the iterable for which
     * <code>predicate(item)</code> is <code>false</code>. If <code>predicate</code> is null
     * (None) return the items that are false.
     */
    public static PyIterator ifilterfalse(PyObject predicate, PyObject iterable) {
        return new FilterIterator(predicate, iterable, false);
    }

    public static PyString __doc__izip = new PyString(
            "izip(iter1 [,iter2 [...]]) --> izip object\n\nReturn an izip object "
                    + "whose .next() method returns a tuple where\nthe i-th element comes from the i-th iterable argument.  "
                    + "The .next()\nmethod continues until the shortest iterable in the argument sequence\nis exhausted and then it "
                    + "raises StopIteration.  Works like the zip()\nfunction but consumes less memory by returning an iterator "
                    + "instead of\na list.");

    /**
     * Create an iterator whose <code>next()</code> method returns a tuple where the i-th element
     * comes from the i-th iterable argument. Continues until the shortest iterable is exhausted.
     * (Code in this method is based on __builtin__.zip()).
     * 
     */
    public static PyIterator izip(PyObject[] argstar) {
        final int itemsize = argstar.length;
        
        if (itemsize == 0) {
            return (PyIterator)(new PyXRange(0).__iter__());            
        }

        // Type check the arguments; they must be sequences.
        final PyObject[] iters = new PyObject[itemsize];

        for (int i = 0; i < itemsize; i++) {
            PyObject iter = argstar[i].__iter__();
            if (iter == null) {
                throw Py.TypeError("izip argument #" + (i + 1)
                        + " must support iteration");
            }
            iters[i] = iter;
        }

        return new ItertoolsIterator() {

            public PyObject __iternext__() {
                if (itemsize == 0)
                    return null;

                PyObject[] next = new PyObject[itemsize];
                PyObject item;

                for (int i = 0; i < itemsize; i++) {

                    item = nextElement(iters[i]);

                    if (item == null) {
                        return null;
                    }
                    next[i] = item;
                }
                return new PyTuple(next);
            }

        };

    }

    public static PyString __doc__starmap = new PyString(
            "starmap(function, sequence) --> starmap object\n\nReturn an "
                    + "iterator whose values are returned from the function evaluated\nwith an argument tuple taken from the "
                    + "given sequence.");

    /**
     * Create an iterator whose <code>next()</code> method returns the result
     * of calling the function (first argument) with a tuple of arguments
     * returned from the iterable (second argument).
     * 
     * @param starargs
     *            [0] = callable function, [1] = iterable with argument tuples
     */
    public static PyIterator starmap(PyObject[] starargs) {
        if (starargs.length != 2) {
            throw Py.TypeError("starmap requires 2 arguments, got "
                    + starargs.length);
        }
        final PyObject callable = starargs[0];
        final PyObject iterator = starargs[1].__iter__();

        return new ItertoolsIterator() {

            public PyObject __iternext__() {
                PyObject args = nextElement(iterator);
                PyObject result = null;

                if (args != null) {
                    if (!args.getClass().isAssignableFrom(PyTuple.class)) {
                        throw Py.TypeError("iterator must return a tuple");
                    }
                    PyTuple argTuple = (PyTuple) args;
                    // convert to array of PyObjects in call to function
                    result = callable.__call__(argTuple.getArray());
                }
                return result;
            }

        };
    }

    /**
     * Iterator base class used by <code>dropwhile()</code> and <code>takewhile</code>.
     */
    static class WhileIterator extends ItertoolsIterator {
        private PyObject iterator;

        private PyObject predicate;

        // flag that indicates if the iterator shoul drop or return arguments "while" the predicate is true
        private boolean drop;

        // flag that is set once the predicate is satisfied
        private boolean predicateSatisfied;

        WhileIterator(PyObject predicate, PyObject iterable, boolean drop) {
            this.predicate = predicate;
            iterator = iterable.__iter__();
            this.drop = drop;
        }

        public PyObject __iternext__() {

            while (true) {
                PyObject element = nextElement(iterator);
                if (element != null) {
                    if (!predicateSatisfied) {
                        // the predicate is not satisfied yet (or still satisfied in the case of drop beeing 
                        // false), so we need to check it
                        if (predicate.__call__(element).__nonzero__() != drop) {
                            predicateSatisfied = drop;
                            return element;
                        }
                        predicateSatisfied = !drop;
                    } else {
                        if (drop) {
                            return element;
                        } else {
                            // end iteration if predicate is false and drop is false
                            return null;
                        }
                    }
                } else {
                    // end iteration
                    return null;
                }

            }
        }
    }

    public static PyString __doc__dropwhile = new PyString(
            "dropwhile(predicate, iterable) --> dropwhile object\n\nDrop items "
                    + "from the iterable while predicate(item) is true.\nAfterwards, return every element until theiterable is exhausted.");

    /**
     * Create an iterator that drops items from the iterable while <code>prdicate(item)</code>
     * equals true. After which every remaining item of the iterable is returned.
     */
    public static PyIterator dropwhile(PyObject predicate, PyObject iterable) {
        return new WhileIterator(predicate, iterable, true);
    }

    public static PyString __doc__takewhile = new PyString(
            "takewhile(predicate, iterable) --> takewhile object\n\nReturn "
                    + "successive entries from an iterable as long as the \npredicate evaluates to true for each entry.");

    /**
     * Create an iterator that returns items from the iterable while <code>predicate(item)</code>
     * is true. After which iteration is stopped.
     */
    public static PyIterator takewhile(PyObject predicate, PyObject iterable) {
        return new WhileIterator(predicate, iterable, false);
    }

    private final static class GroupBy extends ItertoolsIterator {

        private final PyObject iterator;
        private final PyObject keyFunc;
        private PyObject currentKey;
        private PyObject currentValue;
        private PyObject targetKey;

        private GroupBy(PyObject iterable, PyObject key) {
            iterator = iterable.__iter__();
            keyFunc = key;
            targetKey = currentKey = currentValue = new PyXRange(0);
        }

        public PyObject __iternext__() {
            while (currentKey.equals(targetKey)) {
                currentValue = nextElement(iterator);
                if (currentValue == null) {
                    return null;
                }
                if (keyFunc == null) {
                    currentKey = currentValue;
                } else {
                    currentKey = keyFunc.__call__(currentValue);
                }
            }
            targetKey = currentKey;
            return new PyTuple(currentKey, new GroupByIterator());
        }

        private class GroupByIterator extends ItertoolsIterator {

            private boolean completed = false;

            public PyObject __iternext__() {
                final PyObject item = currentValue;
                if (completed) {
                    return null;
                }
                currentValue = nextElement(iterator);
                if (currentValue == null) {
                    completed = true;
                } else {
                    if (keyFunc == null) {
                        currentKey = currentValue;
                    } else {
                        currentKey = keyFunc.__call__(currentValue);
                    }
                }
                if (!currentKey.equals(targetKey)) {
                    completed = true;
                }
                return item;
            }
        }
    }

    public static PyString __doc__groupby = new PyString(
            "groupby(iterable[, keyfunc]) -> create an iterator which returns\n" +
            "(key, sub-iterator) grouped by each value of key(value).");

    /**
     * Create an iterator which returns the pair (key, sub-iterator) grouped by key(value).
     */
    public static PyIterator groupby(PyObject [] args, String [] kws) {
        ArgParser ap = new ArgParser("groupby", args, kws, "iterable", "key");
        if(args.length > 2){
            throw Py.TypeError("groupby takes two arguments, iterable and key");
        }
        PyObject iterable = ap.getPyObject(0);
        PyObject key = ap.getPyObject(1, null);
        return new GroupBy(iterable, key);
    }
    
    public static PyString __doc__tee = new PyString(
            "tee(iterable, n=2) --> tuple of n independent iterators.");

    /**
     * Create a tuple of iterators, each of which is effectively a copy of iterable.
     */
    public static PyTuple tee(PyObject iterable, final int n) {
        return new PyTuple(PyTeeIterator.makeTees(iterable, n));
    }

    /**
     * Create a pair of iterators, each of which is effectively a copy of iterable.
     */
    public static PyTuple tee(PyObject iterable) {
        return tee(iterable, 2);
    }

}
