;;; Arl Standard Library - Sorting
;;;
;;; Custom sorting implementations are necessary because R's built-in sort
;;; functions (sort, order, sort.int) operate on atomic vectors with standard
;;; comparison operators. They cannot accept an arbitrary comparator function
;;; on heterogeneous list elements, which is what Scheme-style sort requires.

(module sort
  (export list-sort sort-by merge-sorted stable-sort)

  (import list :refer (cadr append reverse))
  (import types :refer (null?))
  (import core :refer (r-call))

  ;;' @section Sorting
  ;;' General-purpose sorting with user-supplied comparator functions.
  ;;' `list-sort` uses quicksort (not stable); `stable-sort` uses merge sort
  ;;' and preserves the relative order of equal elements. `merge-sorted` combines
  ;;' two already-sorted lists.

  ;;' @description Sort list using comparison function (simple quicksort implementation).
  ;;' @param lst List to sort
  ;;' @param comparator Binary predicate returning #t when first arg should come before second
  ;;' @examples
  ;;' (list-sort (list 3 1 4 1 5) <)        ; => (1 1 3 4 5)
  ;;' (list-sort (list "b" "a" "c")
  ;;'       (lambda (a b) (< a b)))   ; => ("a" "b" "c")
  ;;' @assert
  ;;' (assert-equal (list 1 1 3 4 5) (list-sort (list 3 1 4 1 5) <))
  ;;' @seealso stable-sort, sort-by
  ;; Sorting
  (define list-sort
    (lambda (lst comparator)
      (define items (_as-list lst))
      (define n (length items))
      (if (<= n 1)
        items
        (begin
          ;; Simple quicksort with index-based partitioning
          (define pivot ([[ items 1))
          ;; Partition into less and greater using cons + reverse (O(n))
          (define less (list))
          (define greater (list))
          (define i 2)
          (while (<= i n)
            (define elem ([[ items i))
            (if (comparator elem pivot)
              (set! less (cons elem less))
              (set! greater (cons elem greater)))
            (set! i (+ i 1)))
          (append (list-sort (rev less) comparator)
                  (list pivot)
                  (list-sort (rev greater) comparator))))))

  ;;' @description Sort list by applying key-fn to each element, then comparing.
  ;;' @param lst List to sort
  ;;' @param key-fn Function to extract a comparison key from each element
  ;;' @param comparator Binary predicate applied to extracted keys
  ;;' @examples
  ;;' (sort-by (list (list "b" 2) (list "a" 1)) car
  ;;'          (lambda (a b) (< a b)))  ; => (("a" 1) ("b" 2))
  ;;' (sort-by (list 3 -1 2) abs <)     ; => (-1 2 3)
  ;;' @assert
  ;;' (assert-equal (list -1 2 3) (sort-by (list 3 -1 2) abs <))
  ;;' @seealso list-sort, stable-sort
  (define sort-by
    (lambda (lst key-fn comparator)
      (define items (_as-list lst))
      (if (<= (length items) 1)
        items
        (begin
          ;; Create list of (key, value) entries
          (define pairs (lapply items (lambda (x) (list (key-fn x) x))))
          ;; Sort by first element of each entry
          (define sorted-pairs (list-sort pairs (lambda (a b) (comparator (car a) (car b)))))
          ;; Extract values
          (lapply sorted-pairs cadr)))))

  ;;' @description Merge two sorted lists into one sorted list.
  ;;' Stable: when elements are equal (neither comparator direction
  ;;' is true), takes from list1 first.
  ;;' @param list1 First sorted list
  ;;' @param list2 Second sorted list
  ;;' @param comparator Binary predicate used to order elements
  ;;' @examples
  ;;' (merge-sorted (list 1 3 5) (list 2 4 6) <)  ; => (1 2 3 4 5 6)
  ;;' (merge-sorted (list 1 2) (list) <)           ; => (1 2)
  ;;' @assert
  ;;' (assert-equal (list 1 2 3 4 5 6) (merge-sorted (list 1 3 5) (list 2 4 6) <))
  ;;' (assert-equal (list 1 2) (merge-sorted (list 1 2) (list) <))
  ;;' @seealso stable-sort, list-sort
  (define merge-sorted
    (lambda (list1 list2 comparator)
      (define l1 (_as-list list1))
      (define l2 (_as-list list2))
      (define n1 (length l1))
      (define n2 (length l2))
      ;; Build result in reverse using cons, then reverse at end (O(n))
      (define acc (list))
      (define i1 1)
      (define i2 1)
      (while (and (<= i1 n1) (<= i2 n2))
        ;; Take from l2 only when l2 is strictly less than l1; otherwise take from l1 (stable)
        (if (comparator ([[ l2 i2) ([[ l1 i1))
          (begin
            (set! acc (cons ([[ l2 i2) acc))
            (set! i2 (+ i2 1)))
          (begin
            (set! acc (cons ([[ l1 i1) acc))
            (set! i1 (+ i1 1)))))
      (while (<= i1 n1)
        (set! acc (cons ([[ l1 i1) acc))
        (set! i1 (+ i1 1)))
      (while (<= i2 n2)
        (set! acc (cons ([[ l2 i2) acc))
        (set! i2 (+ i2 1)))
      (rev acc)))

  ;;' @description Stable sort - preserves order of equal elements. Uses merge sort internally.
  ;;' @param lst List to sort
  ;;' @param comparator Binary predicate returning #t when first arg should come before second
  ;;' @examples
  ;;' (stable-sort (list 3 1 4 1 5) <)  ; => (1 1 3 4 5)
  ;;' @assert
  ;;' (assert-equal (list 1 1 3 4 5) (stable-sort (list 3 1 4 1 5) <))
  ;;' @note Preserves the relative order of elements that compare equal, unlike `list-sort`.
  ;;' @seealso list-sort, merge-sorted, sort-by
  (define stable-sort
    (lambda (lst comparator)
      (define items (_as-list lst))
      (if (<= (length items) 1)
        items
        (begin
          (define mid (as.integer (/ (length items) 2)))
          (define left (r-call "[" (list items (seq_len mid))))
          (define right (r-call "[" (list items (seq (+ mid 1) (length items)))))
          (merge-sorted (stable-sort left comparator)
                 (stable-sort right comparator)
                 comparator)))))
)
