;;; Arl Standard Library - Functional Helpers

(module functional
  (export map mapcat filter remove reduce foldl foldr every? any? complement compose partial
          curry juxt constantly iterate iterate-until memoize
          for-each count map-indexed group-by frequencies)

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

  ;;' @section Core Higher-Order Functions
  ;;' The fundamental building blocks for functional programming in Arl.
  ;;' These functions operate on lists and follow the convention of taking
  ;;' the function argument first, then the data argument.

  ;;' @description Apply function to each element. With one list, maps fn over it.
  ;;' With multiple lists, maps fn over corresponding elements (like Scheme's map).
  ;;' @param fn Function to apply to each element (or corresponding elements)
  ;;' @param lst First list to transform
  ;;' @param rest Additional lists for multi-list mapping
  ;;' @examples
  ;;' (map (lambda (x) (* x 2)) (list 1 2 3))  ; => (2 4 6)
  ;;' (map + (list 1 2 3) (list 4 5 6))         ; => (5 7 9)
  ;;' @assert
  ;;' (assert-equal (list 2 4 6) (map (lambda (x) (* x 2)) (list 1 2 3)))
  ;;' (assert-equal (list 5 7 9) (map + (list 1 2 3) (list 4 5 6)))
  ;;' @seealso filter, mapcat, reduce
  (define map
    (lambda (fn lst . rest)
      (if (= (length rest) 0)
        (lapply (_as-list lst) fn)
        (do.call Map (cons fn (cons (_as-list lst) (lapply rest _as-list)))))))

  ;;' @description Map then concatenate results. Uses do.call(c, ...) for O(n) concatenation.
  ;;' @param fn Function that returns a list for each element
  ;;' @param lst List to transform
  ;;' @examples
  ;;' (mapcat (lambda (x) (list x (* x 10))) (list 1 2 3))  ; => (1 10 2 20 3 30)
  ;;' (mapcat (lambda (x) (list x x)) (list 1 2 3))          ; => (1 1 2 2 3 3)
  ;;' @assert
  ;;' (assert-equal (list 1 10 2 20 3 30) (mapcat (lambda (x) (list x (* x 10))) (list 1 2 3)))
  ;;' @seealso map, flatten
  (define mapcat
    (lambda (fn lst)
      (define results (map fn lst))
      (if (null? results)
        (list)
        (do.call c (lapply results (lambda (x) (_as-list x)))))))

  ;;' @description Filter items by predicate.
  ;;' @param pred Predicate function; elements where it returns #t are kept
  ;;' @param lst List to filter
  ;;' @examples
  ;;' (filter (lambda (x) (> x 2)) (list 1 2 3 4 5))  ; => (3 4 5)
  ;;' (filter even? (list 1 2 3 4 5 6))                 ; => (2 4 6)
  ;;' (filter string? (list 1 "a" 2 "b"))               ; => ("a" "b")
  ;;' @assert
  ;;' (assert-equal (list 3 4 5) (filter (lambda (x) (> x 2)) (list 1 2 3 4 5)))
  ;;' @seealso remove, take-while, drop-while
  (define filter
    (lambda (pred lst)
      (Filter pred (_as-list lst))))

  ;;' @description Remove items where predicate is true.
  ;;' @param pred Predicate function; elements where it returns #t are removed
  ;;' @param lst List to filter
  ;;' @examples
  ;;' (remove even? (list 1 2 3 4 5 6))                 ; => (1 3 5)
  ;;' (remove (lambda (x) (> x 3)) (list 1 2 3 4 5))   ; => (1 2 3)
  ;;' @assert
  ;;' (assert-equal (list 1 3 5) (remove even? (list 1 2 3 4 5 6)))
  ;;' @seealso filter
  (define remove
    (lambda (pred lst)
      (Filter (lambda (x) (not (isTRUE (pred x)))) (_as-list lst))))

  ;;' @description Reduce list with function.
  ;;' @param fn Binary function taking accumulator and current element
  ;;' @param lst List to reduce
  ;;' @param rest Optional initial accumulator value
  ;;' @examples
  ;;' (reduce + (list 1 2 3 4))       ; => 10
  ;;' (reduce * (list 1 2 3 4))       ; => 24
  ;;' (reduce + (list 1 2 3) 100)     ; => 106 (with initial value)
  ;;' (reduce string-append (list "a" "b" "c"))  ; => "abc"
  ;;' @assert
  ;;' (assert-equal 10 (reduce + (list 1 2 3 4)))
  ;;' (assert-equal 24 (reduce * (list 1 2 3 4)))
  ;;' (assert-equal 106 (reduce + (list 1 2 3) 100))
  ;;' @seealso foldl, foldr, map
  (define reduce
    (lambda (fn lst . rest)
      (if (= (length rest) 0)
        (Reduce fn (_as-list lst))
        (Reduce fn (_as-list lst) :init (car rest)))))

  ;;' @description Left fold alias for reduce. Applies fn as (fn acc elem).
  ;;' @param fn Binary function taking accumulator and current element
  ;;' @param lst List to fold over
  ;;' @param rest Optional initial accumulator value
  ;;' @examples
  ;;' (foldl + (list 1 2 3 4))        ; => 10
  ;;' (foldl - (list 1 2 3) 10)       ; => 4 (10 - 1 - 2 - 3)
  ;;' (foldl (lambda (acc x) (cons x acc)) (list 1 2 3) (list))  ; => (3 2 1)
  ;;' @assert
  ;;' (assert-equal 10 (foldl + (list 1 2 3 4)))
  ;;' (assert-equal 4 (foldl - (list 1 2 3) 10))
  ;;' (assert-equal (list 3 2 1) (foldl (lambda (acc x) (cons x acc)) (list 1 2 3) (list)))
  ;;' @seealso foldr, reduce
  (define foldl
    (lambda (fn lst . rest)
      (if (= (length rest) 0)
        (Reduce fn (_as-list lst))
        (Reduce fn (_as-list lst) :init (car rest)))))

  ;;' @description Right fold. Applies fn as (fn elem acc).
  ;;' @param fn Binary function taking current element and accumulator
  ;;' @param lst List to fold over from the right
  ;;' @param rest Optional initial accumulator value
  ;;' @examples
  ;;' (foldr + (list 1 2 3 4))        ; => 10
  ;;' (foldr - (list 1 2 3) 10)       ; => -8 (1 - (2 - (3 - 10)))
  ;;' (foldr cons (list 1 2 3) (list)) ; => (1 2 3) (preserves order)
  ;;' @assert
  ;;' (assert-equal 10 (foldr + (list 1 2 3 4)))
  ;;' (assert-equal -8 (foldr - (list 1 2 3) 10))
  ;;' (assert-equal (list 1 2 3) (foldr cons (list 1 2 3) (list)))
  ;;' @seealso foldl, reduce
  (define foldr
    (lambda (fn lst . rest)
      (if (= (length rest) 0)
        (Reduce fn (_as-list lst) :right #t)
        (Reduce fn (_as-list lst) :init (car rest) :right #t))))

  ;;' @description Return #t if predicate true for all items.
  ;;' @param pred Predicate function to test each element
  ;;' @param lst List to check
  ;;' @examples
  ;;' (every? positive? (list 1 2 3))       ; => #t
  ;;' (every? even? (list 2 4 6))           ; => #t
  ;;' (every? even? (list 2 3 6))           ; => #f
  ;;' (every? string? (list "a" "b" "c"))   ; => #t
  ;;' @assert
  ;;' (assert-true (every? positive? (list 1 2 3)))
  ;;' (assert-false (every? even? (list 2 3 6)))
  ;;' @seealso any?, filter
  (define every?
    (lambda (pred lst)
      (all (vapply (_as-list lst) (lambda (x) (isTRUE (pred x))) (logical 1)))))

  ;;' @description Return #t if predicate true for any item.
  ;;' @param pred Predicate function to test each element
  ;;' @param lst List to check
  ;;' @examples
  ;;' (any? even? (list 1 3 4 5))           ; => #t
  ;;' (any? negative? (list 1 2 3))         ; => #f
  ;;' (any? string? (list 1 "a" 2))         ; => #t
  ;;' @assert
  ;;' (assert-true (any? even? (list 1 3 4 5)))
  ;;' (assert-false (any? negative? (list 1 2 3)))
  ;;' @seealso every?, filter
  (define any?
    (lambda (pred lst)
      (any (vapply (_as-list lst) (lambda (x) (isTRUE (pred x))) (logical 1)))))

  ;;' @description Negate predicate.
  ;;' @param pred Predicate function to negate
  ;;' @examples
  ;;' (define not-even? (complement even?))
  ;;' (not-even? 3)                          ; => #t
  ;;' (not-even? 4)                          ; => #f
  ;;' (filter (complement null?) (list 1 #nil 2 #nil 3))  ; => (1 2 3)
  ;;' @assert
  ;;' (assert-true ((complement even?) 3))
  ;;' (assert-false ((complement even?) 4))
  ;;' @seealso not, filter
  (define complement
    (lambda (pred)
      (lambda (. args)
        (not (do.call pred args)))))

  ;;' @description Compose two functions.
  ;;' @param f Outer function (applied second)
  ;;' @param g Inner function (applied first)
  ;;' @examples
  ;;' (define add1-then-double (compose (lambda (x) (* x 2)) (lambda (x) (+ x 1))))
  ;;' (add1-then-double 3)                   ; => 8  ((3+1)*2)
  ;;' (define abs-sum (compose abs +))
  ;;' (abs-sum -3 -4)                        ; => 7
  ;;' @assert
  ;;' (assert-equal 8 ((compose (lambda (x) (* x 2)) (lambda (x) (+ x 1))) 3))
  ;;' @seealso partial
  (define compose
    (lambda (f g)
      (lambda (. args)
        (f (do.call g args)))))

  ;;' @description Partially apply function.
  ;;' @param fn Function to partially apply
  ;;' @param captured Arguments to bind to fn's first positions
  ;;' @examples
  ;;' (define add5 (partial + 5))
  ;;' (add5 3)                               ; => 8
  ;;' (define greet (partial string-append "Hello, "))
  ;;' (greet "world")                        ; => "Hello, world"
  ;;' @assert
  ;;' (assert-equal 8 ((partial + 5) 3))
  ;;' @seealso compose
  (define partial
    (lambda (fn . captured)
      (lambda (. more)
        (do.call fn (append captured more)))))

  ;; ============================================================================
  ;; Advanced Functional Programming
  ;; ============================================================================

  ;;' @section Advanced Functional Programming
  ;;' Higher-order combinators for advanced functional patterns including
  ;;' currying, juxtaposition, memoization, and iteration.

  ;;' @description Curry a function - enables partial application
  ;;' with optional initial arguments.
  ;;' @param fn Function to curry (must have fixed arity)
  ;;' @param initial-args Optional arguments to apply immediately
  ;;' @examples
  ;;' (define add (curry (lambda (a b) (+ a b))))
  ;;' (define add5 (add 5))
  ;;' (add5 3)                               ; => 8
  ;;' (define multiply (curry (lambda (a b) (* a b))))
  ;;' ((multiply 3) 4)                       ; => 12
  ;;' @assert
  ;;' (assert-equal 8 ((curry (lambda (a b) (+ a b)) 5) 3))
  ;;' @seealso partial
  (define curry
    (lambda (fn . initial-args)
      (define arity (length (formals fn)))
      (define curry-helper
        (lambda (collected-args)
          (lambda (. new-args)
            (define all-args (append collected-args new-args))
            (if (>= (length all-args) arity)
              (do.call fn all-args)
              (curry-helper all-args)))))
      (curry-helper initial-args)))

  ;;' @description Juxtaposition - apply multiple functions to same args, return list of results.
  ;;' @param fns Functions to apply to the arguments
  ;;' @examples
  ;;' (define stats (juxt min max mean))
  ;;' (stats (c 1 2 3 4 5))                  ; => (1 5 3)
  ;;' (define first-and-last (juxt car last))
  ;;' (first-and-last (list 1 2 3))          ; => (1 3)
  ;;' @assert
  ;;' (assert-equal (list 1 5 3) ((juxt min max mean) (c 1 2 3 4 5)))
  ;;' (assert-equal (list 1 3) ((juxt car last) (list 1 2 3)))
  ;;' @seealso map, compose
  (define juxt
    (lambda (. fns)
      (lambda (. args)
        (map (lambda (f) (do.call f args)) fns))))

  ;;' @description Return a function that always returns the given value.
  ;;' @param value Value to return regardless of arguments
  ;;' @examples
  ;;' (define always-42 (constantly 42))
  ;;' (always-42)                             ; => 42
  ;;' (always-42 "ignored" "args")            ; => 42
  ;;' (map (constantly 0) (list 1 2 3))       ; => (0 0 0)
  ;;' @assert
  ;;' (assert-equal 42 ((constantly 42)))
  ;;' (assert-equal (list 0 0 0) (map (constantly 0) (list 1 2 3)))
  (define constantly
    (lambda (value)
      (lambda (. args)
        value)))

  ;;' @description Apply function n times to initial value.
  ;;' @param fn Unary function to apply repeatedly
  ;;' @param n Number of times to apply fn
  ;;' @param init Starting value
  ;;' @examples
  ;;' (iterate (lambda (x) (* x 2)) 5 1)     ; => 32  (1 -> 2 -> 4 -> 8 -> 16 -> 32)
  ;;' (iterate inc 3 0)                       ; => 3   (0 -> 1 -> 2 -> 3)
  ;;' (iterate (lambda (x) (* x x)) 2 2)     ; => 16  (2 -> 4 -> 16)
  ;;' @assert
  ;;' (assert-equal 32 (iterate (lambda (x) (* x 2)) 5 1))
  ;;' (assert-equal 3 (iterate inc 3 0))
  ;;' (assert-equal 16 (iterate (lambda (x) (* x x)) 2 2))
  ;;' @seealso iterate-until, reduce
  (define iterate
    (lambda (fn n init)
      (if (<= n 0)
        init
        (iterate fn (- n 1) (fn init)))))

  ;;' @description Apply function starting from init, collecting
  ;;' values until predicate becomes true. Includes all values where
  ;;' predicate is false, stops when next value would satisfy predicate.
  ;;' @param fn Unary function to apply repeatedly
  ;;' @param init Starting value
  ;;' @param pred Predicate that signals when to stop collecting
  ;;' @examples
  ;;' (iterate-until (lambda (x) (* x 2)) 1 (lambda (x) (> x 100)))
  ;;'   ; => (1 2 4 8 16 32 64)  (stops before 128 which exceeds 100)
  ;;' (iterate-until inc 0 (lambda (x) (>= x 5)))
  ;;'   ; => (0 1 2 3 4)
  ;;' @assert
  ;;' (assert-equal (list 1 2 4 8 16 32 64) (iterate-until (lambda (x) (* x 2)) 1 (lambda (x) (> x 100))))
  ;;' (assert-equal (list 0 1 2 3 4) (iterate-until inc 0 (lambda (x) (>= x 5))))
  ;;' @seealso iterate, take-while
  (define iterate-until
    (lambda (fn init pred)
      (define current init)
      (define acc (list))
      (define done #f)
      (while (not done)
        (begin
          (set! acc (cons current acc))
          (define next (fn current))
          (if (pred next)
            (set! done #t)
            (set! current next))))
      (rev acc)))

  ;;' @description Memoize function - cache results for previously seen arguments.
  ;;' @param fn Function whose results to cache by arguments
  ;;' @examples
  ;;' (define slow-fib (lambda (n) (if (< n 2) n (+ (slow-fib (- n 1)) (slow-fib (- n 2))))))
  ;;' (define fast-fib (memoize slow-fib))
  ;;' (fast-fib 10)                          ; => 55
  ;;' (fast-fib 10)                          ; => 55 (cached, instant)
  ;;' @assert
  ;;' (assert-equal 4 ((memoize (lambda (x) (* x x))) 2))
  ;;' (assert-equal 4 ((memoize (lambda (x) (* x x))) 2))
  (define memoize
    (lambda (fn)
      (define cache (r-call "new.env" (list)))
      (lambda (. args)
        (define key (r-call "paste" (list args :collapse "|")))
        (if (r-call "exists" (list key :envir cache))
          (r-call "get" (list key :envir cache))
          (begin
            (define result (do.call fn args))
            (r-call "assign" (list key result :envir cache))
            result)))))

  ;; ============================================================================
  ;; Iteration and Counting
  ;; ============================================================================

  ;;' @section Iteration and Counting
  ;;' Functions for side-effecting iteration and counting over sequences.

  ;;' @description Apply function to each element for side effects.
  ;;' Returns #nil. Supports multiple lists like multi-list map.
  ;;' @param fn Function to apply for side effects
  ;;' @param lst First list
  ;;' @param rest Additional lists for multi-list iteration
  ;;' @examples
  ;;' (for-each print '(1 2 3))  ; prints 1 2 3, returns #nil
  ;;' @assert
  ;;' (assert-equal #nil (for-each (lambda (x) x) (list 1 2 3)))
  ;;' @seealso map
  (define for-each
    (lambda (fn lst . rest)
      (if (= (length rest) 0)
        (lapply (_as-list lst) fn)
        (do.call Map (cons fn (cons (_as-list lst) (lapply rest _as-list)))))
      #nil))

  ;;' @description Count elements matching predicate.
  ;;' @param pred Predicate function to test each element
  ;;' @param lst List to count matching elements in
  ;;' @examples
  ;;' (count even? (list 1 2 3 4))  ; => 2
  ;;' (count even? (list))           ; => 0
  ;;' @assert
  ;;' (assert-equal 2 (count even? (list 1 2 3 4)))
  ;;' (assert-equal 0 (count even? (list)))
  ;;' @seealso filter, every?, any?
  (define count
    (lambda (pred lst)
      (sum (vapply (_as-list lst) (lambda (x) (isTRUE (pred x))) (logical 1)))))

  ;;' @description Map with index. Calls (fn index element) for each element,
  ;;' with 0-based indices.
  ;;' @param fn Function taking index and element
  ;;' @param lst List to map over
  ;;' @examples
  ;;' (map-indexed list '(a b c))  ; => ((0 a) (1 b) (2 c))
  ;;' @assert
  ;;' (assert-equal (list (list 0 'a) (list 1 'b) (list 2 'c)) (map-indexed list (list 'a 'b 'c)))
  ;;' @seealso map
  (define map-indexed
    (lambda (fn lst)
      (define items (_as-list lst))
      (define n (length items))
      (if (= n 0)
        (list)
        (lapply (seq_len n)
          (lambda (i) (fn (- i 1) (r-call "[[" (list items i))))))))

  ;;' @description Group elements by key function. Returns a dict mapping
  ;;' each key (as string) to a list of matching elements.
  ;;' @param key-fn Function to compute the grouping key for each element
  ;;' @param lst List to group
  ;;' @examples
  ;;' (group-by even? '(1 2 3 4))  ; => dict with TRUE->(2 4), FALSE->(1 3)
  ;;' @seealso frequencies, filter
  (define group-by
    (lambda (key-fn lst)
      (define items (_as-list lst))
      (define n (length items))
      (define result (r-call "new.env" (list :hash #t :parent (r-call "emptyenv" (list)))))
      (r-call "class<-" (list result (c (list "arl_dict") (r-call "class" (list result)))))
      (define keys-order (list))
      (define i 1)
      (while (<= i n)
        (begin
          (define elem ([[ items i))
          (define key (r-call "as.character" (list (key-fn elem))))
          (if (r-call "exists" (list key :envir result :inherits #f))
            (begin
              (define current (r-call "get" (list key :envir result :inherits #f)))
              (r-call "assign" (list key (c current (list elem)) :envir result)))
            (begin
              (set! keys-order (c keys-order (list key)))
              (r-call "assign" (list key (list elem) :envir result))))
          (set! i (+ i 1))))
      (r-call "assign" (list ".__keys" (r-call "as.character" (list keys-order)) :envir result))
      result))

  ;;' @description Count occurrences of each element. Returns a dict mapping
  ;;' each element (as string) to its count.
  ;;' @param lst List to count element frequencies in
  ;;' @examples
  ;;' (frequencies '(a b a c b a))  ; => dict with a->3 b->2 c->1
  ;;' @seealso group-by, count
  (define frequencies
    (lambda (lst)
      (define items (_as-list lst))
      (define n (length items))
      (define result (r-call "new.env" (list :hash #t :parent (r-call "emptyenv" (list)))))
      (r-call "class<-" (list result (c (list "arl_dict") (r-call "class" (list result)))))
      (define keys-order (list))
      (define i 1)
      (while (<= i n)
        (begin
          (define elem ([[ items i))
          (define key (r-call "as.character" (list elem)))
          (if (r-call "exists" (list key :envir result :inherits #f))
            (begin
              (define current (r-call "get" (list key :envir result :inherits #f)))
              (r-call "assign" (list key (+ current 1) :envir result)))
            (begin
              (set! keys-order (c keys-order (list key)))
              (r-call "assign" (list key 1 :envir result))))
          (set! i (+ i 1))))
      (r-call "assign" (list ".__keys" (r-call "as.character" (list keys-order)) :envir result))
      result))

)
