(define-module (language python module itertools)
  #:use-module (oop pf-objects)
  #:use-module (language python for)
  #:use-module (language python try)
  #:use-module (language python yield)
  #:use-module (language python def)
  #:use-module (language python exceptions)
  #:use-module ((language python module python)
		#:select (iter zip next))
  
  #:export (count cycle repeat accumulate chain compress dropwhile
		  filterfalse groupby isslice starmap takewhile
		  tee zip_longest product permutation combination
		  combination_with_replacement))

(define count
  (lambda* (start #:optional (step 1))
    ((make-generator ()
     (lambda (yield)
       (let lp ((i start))
         (yield i)
         (lp (+ i step))))))))

(define cycle
  (make-generator (p)
   (lambda (yield p)
     (let lp ()
       (for ((x : p)) () (yield x))
       (lp)))))

(define repeat
  (lambda* (e #:optional (n -1))
    ((make-generator ()
     (lambda (yield)                
       (let lp ((i 0))
         (if (not (= i n))
             (begin
               (yield e)
               (lp (+ i 1))))))))))

(define accumulate
  (lambda* (p #:optional (f +))
    ((make-generator ()
     (lambda (yield)                
       (for ((x : p)) ((s 0) (first? #t))
            (if first?
                (begin
                  (yield x)
                  (values x #f))
                (let ((s (f x s)))
                  (yield s)
                  (values s #f)))))))))
  
(define-python-class chain ()
  (define __call__
    (lambda (self . l)
      ((make-generator ()
       (lambda (yield)
         (let lp ((l l))
	   (if (pair? l)
               (begin
                 (for ((x : (car l))) ()
                      (yield x))
                 (lp (cdr l))))))))))

  (define from_iterable
    (make-generator (i)
     (lambda (yield i)
       (for ((ii : i)) ()
	    (for ((x : ii)) ()
		 (yield x)))))))

(define compress
  (make-generator (data selectors)
   (lambda (yield data selectors)
     (for ((d : data) (s : selectors)) ()
	  (if s (yield d))))))

(define dropwhile
  (make-generator (pred seq)
   (lambda (yield pred seq)
     (for ((x : seq)) ((start? #f))
	  (if start?
	      (begin
		(yield x)
		#t)
	      (if (pred x)
		  #f
		  (begin
		    (yield x)
		    #t)))))))

(define filterfalse
  (make-generator (pred seq)
   (lambda (yield pred seq)
     (for ((x : seq)) ()
	  (if (not (pred x)) (yield x))))))

(define none (list 'none))
(define groupby
  (lambda* (seq #:optional (key (lambda (x) x)))
    ((make-generator ()
     (lambda (yield)                
       (for ((x : seq)) ((k none) (l '()))
            (if (eq? k none)
                (values (key x) (list x))
                (let ((kk (key x)))
                  (if (equal? k kk)
                      (values k (cons x l))
                      (begin
                        (yield k (reverse l))
                        (values kk (list x))))))
            #:final
            (if (not (eq? k none))
                (yield k (reverse l)))))))))
  
     
(define isslice
  (lambda* (seq #:optional (start 0) (stop -1) (step 1))
    ((make-generator ()
     (lambda (yield)
       (for ((x : seq) (i : (count 0))) ()
            (if (= i stop) (break))
            (if (and (>= i start)
                     (= (modulo (- i start) step) 0))
                (yield x))))))))

(define starmap
  (make-generator (f seq)
   (lambda (yield f seq)
     (for ((x : seq)) () (yield (apply f x))))))

(define takewhile
  (make-generator (pred seq)
   (lambda (yield pred seq)
     (for ((x : seq)) ()
	  (if (not (pred x)) (break))
	  (yield x)))))


(define tee
  (lambda* (it #:optional (n 2))
    (define (clone it)
       (let ((l '())
	     (i  0)
	     (r '()))
	 (define (mk)
	   ((make-generator ()
	      (lambda (yield)
		(let lp ((head #f))
		  (if (and head (= i head))
		      (let* ((x  (call-with-values (lambda () (next it))
                                   (lambda x x)))
			     (i0 (+ i 1)))
			
			(set! r (cons x r))
			(set! i i0)
			(apply yield x)
			(lp i0))
		      (if (pair? l)
			  (let ((x (car l)))
			    (set! l (cdr l))
			    (apply yield x)
			    (lp #f))
			  (if (null? r)
			      (lp i)
			      (begin
				(set! l (reverse r))
				(set! r '())
				(lp #f))))))))))
	 (values (mk) (mk))))

     (if (<= n 0)
	 '()
	 (let lp ((i 1) (it (iter it)))
	   (if (< i n)
	       (call-with-values (lambda () (clone it))
		 (lambda (it1 it2)
		   (cons it1 (lp (+ i 1) it2))))
	       (list it))))))
	   
(define zip_longest
  (lam ((* l) (= fillvalue None))
    ((make-generator
      (lambda (yield)
	(define mkit
          (lambda (it)
	   ((make-generator ()
	     (lambda (yield)
	       (for ((x : it)) ()
                    (yield (cons 1 x)))
               (let lp ()
                 (yield (cons 0 0))
                 (lp)))))))
	
	(for ((x : (apply zip (map mkit l)))) ()
	     (if (= (apply + (map car x)) 0)
		 (break)
		 (yield (map (lambda (y) (if (= (car y) 0) fillvalue (cdr y)))
			     x)))))))))
	
  
(def (product (* iterables) (= repeat 1))
     ((make-generator ()
      (lambda (yield)       
	(let* ((iterables (let lp ((i 1) (l (map iter iterables)))
                            (if (= i repeat)
                                l
                                (let* ((t (map tee  l))
                                       (h (map car  t))
                                       (r (map cadr t)))
                                  (append h (lp (+ i 1) r)))))))
          (let lp ((l iterables) (ret '()))
            (if (pair? l)
                (let ((x.y (tee (car l))))
                  (for ((x : (car x.y))) ((l (cdr l)))
                       (lp l (cons x ret))
                       #:final (cons (cadr x.y) l)))
                (yield (reverse ret)))))))))
                       
(def (permutation it (= r None))
     ((make-generator ()
      (lambda (yield)
	(let* ((ll (for ((x : it)) ((l '())) (cons x l) #:final (reverse l)))
	       (N  (length ll)))
	  (let lp ((l ll) (ri '()) (rl '()) (i 0) (depth 0))
	    (if (pair? l)
		(if (and (or (eq? r None) (< depth r)) (< depth N))
		    (if (member i ri)
			(lp (cdr l) ri rl (+ i 1) depth)
			(let ((x (car l)))
			  (lp ll (cons i ri) (cons x rl) 0 (+ depth 1))
			  (lp (cdr l) ri rl (+ i 1) depth)))
		    (yield (reverse rl))))))))))

(def (combination it r)
     ((make-generator ()
      (lambda (yield)
	(let* ((ll (for ((x : it)) ((l '()))
			(cons x l)
			#:final (reverse l)))
	      (N  (length ll)))
	  (let lp ((l ll) (rl '()) (depth 0))
	    (if (< depth r)
		(if (>= (length l) (- r depth))
		    (let ((x (car l)))
		      (lp (cdr l) (cons x rl) (+ depth 1))
		      (lp (cdr l) rl          depth)))
		(yield (reverse rl)))))))))

(def (combination_with_replacement it r)
     ((make-generator ()
      (lambda (yield)
	(let* ((ll (for ((x : it)) ((l '()))
			(cons x l)
			#:final (reverse l)))
	      (N  (length ll)))
	  (let lp ((l ll) (rl '()) (depth 0))
	    (if (< depth r)
		(if (pair? l)
		    (let ((x (car l)))
		      (lp l (cons x rl) (+ depth 1))
		      (lp (cdr l) rl          depth)))
		(yield (reverse rl)))))))))