(module cardinality-constraints mzscheme
  
  (require (lib "etc.ss")
           (prefix srfi1: (lib "1.ss" "srfi")))
  
  (provide encode-atmost1/lprop
           encode-atmost1/needed-vars
           encode-atmost1)
  
  (define (do-pairs lst)
     (let loop1 ([outlst lst])
       (if (null? (cdr outlst))
           '()
           (let loop2 ([inlst (cdr outlst)])
             (if (null? inlst)
                 (loop1 (cdr outlst))
                 (cons (cons (car inlst) (car outlst))
                       (loop2 (cdr inlst))))))))
  
  (define (encode-atmost1/needed-vars . ops)
    (let ([len (length ops)])
      (if (<= len 5) 0 (- len 1))))
          
  ;; Returns a cnf formula with the encoding for the at most 1 of the operators and uses
  ;; the list of extra-vars as extra-variables. the number of extra variables needs to be the size of ops minus one.
  (define (encode-atmost1 extra-vars . ops)
    (unless (srfi1:every integer? ops)
      (error 'encode-atmost1 "Expected numbers as arguments, got (displayed as list): ~a" ops))
    (let ([n-ops (length ops)])
      (if (and (> n-ops 5) (not (= (length extra-vars) (- n-ops 1))))
          (error 'encode-atmost1 "Expected ~a extra vars for ~a ops, got ~a." (- n-ops 1) n-ops (length extra-vars)))
      (if (<= n-ops 5)
          (map (λ (p) (list (- (car p)) (- (cdr p)))) (do-pairs ops))
          (let ([ops-vec (list->vector ops)]
                [aux-vec (list->vector extra-vars)])
            (srfi1:cons* (list (- (vector-ref ops-vec 0)) (vector-ref aux-vec 0))
                         (list (- (vector-ref ops-vec (- n-ops 1))) (- (vector-ref aux-vec (- n-ops 2))))
                         (let loop ([i 1] [clauses null])
                           (if (= i (- n-ops 1))
                               clauses
                               (loop
                                (+ i 1)
                                (srfi1:cons*
                                 (list (- (vector-ref ops-vec i)) (vector-ref aux-vec i))
                                 (list (- (vector-ref aux-vec (- i 1))) (vector-ref aux-vec i))
                                 (list (- (vector-ref ops-vec i)) (- (vector-ref aux-vec (- i 1))))
                                 clauses)))))))))
  
  ;; Returns a propositional logic formula with the encoding for at most 1 of the operators
  (define (encode-atmost1/lprop . ops)
    (let ([n-ops (length ops)])
      (if (<= n-ops 5) ;; In this case do pairwise encoding
          `(and ,@(map (λ (p) `(or (not ,(car p)) (not ,(cdr p)))) (do-pairs ops)))
          (let ([ops-vec (list->vector ops)]
                [aux-vec (list->vector (build-list (- n-ops 1) (λ (i) (gensym 'aux:))))])
            `(and (=> ,(vector-ref ops-vec 0) ,(vector-ref aux-vec 0))
                  (=> ,(vector-ref ops-vec (- n-ops 1)) (not ,(vector-ref aux-vec (- n-ops 2))))
                  ,@(let loop ([i 1] [clauses null])
                      (if (= i (- n-ops 1))
                          clauses
                          (loop
                           (+ i 1)
                           (srfi1:cons*
                            `(=> ,(vector-ref ops-vec i) ,(vector-ref aux-vec i))
                            `(=> ,(vector-ref aux-vec (- i 1)) ,(vector-ref aux-vec i))
                            `(=> ,(vector-ref ops-vec i) (not ,(vector-ref aux-vec (- i 1))))
                            clauses)))))))))
  
  )