(module scrambler
        (extern
         (include "stdlib.h")
         (c-rand::int () "rand")
         (c-srand::void (::uint) "srand"))
        (main main))

;; --------------------------------
;; Configuration
(define *reserved* '(let flet true false not and or xor iff implies if_then_else exists forall theory ite = < > <= >= + - * / ~ distinct U Int Real Array Array1 Array2 select store concat extract bvor bvand bvnot bvxor bvneg bvsub bvadd bvmul __mk_bv bvslt bvsleq bvsgt bvsgeq bvlt bvleq bvgt bvgeq sign_extend rotate_left rotate_right repeat shift_right0 shift_right1 shift_left1 shift_left0 bvnand bvnor bit0 bit1 fill))
;; add bitvector sorts
(let loop ((i 1))
	(when (<= i 32)
		(let* ((str (string-append "BitVec@O" (number->string i) "@C"))
					 (sym (string->symbol str)))
			(set! *reserved* (cons sym *reserved*))
			(loop (+ i 1)))))
(define *ac* '(and or iff xor + * bvadd bvmul bvand bvor bvxor bvnand bvnor))
(define *tmp-file* "/tmp/smt-tmp-file.txt")
;; --------------------------------

;; --------------------------------
;; Support
(define (init!)
  (for-each (lambda (sym)
              (putprop! sym 'reserved #t))
            *reserved*)
  (for-each (lambda (sym)
              (putprop! sym 'ac #t))
            *ac*))

(define (reserved? sym) 
  (getprop sym 'reserved))
  
(define (ac? sym)
  (getprop sym 'ac))

(init!)

(define (print-error . args)
  (with-output-to-port (current-error-port)
    (lambda ()
      (apply print args))))

(define (sign-error . args)
  (apply print-error args)
  (exit -1))

(define (rand n)
  (modulo (c-rand) n))

(define (srand seed)
  (c-srand seed)
  #unspecified)

(define (permute-vector! vect)
  (let ((n (vector-length vect)))
    (let loop ((i 0))
      (when (< i n)
        (let ((j (+ (rand (- n i)) i)))
          (unless (= i j)
            (let ((tmp (vector-ref vect i)))
              (vector-set! vect i (vector-ref vect j))
              (vector-set! vect j tmp))))
        (loop (+ i 1))))))

(define (permute-list l)
  (let ((v (list->vector l)))
    (permute-vector! v)
    (vector->list v)))

(define (identity x) x)
;; --------------------------------

;; --------------------------------
;; Symbol renaming
(define *map* (make-hashtable))
(define *already-created* (make-hashtable))
(define (gen-new-sym prefix)
  (let* ((new-str (string-append prefix (integer->string (rand 10000))))
         (new-sym (string->symbol new-str)))
    (cond
     ((hashtable-get *already-created* new-sym)
      (gen-new-sym new-str))
     (else
      (hashtable-put! *already-created* new-sym #t)
      new-sym))))
(define (map-sym sym)
  (cond
   ((reserved? sym) sym)
   ((hashtable-get *map* sym) => identity)
   (else
    (let* ((str (symbol->string sym))
           (first-char (string-ref str 0))
           (prefix (cond
                    ((eq? first-char #\?)
                     "?t")
                    ((eq? first-char #\$)
                     "$p")
                    (else
                     "x")))
           (new-sym (gen-new-sym prefix)))
      (hashtable-put! *map* sym new-sym)
      new-sym))))
;; --------------------------------


;; --------------------------------
;; Load a SMT-LIB file
(define *logic* #f)
(define *extrasorts* '())
(define *extrafuns* '())
(define *extrapreds* '())
(define *assumptions* '())
(define *formula* #f)
(define (load-smt! file-name)
  (let ((sed-cmd (string-append "sed 's|(benchmark .*|(benchmark smtcomp|;s|bv\\([0-9][0-9]*\\)|(__mk_bv \\\"\\1\\\")|g;s|extract\\[\\([0-9][0-9]*\\):\\([0-9][0-9]*\\)\\]|(extract \\1 \\2)|g;s|fill\\[\\([0-9][0-9]*\\)\\]|(fill \\1)|g;s|{|(\\\"|g;s|}|\\\")|g;s| \\([0-9][0-9]*\\.[0-9][0-9]*\\)| \\\"\\1\\\"|g;s| \\([0-9][0-9]*\\)| \\\"\\1\\\"|g;s|\\[|@O|g;s|\\]|@C|g' \"" file-name "\" > " *tmp-file*)))
    ;; (print sed-cmd)
    (unless (= (system sed-cmd) 0)
      (sign-error "Failed to execute: " sed-cmd))
    (let ((sexpr (with-input-from-file *tmp-file*
                   (lambda ()
                     (read)))))
			;; (pp sexpr)
			;; (exit 0)
      (process-sexpr! sexpr))))
(define (process-sexpr! sexpr)
  (match-case sexpr
    ((benchmark ?name . ?args)
     (let loop ((args args))
       (unless (or (null? args) (null? (cdr args)))
         (let ((id  (car args))
               (arg (cadr args)))
           (cond
            ((eq? id ':extrasorts)
             (set! *extrasorts* (append! *extrasorts* arg)))
            ((eq? id ':extrapreds)
             (set! *extrapreds* (append! *extrapreds* arg)))
            ((eq? id ':extrafuns)
             (set! *extrafuns* (append! *extrafuns* arg)))
            ((eq? id ':assumption)
             (set! *assumptions* (cons arg *assumptions*)))
            ((eq? id ':logic)
             (set! *logic* arg))
            ((eq? id ':formula)
             (when *formula* (sign-error "File has more than one formula."))
             (set! *formula* arg))
            ((memq id '(:status :source :notes :category :difficulty))
             ;; ignore
             #unspecified)
            (else
             (sign-error "Invalid SMT-LIB benchmark: " id)))
           (loop (cddr args))))))
    (?-
     (sign-error "Invalid SMT-LIB benchmark."))))
;; --------------------------------

;; --------------------------------
;; Scramble
(define (scramble sexpr)
  (match-case sexpr
    ((?op . ?args)
     (let ((new-op (scramble op))
           (new-args (map scramble args)))
       (if (and (symbol? new-op) (ac? new-op))
         `(,new-op ,@(permute-list new-args))
         `(,new-op ,@new-args))))
    ((? symbol?)
     (map-sym sexpr))
    (?-
     sexpr)))
;; --------------------------------
  
;; --------------------------------
;; Print benchmark
(define (print-benchmark)
	(with-output-to-file *tmp-file*
		(lambda ()
			(print "(benchmark smtcomp")
			(print ":source { }")
			(print ":status unknown")
			(print ":logic " *logic*)
			(unless (null? *extrasorts*)
				(print ":extrasorts " (map scramble (permute-list *extrasorts*))))
			(unless (null? *extrapreds*)
				(print ":extrapreds " (map scramble (permute-list *extrapreds*))))
			(unless (null? *extrafuns*)
				(print ":extrafuns " (map scramble (permute-list *extrafuns*))))
			(for-each (lambda (a)
									(print ":assumption " a))
								(permute-list (map scramble *assumptions*)))
			(display* ":formula " (scramble *formula*))
			(print ")")))
	(let ((sed-cmd (string-append "sed 's|(__mk_bv \\([0-9][0-9]*\\))|bv\\1|g;s|@O|\\[|g;s|@C|\\]|g;s|(extract \\([0-9][0-9]*\\) \\([0-9][0-9]*\\))|extract\\[\\1:\\2\\]|g;s|(fill \\([0-9][0-9]*\\))|fill\\[\\1\\]|g' \"" *tmp-file* "\"")))
		(unless (= (system sed-cmd) 0)
			(sign-error "Failed to execute: " sed-cmd))))


(define (main args)
  (unless (= (length args) 3)
    (sign-error "Usage: scrambler <seed> <file-name.smt>"))
  (let ((seed (string->integer (cadr args)))
        (file-name (caddr args)))
    (srand seed)
    (load-smt! file-name)
    (print-benchmark)))