#lang scheme
(require (for-syntax scheme))

(define (timed f)
  (lambda as
    (time (apply f as))))

(define (parse fname)
  (with-input-from-file
      fname
    (lambda ()
      (let loop ()
        (let ([c (read)])
          (if (eof-object? c)
              empty
              (list* c (loop))))))))

(define qsort1
  (match-lambda
    [(list) empty]
    [l
     (define x (list-ref l (floor (* (length l) 1/2))))
     (define-values (pre eq post)
       (for/fold ([pre empty]
                  [eq empty]
                  [post empty])
         ([y (in-list l)])
         (cond
           [(= y x)
            (values pre (list* y eq) post)]
           [(< y x)
            (values (list* y pre) eq post)]
           [else
            (values pre eq (list* y post))])))
     (append (qsort1 pre)
             eq
             (qsort1 post))]))
(define sort1 (compose (timed qsort1) parse))

(define qsort2
  (match-lambda
    [(list) empty]
    [l
     (define e1 (list-ref l (floor (* (length l) 1/3))))
     (define e2 (list-ref l (floor (* (length l) 2/3))))
     (define x1 (min e1 e2))
     (define x2 (max e1 e2))
     (define-values (pre1 eq1 pre2 eq2 post2)
       (for/fold ([pre1 empty]
                  [eq1 empty]
                  [pre2 empty]
                  [eq2 empty]
                  [post2 empty])
         ([y (in-list l)])
         (cond
           [(= y x1)
            (values pre1 (list* y eq1) pre2 eq2 post2)]
           [(= y x2)
            (values pre1 eq1 pre2 (list* y eq2) post2)]
           [(< y x1)
            (values (list* y pre1) eq1 pre2 eq2 post2)]
           [(< y x2)
            (values pre1 eq1 (list* y pre2) eq2 post2)]
           [else
            (values pre1 eq1 pre2 eq2 (list* y post2))])))
     (append (qsort2 pre1)
             eq1
             (qsort2 pre2)
             eq2
             (qsort2 post2))]))
(define sort2 (compose (timed qsort2) parse))

(define qsort3
  (match-lambda
    [(list) empty]
    [l
     (define e1 (list-ref l (floor (* (length l) 1/4))))
     (define e2 (list-ref l (floor (* (length l) 2/4))))
     (define e3 (list-ref l (floor (* (length l) 2/4))))
     (define x1 (min e1 e2 e3))
     (define x3 (max e1 e2 e3))
     (define x2
       (cond
         [(not (or (= e1 x1) (= e1 x3))) e1]
         [(not (or (= e2 x1) (= e2 x3))) e2]
         [else x3]))
     (define-values (pre1 eq1 pre2 eq2 pre3 eq3 post3)
       (for/fold ([pre1 empty]
                  [eq1 empty]
                  [pre2 empty]
                  [eq2 empty]
                  [pre3 empty]
                  [eq3 empty]
                  [post3 empty])
         ([y (in-list l)])
         (cond
           [(= y x1)
            (values pre1 (list* y eq1) pre2 eq2 pre3 eq3 post3)]
           [(= y x2)
            (values pre1 eq1 pre2 (list* y eq2) pre3 eq3 post3)]
           [(= y x3)
            (values pre1 eq1 pre2 eq2 pre3 (list* y eq3) post3)]
           [(< y x1)
            (values (list* y pre1) eq1 pre2 eq2 pre3 eq3 post3)]
           [(< y x2)
            (values pre1 eq1 (list* y pre2) eq2 pre3 eq3 post3)]
           [(< y x3)
            (values pre1 eq1 pre2 eq2 (list* y pre3) eq3 post3)]
           [else
            (values pre1 eq1 pre2 eq2 pre3 eq3 (list* y post3))])))
     (append (qsort3 pre1)
             eq1
             (qsort3 pre2)
             eq2
             (qsort3 pre3)
             eq3
             (qsort3 post3))]))
(define sort3 (compose (timed qsort3) parse))

(define-syntax (define-qsort stx)
  (syntax-case stx ()
    [(_ qsortN pivots-stx)
     (local [(define pivots (syntax->datum #'pivots-stx))
             (define (split-ids ids)
               (for/list ([i (in-range pivots)])
                 (list (for/list ([j (in-range pivots)]
                                  [pre_j ids]
                                  #:when (j . < . i))
                         pre_j)
                       (for/list ([j (in-range pivots)]
                                  [pre_j ids]
                                  #:when (i . < . j))
                         pre_j))))]
       (with-syntax
           ([(i ...) (for/list ([i (in-range pivots)]) (datum->syntax stx i))]
            [(e_i ...) (generate-temporaries (for/list ([i (in-range pivots)]) (format "pivot~a_" i)))]
            [(p_i ...) (generate-temporaries (for/list ([i (in-range pivots)]) (format "sorted-pivot~a_" i)))]
            [(pre_i ...) (generate-temporaries (for/list ([i (in-range pivots)]) (format "pre~a_" i)))]
            [(eq_i ...) (generate-temporaries (for/list ([i (in-range pivots)]) (format "eq~a_" i)))])
         (with-syntax
             ([(((pre_li ...)
                 (pre_gi ...))
                ...)
               (split-ids (syntax->list #'(pre_i ...)))]
              [(((eq_li ...)
                 (eq_gi ...))
                ...)
               (split-ids (syntax->list #'(eq_i ...)))])
           (syntax/loc stx
             (define qsortN
               (match-lambda
                 [(list) empty]
                 [l
                  (let* ([e_i (list-ref l i)]
                         ...
                         [sorted-pivots (qsort1 (list e_i ...))]
                         [p_i (list-ref sorted-pivots i)]
                         ...)
                    (let-values 
                        ([(pre_i ... eq_i ... post_n)
                          (for/fold ([pre_i empty] ... 
                                     [eq_i empty] ...
                                     [post_n empty])
                            ([y (in-list l)])
                            (cond
                              [(= y p_i)
                               (values pre_i ... 
                                       eq_li ... 
                                       (list* y eq_i) 
                                       eq_gi ...
                                       post_n)]
                              ...
                              [(< y p_i)
                               (values pre_li ...
                                       (list* y pre_i)
                                       pre_gi ...
                                       eq_i ...
                                       post_n)]
                              ...
                              [else
                               (values pre_i ...
                                       eq_i ...
                                       (list* y post_n))]))])
                      (append* (append (qsortN pre_i) eq_i)
                               ...
                               (list (qsortN post_n)))))]))))))]))

(define-qsort qsorta 10)
(define sorta (compose (timed qsorta) parse))

(define (generate fname)
  (with-output-to-file
      fname
    (lambda ()
      (for ([i (in-range (* 1.5 (expt 10 6)))])
        (write (random 100))
        (display #\space)))
    #:exists 'replace))

(define mode (make-parameter sort1))
(define quiet (make-parameter (lambda (x) x)))

(command-line #:program "qsort3"
              #:once-each
              ["--q" "Quiet mode"
                     (quiet void)]
              #:once-any
              ["--g" "Generate a number list"
                     (mode (timed generate))]
              ["--1" "Sort with one pivot"
                     (mode sort1)]
              ["--2" "Sort with two pivots"
                     (mode sort2)]
              ["--3" "Sort with three pivots"
                     (mode sort3)]
              ["--a" "Sort with auto pivots"
                     (mode sorta)]
              #:args (filename)
              ((quiet)
               ((mode) filename)))