(define-module (language python module heapq)
  #:use-module (ice-9 control)
  #:use-module (language python for)
  #:use-module (language python yield)
  #:use-module (language python try)
  #:use-module (language python exceptions)
  #:use-module (language python list)
  #:use-module (language python def)
  #:use-module (oop pf-objects)
  #:use-module ((language python module python)
		#:select (zip sorted range enumerate iter max min))
  #:export (heappush heappop heapify heapreplace merge nlargest nsmallest
		     heappushpop))

(define-syntax-rule (aif it p a b) (let ((it p)) (if it a b)))

(define (heappush heap item)
  (pylist-append! heap item)
  (siftdown heap 0 (- (len heap) 1)))

(define (heappop heap)
  (let ((lastelt (pylist-pop! heap)))
    (if (> (len heap) 0)
	(let ((ret (pylist-ref heap 0)))
	  (pylist-set! heap 0 lastelt)
	  (siftup heap 0)
	  ret)
	lastelt)))

(define (heapreplace heap item)
  (let ((ret (pylist-ref heap 0)))
    (pylist-set! heap 0 item)
    (siftup heap 0)
    ret))

(define (heapify x)
  (let lp ((i (- (floor-quotient (len x) 2) 1)))
    (if (>= i 0)
	(begin
	  (siftup x i)
	  (lp (- i 1))))))

(define (heappop_max heap)
  (let ((lastelt (pylist-pop! heap)))
    (if (> (len heap) 0)
	(let ((ret (pylist-ref heap 0)))
	  (pylist-set! heap 0 lastelt)
	  (siftup_max heap 0)
	  ret)
	lastelt)))

(define (heapreplace_max heap item)
  (let ((ret (pylist-ref heap 0)))
    (pylist-set! heap 0 item)
    (siftup_max heap 0)
    ret))	

(define (heapify_max x)
  (let lp ((i (- (floor-quotient (len x) 2) 1)))
    (if (>= i 0)
	(begin
	  (siftup_max x i)
	  (lp (- i 1))))))

(define (siftdown heap startpos pos)
  (define newitem (pylist-ref heap pos))
  
  (let lp ((pos pos))
    (if (> pos startpos)
	(let ()
	  (define parentpos (ash (- pos 1) 1))
	  (define parent    (pylist-ref heap parentpos))
	  (if (< newitem parent)
	      (begin
		(pylist-set! heap pos parent)
		(lp parentpos))
	      (pylist-set! heap pos newitem)))
	(pylist-set! heap pos newitem))))

(define (siftup heap pos)
  (define endpos   (len heap))
  (define startpos pos)
  (define newitem  (pylist-ref heap pos))
  (let lp ((pos pos) (childpos (+ (* 2 pos) 1)))
    (if (< childpos endpos)
	(let* ((rightpos (+ childpos 1))
	       (childpos
		(if (and (< rightpos endpos)
			 (not (< (pylist-ref heap childpos)
				 (pylist-ref heap rightpos))))
		    rightpos
		    childpos)))
	  (pylist-set! heap pos (pylist-ref heap childpos))
	  (lp childpos (+ (* 2 childpos) 1)))
	(begin
	  (pylist-set! heap pos newitem)
	  (siftdown heap startpos pos)))))

(define (siftdown_max heap startpos pos)
  (define newitem (pylist-ref heap pos))
  
  (let lp ((pos pos))
    (if (> pos startpos)
	(let ()
	  (define parentpos (ash (- pos 1) 1))
	  (define parent    (pylist-ref heap parentpos))
	  (if (< parent newitem)
	      (begin
		(pylist-set! heap pos parent)
		(lp parentpos))
	      (pylist-set! heap pos newitem)))
	(pylist-set! heap pos newitem))))

(define (siftup_max heap pos)
  (define endpos   (len heap))
  (define startpos pos)
  (define newitem  (pylist-ref heap pos))
  (let lp ((pos pos) (childpos (+ (* 2 pos) 1)))
    (if (< childpos endpos)
	(let* ((rightpos (+ childpos 1))
	       (childpos
		(if (and (< rightpos endpos)
			 (not (< (pylist-ref heap rightpos)
				 (pylist-ref heap childpos))))
		    rightpos
		    childpos)))
	  (pylist-set! heap pos (pylist-ref heap childpos))
	  (lp childpos (+ (* 2 childpos) 1)))
	(begin
	  (pylist-set! heap pos newitem)
	  (siftdown_max heap startpos pos)))))

(def (merge (* iterables) (= key (lambda (x) x)) (= reverse #f))
    ((make-generator ()
       (lambda (yield)
	 (define h (py-list))
	 (define (h_append x) (pylist-append! h x))
	 (define-values (-heapify -heappop -heapreplace direction)
	   (if reverse
	       (values heapify_max
		       heappop_max
		       heapreplace_max
		       -1)
	       (values heapify
		       heappop
		       heapreplace
		       1)))
	   
       	 (for ((order it : (enumerate (map iter iterables)))) ()
	      (define next- (aif x (ref it '__next__)
				 x
				 (lambda () (next it))))
	      (try
	       (lambda ()
		 (h_append (vector (key (next-))
				   (* order direction)
				   next-)))
	       (#:except StopIteration =>
			 (lambda x (-heappop h)))))

	 (-heapify h)
	 (let lp ()
	   (if (> (len h) 1)
	       (try
		(lambda ()
		  (let lp ()
		    (let ((s (pylist-ref h 0)))
		      (yield (vector-ref s 0))
		      (vector-set! s 0 (key ((vector-ref s 2))))
		      (-heapreplace h s)
		      (lp))))
		(#:except StopIteration =>
			  (lambda x (-heappop h) (lp))))))

	 (if (> (len h) 0)
	     (let* ((x (pylist-ref h 0))
		    (n (pylist-ref x 2)))
	       (let lp ((i (vector-ref x 0)))
		 (yield i)
		 (lp (n)))))))))
		 


	 
(define-syntax-rule (mkn nsmallest heapify_max heapreplace_max q< ran o ++)
  (define* (nsmallest n iterable #:key (key (lambda (x) x)))
    (let/ec return
      (if (equal? n 1)
	  (return
	   (let* ((sent (list 1))
		  (res  (min (iter iterable) #:default sent #:key key)))
	     (if (eq? sent res)
		 '()
		 (list res)))))

      (let ((size (try
		   (lambda () (len iterable))
		   (#:except (or TypeError AttributeError) =>
			     (lambda x #f)))))
	(if (and size (>= n size))
	    (return (pylist-slice (sorted iterable #:key key) 0
				  ((@ (guile) min) n size)
				  1))))
  
      (let ((it     (iter iterable))
	    (result (py-list)))
    
	(for ((i elem : (zip (ran n) it))) ()
	     (pylist-append! result (vector (key elem) i elem)))

	(if (= 0 (len result))
	    (return result))

	(heapify_max result)

	(for ((elem : it)) ((top   (vector-ref (pylist-ref result 0) 0))
			    (order (o n)))
	     (let ((k (key elem)))
	       (if (q< k top)
		   (begin
		     (heapreplace_max result (vector k order elem))
		     (values (vector-ref (pylist-ref result 0) 0)
			     (++ order)))
		   (values top order))))
	(pylist-sort! result)
	(return
	 (let ((r (py-list)))
	   (for ((x : result)) ()
		(pylist-append! r (vector-ref x 2)))
	   r))))))

(mkn nsmallest heapify_max heapreplace_max < (lambda (n) (range n))
     (lambda (n) n) (lambda (x) (+ x 1)))

(mkn nlargest  heapify     heapreplace     > (lambda (n) (range 0 (- n) -1))
     (lambda (n) (- n)) (lambda (x) (- x 1)))