(define-module (language python module itertools)
  #:use-module (language python for)
  #:use-module (language python yield)
  #:use-module (language python def)
  #:use-module (language python module copy)
  #:use-module (language python module python)
  #:use-module (language python module copy)
  #:export (count cycle repeat accumulate chain compress dropwhile
		  filterfalse groupby isslice starmap takewhile
		  tee zip_longest))

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

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

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

(define accumulate
  (make-generator
   (lambda* (yield p #:optional (f +))
     (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__
    (make-generator
     (lambda (yield . l)
       (let lp ((l l))
	 (if (pair? l)
	     (begin
	       (for ((x : (car l))) ()
		    (yield x))
	       (lp (cdr l))))))))
  (define from_iterable
    (make-generator
     (lambda (yield i)
       (for ((ii : i)) ()
	    (for ((x : ii)) ()
		 (yield x)))))))

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

(define dropwhile
  (make-generator
   (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
   (lambda (yield pred seq)
     (for ((x : seq))
	  (if (not (f x)) (yield x))))))

(define none (list 'none))
(define groupby
  (make-generator
   (lambda* (yield seq #:optional (key (lambda (x) x)))
     (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 l (reverse l)))))))
     
     
(define isslice
  (make-generator
   (lambda* (yield seq #:optional (start 0) (stop -1) (step 1))
     (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
   (lambda (yield f seq)
     (for ((x : seq)) () (yield (f x))))))

(define takewhile
  (make-generator
   (lambda (yield pred seq)
     (for ((x : seq)) ()
	  (if (not (pred x)) (break))
	  (yield x)))))
  
(define tee
  (make-generator
   (lambda (yield it n)
     (let lp ((i 0))
       (if (< i n)
	   (cons (deepcopy it)
		 (lp (+ i 1)))
	   '())))))

(define zip_longest
  (make-generator
   (lam (yield (* l) (= fillvalue None))
	(define mkit
	  (make-generator
	   (lambda (yield it)
	     (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 (x) (if (= (car x) 0) fillvalue (cdr x)))
			       x))))))))