;;; Arl Standard Library - Sequence Helpers

(module sequences
  (export take drop take-while drop-while
          partition flatten repeatedly repeat zip
          member contains?
          length= length> length<
          find distinct split-at split-with interpose partition-by)

  (import types :refer (list?))
  (import core :refer (r-call))
  (import equality :refer (equal?))


  ;;' @section Sequence Helpers
  ;;' Higher-level operations for slicing, chunking, and transforming
  ;;' sequences. These complement the core list operations in the `list`
  ;;' module with lazy-style idioms common in functional programming.

  ;;' @description Take first n items.
  ;;' @param n Number of elements to take
  ;;' @param lst Source list
  ;;' @examples
  ;;' (take 3 (list 1 2 3 4 5))  ; => (1 2 3)
  ;;' (take 0 (list 1 2 3))      ; => ()
  ;;' (take 10 (list 1 2))       ; => (1 2)
  ;;' @assert
  ;;' (assert-equal (list 1 2 3) (take 3 (list 1 2 3 4 5)))
  ;;' (assert-equal (list) (take 0 (list 1 2 3)))
  ;;' (assert-equal (list 1 2) (take 10 (list 1 2)))
  ;;' @seealso drop, take-while
  (define take
    (lambda (n lst)
      (define items (_as-list lst))
      (if (<= n 0)
        (list)
        (r-call "[" (list items (seq_len (min n (length items))))))))

  ;;' @description Drop first n items.
  ;;' @param n Number of elements to skip
  ;;' @param lst Source list
  ;;' @examples
  ;;' (drop 2 (list 1 2 3 4 5))  ; => (3 4 5)
  ;;' (drop 0 (list 1 2 3))      ; => (1 2 3)
  ;;' (drop 10 (list 1 2))       ; => ()
  ;;' @assert
  ;;' (assert-equal (list 3 4 5) (drop 2 (list 1 2 3 4 5)))
  ;;' (assert-equal (list 1 2 3) (drop 0 (list 1 2 3)))
  ;;' (assert-equal (list) (drop 10 (list 1 2)))
  ;;' @seealso take, drop-while
  (define drop
    (lambda (n lst)
      (define items (_as-list lst))
      (if (<= n 0)
        items
        (if (>= n (length items))
          (list)
          (r-call "[" (list items (seq (+ n 1) (length items))))))))

  ;;' @description Take items while predicate is true.
  ;;' @param pred Predicate function; taking stops at first #f
  ;;' @param lst Source list
  ;;' @examples
  ;;' (take-while odd? (list 1 3 5 2 4))   ; => (1 3 5)
  ;;' (take-while even? (list 1 2 3))       ; => ()
  ;;' @assert
  ;;' (assert-equal (list 1 3 5) (take-while odd? (list 1 3 5 2 4)))
  ;;' (assert-equal (list) (take-while even? (list 1 2 3)))
  ;;' @seealso take, drop-while, filter
  (define take-while
    (lambda (pred lst)
      (define items (_as-list lst))
      (define n (length items))
      (define i 1)
      (while (and (<= i n) (isTRUE (pred ([[ items i))))
        (set! i (+ i 1)))
      (if (= i 1) (list)
        (r-call "[" (list items (seq_len (- i 1)))))))

  ;;' @description Drop items while predicate is true.
  ;;' @param pred Predicate function; dropping stops at first #f
  ;;' @param lst Source list
  ;;' @examples
  ;;' (drop-while odd? (list 1 3 5 2 4))   ; => (2 4)
  ;;' (drop-while even? (list 1 2 3))       ; => (1 2 3)
  ;;' @assert
  ;;' (assert-equal (list 2 4) (drop-while odd? (list 1 3 5 2 4)))
  ;;' (assert-equal (list 1 2 3) (drop-while even? (list 1 2 3)))
  ;;' @seealso drop, take-while, filter
  (define drop-while
    (lambda (pred lst)
      (define items (_as-list lst))
      (define n (length items))
      (define i 1)
      (while (and (<= i n) (isTRUE (pred ([[ items i))))
        (set! i (+ i 1)))
      (if (> i n) (list)
        (r-call "[" (list items (seq i n))))))

  ;;' @description Split list into chunks of size n.
  ;;' @param n Chunk size
  ;;' @param lst Source list
  ;;' @param step Offset between chunk starts (default n, i.e. no overlap)
  ;;' @examples
  ;;' (partition 2 (list 1 2 3 4 5 6))       ; => ((1 2) (3 4) (5 6))
  ;;' (partition 3 (list 1 2 3 4 5 6) 2)     ; => ((1 2 3) (3 4 5))
  ;;' (partition 2 (list 1 2 3 4 5))          ; => ((1 2) (3 4))
  ;;' @assert
  ;;' (assert-equal (list (list 1 2) (list 3 4) (list 5 6)) (partition 2 (list 1 2 3 4 5 6)))
  ;;' (assert-equal (list (list 1 2) (list 3 4)) (partition 2 (list 1 2 3 4 5)))
  ;;' @seealso take, drop
  (define partition
    (lambda (n lst (step n))
      (define items (_as-list lst))
      (if (<= n 0)
        (stop "partition requires positive n and step")
        (if (<= step 0)
          (stop "partition requires positive n and step")
          (if (= (length items) 0)
            (list)
            (begin
              (define acc (list))
              (define idx 1)
              (define n-items (length items))
              ;; Build in reverse with cons, then rev — O(n) total
              (while (<= idx n-items)
                (define end (+ idx (- n 1)))
                (if (<= end n-items)
                  (set! acc (cons (r-call "[" (list items (seq idx end))) acc)))
                (set! idx (+ idx step)))
              (rev acc)))))))

  ;;' @description Flatten nested lists.
  ;;' @param lst Possibly nested list to flatten into a single-level list
  ;;' @examples
  ;;' (flatten (list 1 (list 2 3) (list 4 (list 5))))  ; => (1 2 3 4 5)
  ;;' (flatten (list 1 2 3))                            ; => (1 2 3)
  ;;' @assert
  ;;' (assert-equal (list 1 2 3 4 5) (flatten (list 1 (list 2 3) (list 4 (list 5)))))
  ;;' (assert-equal (list 1 2 3) (flatten (list 1 2 3)))
  ;;' @seealso mapcat, append
  (define flatten
    (lambda (lst)
      (define items (_as-list lst))
      (define n (length items))
      (if (= n 0) (list)
        (begin
          (define parts (list))
          (define i 1)
          (while (<= i n)
            (define item ([[ items i))
            (if (list? item)
              (set! parts (cons (flatten item) parts))
              (set! parts (cons (list item) parts)))
            (set! i (+ i 1)))
          (do.call c (rev parts))))))

  ;;' @description Call fn n times collecting results.
  ;;' @param n Number of times to call fn
  ;;' @param fn Zero-argument function to call repeatedly
  ;;' @examples
  ;;' (repeatedly 3 (lambda () 42))  ; => (42 42 42)
  ;;' @assert
  ;;' (assert-equal (list 42 42 42) (repeatedly 3 (lambda () 42)))
  ;;' @seealso repeat, map
  (define repeatedly
    (lambda (n fn)
      (lapply (seq_len n) (lambda (i) (fn)))))

  ;;' @description Repeat value n times.
  ;;' @param n Number of repetitions
  ;;' @param value Value to repeat
  ;;' @examples
  ;;' (repeat 4 'x)   ; => (x x x x)
  ;;' (repeat 3 0)     ; => (0 0 0)
  ;;' @assert
  ;;' (assert-equal (list 'x 'x 'x 'x) (repeat 4 'x))
  ;;' (assert-equal (list 0 0 0) (repeat 3 0))
  ;;' @seealso repeatedly
  (define repeat
    (lambda (n value)
      (replicate n value :simplify #f)))

  ;;' @description Zip lists into list of tuples.
  ;;' @param lists Two or more lists to interleave element-wise
  ;;' @examples
  ;;' (zip (list 1 2 3) (list 'a 'b 'c))          ; => ((1 a) (2 b) (3 c))
  ;;' (zip (list 1 2) (list 'a 'b) (list 'x 'y))  ; => ((1 a x) (2 b y))
  ;;' @assert
  ;;' (assert-equal (list (list 1 'a) (list 2 'b) (list 3 'c)) (zip (list 1 2 3) (list 'a 'b 'c)))
  ;;' @seealso map
  (define zip
    (lambda (. lists)
      (if (= (length lists) 0)
        (list)
        (begin
          (define items (lapply lists _as-list))
          (define min-len (min (vapply items length (integer 1))))
          (if (= min-len 0)
            (list)
            (begin
              ;; Collect ith elements from lists.
              (define zip-at
                (lambda (i)
                  (lapply items (lambda (lst) (r-call "[[" (list lst i))))))
              (lapply (seq_len min-len) zip-at)))))))

  ;;' @description Return sublist of lst starting at first element
  ;;' equal to x, or #f if not found. Uses equal? by default;
  ;;' pass :use-identical #t for R's identical().
  ;;' @param x Value to search for
  ;;' @param lst List to search in
  ;;' @param use-identical Use R's identical() instead of equal? (default #f)
  ;;' @examples
  ;;' (member 3 (list 1 2 3 4 5))   ; => (3 4 5)
  ;;' (member 9 (list 1 2 3))       ; => #f
  ;;' @assert
  ;;' (assert-equal (list 3 4 5) (member 3 (list 1 2 3 4 5)))
  ;;' (assert-false (member 9 (list 1 2 3)))
  ;;' @seealso contains?
  (define member
    (lambda (x lst (use-identical #f))
      (define items (_as-list lst))
      (define n (length items))
      (define result #f)
      (define i 1)
      (while (and (<= i n) (not result))
        (if (if use-identical
                (identical ([[ items i) x)
                (equal? ([[ items i) x))
          (set! result (r-call "[" (list items (seq i n))))
          #nil)
        (set! i (+ i 1)))
      result))

  ;;' @description Return #t if lst contains element equal to x.
  ;;' @param x Value to search for
  ;;' @param lst List to search in
  ;;' @examples
  ;;' (contains? 2 (list 1 2 3))   ; => #t
  ;;' (contains? 9 (list 1 2 3))   ; => #f
  ;;' @assert
  ;;' (assert-true (contains? 2 (list 1 2 3)))
  ;;' (assert-false (contains? 9 (list 1 2 3)))
  ;;' @seealso member
  (define contains?
    (lambda (x lst)
      (if (member x lst) #t #f)))

  ;;' @section Length Predicates
  ;;' Efficient length comparisons that short-circuit when possible.

  ;;' @description Return #t if length of x equals n.
  ;;' @param x Sequence to measure
  ;;' @param n Expected length
  ;;' @examples
  ;;' (length= (list 1 2 3) 3)  ; => #t
  ;;' (length= (list 1 2) 3)    ; => #f
  ;;' @assert
  ;;' (assert-true (length= (list 1 2 3) 3))
  ;;' (assert-false (length= (list 1 2) 3))
  ;;' @seealso length>, length<, empty?
  (define length=
    (lambda (x n)
      (isTRUE (== (length x) n))))

  ;;' @description Return #t if length of x is greater than n.
  ;;' @param x Sequence to measure
  ;;' @param n Length to compare against
  ;;' @examples
  ;;' (length> (list 1 2 3) 2)  ; => #t
  ;;' (length> (list 1 2) 3)    ; => #f
  ;;' @assert
  ;;' (assert-true (length> (list 1 2 3) 2))
  ;;' (assert-false (length> (list 1 2) 3))
  ;;' @seealso length=, length<
  (define length>
    (lambda (x n)
      (isTRUE (> (length x) n))))

  ;;' @description Return #t if length of x is less than n.
  ;;' @param x Sequence to measure
  ;;' @param n Length to compare against
  ;;' @examples
  ;;' (length< (list 1 2) 3)    ; => #t
  ;;' (length< (list 1 2 3) 2)  ; => #f
  ;;' @assert
  ;;' (assert-true (length< (list 1 2) 3))
  ;;' (assert-false (length< (list 1 2 3) 2))
  ;;' @seealso length=, length>
  (define length<
    (lambda (x n)
      (isTRUE (< (length x) n))))

  ;; ============================================================================
  ;; Search and Deduplication
  ;; ============================================================================

  ;;' @section Search and Deduplication
  ;;' Functions for finding elements and removing duplicates.

  ;;' @description Return first element matching predicate, or #f if not found.
  ;;' @param pred Predicate function to test each element
  ;;' @param lst List to search
  ;;' @examples
  ;;' (find even? (list 1 3 4 5))  ; => 4
  ;;' (find even? (list 1 3 5))    ; => #f
  ;;' @assert
  ;;' (assert-equal 4 (find even? (list 1 3 4 5)))
  ;;' (assert-false (find even? (list 1 3 5)))
  ;;' @seealso member, contains?, filter
  (define find
    (lambda (pred lst)
      (define items (_as-list lst))
      (define n (length items))
      (define result #f)
      (define i 1)
      (while (and (<= i n) (not result))
        (begin
          (define elem ([[ items i))
          (if (isTRUE (pred elem))
            (set! result elem))
          (set! i (+ i 1))))
      result))

  ;;' @description Remove duplicates preserving first-occurrence order.
  ;;' Uses R's identical() for comparison.
  ;;' @param lst List to deduplicate
  ;;' @examples
  ;;' (distinct (list 1 2 1 3 2))  ; => (1 2 3)
  ;;' @assert
  ;;' (assert-equal (list 1 2 3) (distinct (list 1 2 1 3 2)))
  ;;' @seealso filter
  (define distinct
    (lambda (lst)
      (define items (_as-list lst))
      (define n (length items))
      (if (= n 0) (list)
        (begin
          ;; Use an env as a seen-set for O(n) dedup via hashing.
          ;; Keys are type-prefixed to avoid collisions (e.g. 1 vs "1").
          (define seen (r-call "new.env" (list :hash #t :parent (emptyenv))))
          (define acc (list))
          (define i 1)
          ;; Iterate forward with cons, rev at end — preserves first-occurrence order, O(n) total
          (while (<= i n)
            (define elem ([[ items i))
            (define key (r-call "paste0" (list (r-call "typeof" (list elem)) ":" (r-call "as.character" (list elem)))))
            (if (not (r-call "exists" (list key :envir seen :inherits #f)))
              (begin
                (r-call "assign" (list key #t :envir seen))
                (set! acc (cons elem acc))))
            (set! i (+ i 1)))
          (rev acc)))))

  ;; ============================================================================
  ;; Splitting
  ;; ============================================================================

  ;;' @section Splitting
  ;;' Functions for splitting sequences at positions or predicates.

  ;;' @description Split list into two parts at index n.
  ;;' Returns a list of (take n lst) and (drop n lst).
  ;;' @param n Index to split at
  ;;' @param lst List to split
  ;;' @examples
  ;;' (split-at 2 (list 'a 'b 'c 'd))  ; => ((a b) (c d))
  ;;' @assert
  ;;' (assert-equal (list (list 1 2) (list 3 4)) (split-at 2 (list 1 2 3 4)))
  ;;' @seealso split-with, take, drop
  (define split-at
    (lambda (n lst)
      (list (take n lst) (drop n lst))))

  ;;' @description Split list into two parts: elements satisfying pred at the
  ;;' start, and the rest. Returns (take-while pred lst) and (drop-while pred lst).
  ;;' @param pred Predicate function
  ;;' @param lst List to split
  ;;' @examples
  ;;' (split-with even? (list 2 4 1 3))  ; => ((2 4) (1 3))
  ;;' @assert
  ;;' (assert-equal (list (list 2 4) (list 1 3)) (split-with even? (list 2 4 1 3)))
  ;;' @seealso split-at, take-while, drop-while
  (define split-with
    (lambda (pred lst)
      (list (take-while pred lst) (drop-while pred lst))))

  ;;' @description Insert separator between every element of list.
  ;;' @param sep Value to insert between elements
  ;;' @param lst List to interpose into
  ;;' @examples
  ;;' (interpose 0 (list 1 2 3))  ; => (1 0 2 0 3)
  ;;' @assert
  ;;' (assert-equal (list 1 0 2 0 3) (interpose 0 (list 1 2 3)))
  ;;' @seealso flatten
  (define interpose
    (lambda (sep lst)
      (define items (_as-list lst))
      (define n (length items))
      (if (<= n 1)
        items
        (begin
          ;; Build in reverse with cons, then reverse (O(n))
          (define acc (list))
          (define i n)
          (while (>= i 1)
            (set! acc (cons ([[ items i) acc))
            (if (> i 1)
              (set! acc (cons sep acc)))
            (set! i (- i 1)))
          acc))))

  ;;' @description Split list into runs of consecutive elements producing
  ;;' the same key. Different from partition which splits by size.
  ;;' @param fn Key function applied to each element
  ;;' @param lst List to partition
  ;;' @examples
  ;;' (partition-by even? (list 2 4 1 3 6))  ; => ((2 4) (1 3) (6))
  ;;' @assert
  ;;' (assert-equal (list (list 2 4) (list 1 3) (list 6)) (partition-by even? (list 2 4 1 3 6)))
  ;;' @seealso partition, split-with
  (define partition-by
    (lambda (fn lst)
      (define items (_as-list lst))
      (define n (length items))
      (if (= n 0)
        (list)
        (begin
          ;; Build runs in reverse with cons, rev when finishing each run.
          ;; Collect runs in reverse with cons, rev at end — O(n) total.
          (define acc (list))
          (define current-run (list ([[ items 1)))
          (define current-key (fn ([[ items 1)))
          (define i 2)
          (while (<= i n)
            (begin
              (define elem ([[ items i))
              (define key (fn elem))
              (if (identical key current-key)
                (set! current-run (cons elem current-run))
                (begin
                  (set! acc (cons (rev current-run) acc))
                  (set! current-run (list elem))
                  (set! current-key key)))
              (set! i (+ i 1))))
          (rev (cons (rev current-run) acc))))))
)
