;;; Arl Standard Library - Set (Hash Table)

(module set
  (export set set? set-add set-remove set-contains? set-union set-intersection
          set-difference set->list list->set set-size set-map set-filter)

  (import core :refer (r-call error))


;;' @section Set (Hash-Backed)
;;' Sets are mutable collections of unique items backed by R environments.
;;' Any value can be a set element; uniqueness is determined by serialization.

;;; Internal helpers (defined early for use in constructors)

;;' @description Create a new empty set environment.
(define set-new
  (lambda ()
    (define set-obj (r-call "new.env" (list :hash #t :parent (r-call "emptyenv" (list)))))
    (r-call "class<-" (list set-obj (c (list "arl_set") (r-call "class" (list set-obj)))))
    set-obj))

;;' @description Convert value to string key for set storage via serialization.
;;' @param value Value to serialize into a string key
(define set-key
  (lambda (value)
    (define raw (r-call "serialize" (list value #nil :ascii #t)))
    (r-call "paste" (list (r-call "as.integer" (list raw)) :collapse ","))))

;;; Set constructor

;;' @description Create a hash-backed set of unique items.
;;' @param args Items to include in the set, or a single list to convert
;;' @examples
;;' (set 1 2 3)                ; => set of {1, 2, 3}
;;' (set 1 1 2)                ; => set of {1, 2} (duplicates removed)
;;' (set)                      ; => empty set
;;' (set '(a b c))             ; => set from list
;;' @assert
;;' (assert-true (set? (set 1 2 3)))
;;' (assert-true (set-contains? (set 1 2 3) 2))
;;' (assert-false (set-contains? (set 1 2 3) 99))
;;' (assert-true (set? (set)))
;;' (assert-true (set-contains? (set '(a b c)) 'b))
(define set
  (lambda (. args)
    (define set-obj (set-new))
    (define items
      (if (and (= (length args) 1) (or (is.list (car args)) (is.call (car args))))
        (as.list (car args))
        (as.list args)))
    (define add-item
      (lambda (item)
        (begin
          (define key (set-key item))
          (if (not (r-call "exists" (list key :envir set-obj :inherits #f)))
            (r-call "assign" (list key item :envir set-obj))))))
    (r-call "lapply" (list items add-item))
    set-obj))

;;; Type predicate

;;' @description Return #t if x is a set.
;;' @param x Value to test
;;' @examples
;;' (set? (set 1 2))           ; => #t
;;' (set? (list 1 2))          ; => #f
;;' (set? 42)                  ; => #f
;;' @assert
;;' (assert-true (set? (set 1 2)))
;;' (assert-false (set? (list 1 2)))
;;' (assert-false (set? 42))
(define set?
  (lambda (x)
    (and (r-call "is.environment" (list x)) (r-call "inherits" (list x "arl_set")))))

;;; Set operations

;;' @description Add item to set and return set.
;;' @param set Set to add to
;;' @param item Value to add
;;' @examples
;;' (define s (set 1 2))
;;' (set-add s 3)              ; => s (mutated, now {1, 2, 3})
;;' (set-contains? s 3)        ; => #t
;;' @assert
;;' (define _test-sa (set 1 2))
;;' (set-add _test-sa 3)
;;' (assert-true (set-contains? _test-sa 3))
;;' @seealso set-remove
(define set-add
  (lambda (set item)
    (if (not (set? set))
      (error "set-add requires a set")
      (begin
        (define key (set-key item))
        (if (not (r-call "exists" (list key :envir set :inherits #f)))
          (r-call "assign" (list key item :envir set)))
        set))))

;;' @description Remove item from set and return set.
;;' @param set Set to remove from
;;' @param item Value to remove
;;' @examples
;;' (define s (set 1 2 3))
;;' (set-remove s 2)           ; => s (mutated, now {1, 3})
;;' (set-contains? s 2)        ; => #f
;;' @assert
;;' (define _test-sr (set 1 2 3))
;;' (set-remove _test-sr 2)
;;' (assert-false (set-contains? _test-sr 2))
;;' @seealso set-add
(define set-remove
  (lambda (set item)
    (if (not (set? set))
      set
      (begin
        (define key (set-key item))
        (if (r-call "exists" (list key :envir set :inherits #f))
          (r-call "rm" (list :list (c key) :envir set :inherits #f)))
        set))))

;;' @description Return #t if set contains item.
;;' @param set Set to search in
;;' @param item Value to look for
;;' @examples
;;' (define s (set 1 2 3))
;;' (set-contains? s 2)        ; => #t
;;' (set-contains? s 99)       ; => #f
;;' @assert
;;' (assert-true (set-contains? (set 1 2 3) 2))
;;' (assert-false (set-contains? (set 1 2 3) 99))
(define set-contains?
  (lambda (set item)
    (if (not (set? set))
      #f
      (begin
        (define key (set-key item))
        (r-call "exists" (list key :envir set :inherits #f))))))

;;' @description Return union of two sets.
;;' @param a First set
;;' @param b Second set
;;' @examples
;;' (define a (set 1 2 3))
;;' (define b (set 3 4 5))
;;' (define u (set-union a b))
;;' (set-contains? u 1)        ; => #t
;;' (set-contains? u 5)        ; => #t
;;' @assert
;;' (define _test-su (set-union (set 1 2 3) (set 3 4 5)))
;;' (assert-true (set-contains? _test-su 1))
;;' (assert-true (set-contains? _test-su 5))
;;' @seealso set-intersection, set-difference
(define set-union
  (lambda (a b)
    (define result (set-new))
    (set-copy-into result a)
    (set-copy-into result b)
    result))

;;' @description Return intersection of two sets.
;;' @param a First set
;;' @param b Second set
;;' @examples
;;' (define a (set 1 2 3))
;;' (define b (set 2 3 4))
;;' (define i (set-intersection a b))
;;' (set-contains? i 2)        ; => #t
;;' (set-contains? i 1)        ; => #f
;;' @assert
;;' (define _test-si (set-intersection (set 1 2 3) (set 2 3 4)))
;;' (assert-true (set-contains? _test-si 2))
;;' (assert-false (set-contains? _test-si 1))
;;' @seealso set-union, set-difference
(define set-intersection
  (lambda (a b)
    (define result (set-new))
    (if (and (set? a) (set? b))
      (begin
        (define keys (r-call "ls" (list :envir a :all.names #t :sorted #f)))
        (define add-if-in-b
          (lambda (key)
            (if (r-call "exists" (list key :envir b :inherits #f))
              (begin
                (define value (r-call "get" (list key :envir a :inherits #f)))
                (r-call "assign" (list key value :envir result))))))
        (r-call "lapply" (list (as.list keys) add-if-in-b))
        result))
    result))

;;' @description Return items in a that are not in b.
;;' @param a Set to take items from
;;' @param b Set whose items are excluded
;;' @examples
;;' (define a (set 1 2 3))
;;' (define b (set 2 3 4))
;;' (define d (set-difference a b))
;;' (set-contains? d 1)        ; => #t
;;' (set-contains? d 2)        ; => #f
;;' @assert
;;' (define _test-sd (set-difference (set 1 2 3) (set 2 3 4)))
;;' (assert-true (set-contains? _test-sd 1))
;;' (assert-false (set-contains? _test-sd 2))
;;' @seealso set-union, set-intersection
(define set-difference
  (lambda (a b)
    (define result (set-new))
    (if (set? a)
      (begin
        (define keys (r-call "ls" (list :envir a :all.names #t :sorted #f)))
        (define add-if-not-in-b
          (lambda (key)
            (if (or (not (set? b)) (not (r-call "exists" (list key :envir b :inherits #f))))
              (begin
                (define value (r-call "get" (list key :envir a :inherits #f)))
                (r-call "assign" (list key value :envir result))))))
        (r-call "lapply" (list (as.list keys) add-if-not-in-b))
        result))
    result))

;;; Internal helper

;;' @description Copy items from source into target (internal helper).
;;' @param target Set to copy items into
;;' @param source Set to copy items from
(define set-copy-into
  (lambda (target source)
    (if (and (set? target) (set? source))
      (begin
        (define keys (r-call "ls" (list :envir source :all.names #t :sorted #f)))
        (define copy-key
          (lambda (key)
            (if (not (r-call "exists" (list key :envir target :inherits #f)))
              (begin
                (define value (r-call "get" (list key :envir source :inherits #f)))
                (r-call "assign" (list key value :envir target))))))
        (r-call "lapply" (list (as.list keys) copy-key)))
    target)))

;; ============================================================================
;; Set Conversion and Transformation
;; ============================================================================

;;' @section Set Conversion and Transformation
;;' Functions for converting between sets and lists, and for transforming sets.

;;' @description Extract set elements as a list.
;;' @param s Set to convert
;;' @examples
;;' (set->list (set 1 2 3))  ; => (1 2 3) (order may vary)
;;' @assert
;;' (assert-true (> (length (set->list (set 1 2 3))) 0))
;;' @seealso list->set
(define set->list
  (lambda (s)
    (if (not (set? s))
      (list)
      (begin
        (define keys (r-call "ls" (list :envir s :all.names #t :sorted #f)))
        (if (= (length keys) 0)
          (list)
          (r-call "as.list" (list (r-call "unname" (list (r-call "mget" (list keys :envir s :inherits #f)))))))))))

;;' @description Create a set from a list.
;;' @param lst List of values to include in the set
;;' @examples
;;' (list->set '(1 2 3 2 1))  ; => set of {1, 2, 3}
;;' @seealso set->list, set
(define list->set
  (lambda (lst)
    (define s (set-new))
    (define items (if (is.list lst) lst (as.list lst)))
    (define add-item
      (lambda (item)
        (begin
          (define key (set-key item))
          (if (not (r-call "exists" (list key :envir s :inherits #f)))
            (r-call "assign" (list key item :envir s))))))
    (r-call "lapply" (list items add-item))
    s))

;;' @description Return the number of elements in a set.
;;' @param s Set to measure
;;' @examples
;;' (set-size (set 1 2 3))  ; => 3
;;' (set-size (set))         ; => 0
;;' @assert
;;' (assert-equal 3 (set-size (set 1 2 3)))
;;' (assert-equal 0 (set-size (set)))
;;' @seealso set?
(define set-size
  (lambda (s)
    (if (not (set? s))
      0
      (length (r-call "ls" (list :envir s :all.names #t :sorted #f))))))

;;' @description Apply function to each element, return new set of results.
;;' @param fn Function to apply to each element
;;' @param s Set to map over
;;' @examples
;;' (set-map inc (set 1 2 3))  ; => set of {2, 3, 4}
;;' @seealso set-filter
(define set-map
  (lambda (fn s)
    (define result (set-new))
    (if (set? s)
      (begin
        (define keys (r-call "ls" (list :envir s :all.names #t :sorted #f)))
        (define process-key
          (lambda (key)
            (begin
              (define value (r-call "get" (list key :envir s :inherits #f)))
              (define new-value (fn value))
              (define new-key (set-key new-value))
              (if (not (r-call "exists" (list new-key :envir result :inherits #f)))
                (r-call "assign" (list new-key new-value :envir result))))))
        (r-call "lapply" (list (as.list keys) process-key))))
    result))

;;' @description Keep only elements matching predicate, return new set.
;;' @param pred Predicate function to test each element
;;' @param s Set to filter
;;' @examples
;;' (set-filter even? (set 1 2 3 4))  ; => set of {2, 4}
;;' @seealso set-map
(define set-filter
  (lambda (pred s)
    (define result (set-new))
    (if (set? s)
      (begin
        (define keys (r-call "ls" (list :envir s :all.names #t :sorted #f)))
        (define process-key
          (lambda (key)
            (begin
              (define value (r-call "get" (list key :envir s :inherits #f)))
              (if (isTRUE (pred value))
                (r-call "assign" (list key value :envir result))))))
        (r-call "lapply" (list (as.list keys) process-key))))
    result))
)
