(module mathsat-translator
        (main main))

(define *mathsat-keyword-list*
  '("not" "or"))

(define *line* 1)

(define *mode* 'esolver)

(define (mathsat/init-lexer!)
  (for-each (lambda (word)
              (putprop! (string->symbol word) 'mathsat-reserved #t))
            *mathsat-keyword-list*))

(define *mathsat-lexer*
  (regular-grammar
   ((nonzero-digit (in ("19"))))
   ((in #\newline #a012)
    (set! *line* (+ *line* 1))
    (ignore))
   ((+ (in #\space #\tab))
    (ignore))
   ;; parenthesis
   (#\(
    (cons 'LP *line*))
   (#\)
    (cons 'RP *line*))
   ;; brackets
   (#\[
    (cons 'LB *line*))
   (#\]
    (cons 'RB *line*))
   ;; integer constant
   ((+ digit)
    (cons 'NUM (string->integer (the-string))))
   ("+"
    (cons 'PLUS *line*))
   ("-"
    (cons 'MINUS *line*))
   ("<"
    (cons 'LT *line*))
   ("<="
    (cons 'LE *line*))
   (">"
    (cons 'GT *line*))
   (">="
    (cons 'GE *line*))
   ("="
    (cons 'EQ *line*))
   ((: (or #\_ alpha) (* (or #\_ alpha digit)))
    (let* ((string   (the-string))
           (symbol   (string->symbol string))
           (upsymbol (string->symbol (string-upcase string))))
      (cond
       ((getprop symbol 'mathsat-reserved)
        (cons upsymbol *line*))
       (else
        (cons 'ID symbol)))))))

(define *mathsat-parser*
  (lalr-grammar
   ;;tokens
   (OR LP RP LB RB NUM PLUS MINUS LT LE GT GE EQ NOT ID)
   (start
    ((hints constraints)
     constraints))
   (hints
    ((LP step-hints id-list id-list-list RP)
     ;; ignore hints
     #unspecified))
   (id-list
    ((LP id+ RP)
     (reverse! id+)))
   (id+
    ((ID)
     (list ID))
    ((id+ ID)
     (cons ID id+)))
   (id-list-list
    ((LP id-list+ RP)
     (reverse! id-list+)))
   (id-list+
    ((id-list)
     (list id-list))
    ((id-list+ id-list)
     (cons id-list id-list+)))
   (step-hints
    ((LP LB NUM RB id-list+ RP)
     #unspecified))
   (constraints
    ((LP constraint+ RP)
     `(AND ,@(reverse! constraint+))))
   (constraint+
    ((constraint)
     (list constraint))
    ((constraint+ constraint)
     (cons constraint constraint+)))
   (constraint
    ((LP literal+ RP)
     `(OR ,@(reverse! literal+))))
   (literal+
    ((literal)
     (list literal))
    ((literal+ OR literal)
     (cons literal literal+)))
   (literal
    ((atom)
     atom)
    ((NOT atom)
     `(NOT ,atom)))
   (atom
    ((ID) ID)
    ((LP EQ ID@id1 ID@id2 RP) `(EQ ,id1 ,id2))
    ((LP EQ MINUS ID@id1 ID@id2 MINUS ID@id3 ID@id4 RP) 
     (when (eq? *mode* 'uclid)
       (sign-error "UCLID only supports separation logic."))
     `(EQ (SUB ,id1 ,id2) (SUB ,id3 ,id4)))
    ((LP GT ID@id1 ID@id2 RP) `(GT ,id1 ,id2))
    ((LP LT ID@id1 ID@id2 RP) `(LT ,id1 ,id2))
    ((LP GE ID@id1 ID@id2 RP) `(GE ,id1 ,id2))
    ((LP LE ID@id1 ID@id2 RP) `(LE ,id1 ,id2))
    ((LP GE MINUS ID@id1 ID@id2 NUM RP) `(GE (SUB ,id1 ,id2) ,NUM))
    ((LP LE MINUS ID@id1 ID@id2 NUM RP) `(LE (SUB ,id1 ,id2) ,NUM))
    ((LP GT MINUS ID@id1 ID@id2 NUM RP) `(GT (SUB ,id1 ,id2) ,NUM))
    ((LP LT MINUS ID@id1 ID@id2 NUM RP) `(LT (SUB ,id1 ,id2) ,NUM)))
   ))

(define (mathsat/parse)
  (read/lalrp
   *mathsat-parser*
   *mathsat-lexer*
   (current-input-port)))

(define (mathsat/parse-string astring)
  (with-input-from-string
      astring
    mathsat/parse))

(define (mathsat/parse-file filename)
  (with-input-from-file
      filename
    mathsat/parse))

(define (collect-var-types ast)
  (let ((table (make-hashtable)))
    (let loop ((ast ast)
               (expected-type 'bool))
      (cond
       ((list? ast)
        (let ((c (car ast)))
          (cond
           ((memq c '(ADD SUB GT GE LT LE EQ NEQ))
            (for-each (lambda (child) (loop child 'real))
                      (cdr ast)))
           (else
            (for-each (lambda (child) (loop child 'bool))
                      (cdr ast))))))
       ((symbol? ast)
        (unless (memq ast '(true false))
          (hashtable-put! table ast expected-type)))))
    table))

(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 *ics-op-table* '((ADD . +) (SUB . -) (LT . <) (LE . <=) (GT . >) (GE . >=) (EQ . =) (NEQ . <>)
                         (AND . "&") (OR . "|") (IFF . <=>) (IMPLIES . =>)))

(define (display-infix-core op arg-list o-par c-par loop)
  (display o-par)
  (loop (car arg-list))
  (for-each (lambda (arg)
              (display* " " op " ")
              (loop arg))
            (cdr arg-list))
  (display c-par))

(define (display-infix op arg-list loop)
  (display-infix-core op arg-list "(" ")" loop))

(define (display-infix-prop op arg-list loop)
  (display-infix-core op arg-list "[" "]" loop))

(define (mathsat->ics ast)
  (print "sat ")
  (let* ((var-type-table (collect-var-types ast))
         (pp-ast (lambda (ast)
                   (let loop ((ast ast))
                     (cond 
                      ((list? ast)
                       (let ((c (car ast)))
                         (cond
                          ((memq c '(ADD SUB))
                           (display-infix (cdr (assq c *ics-op-table*)) (cdr ast) loop))
                          ((memq c '(OR AND IMPLIES IFF LT LE GT GE EQ NEQ))
                           (display-infix-prop (cdr (assq c *ics-op-table*)) (cdr ast) loop))
                          ((eq? c 'NOT)
                           [assert (ast) (= (length ast) 2)]
                           (display "[~ ")
                           (loop (cadr ast))
                           (display "]"))
                          (else
                           (error 'mathsat->ics "unexpected ast" ast)))))
                      (else
                       (display ast)))))))
    (pp-ast ast)
    (print " .")
    ))

(define *esolver-op-table* '((ADD . +) (SUB . -) (LT . <) (LE . <=) (GT . >) (GE . >=) (EQ . =) (NEQ . /=)
                             (AND . "and") (OR . "or") (IFF . =) (IMPLIES . =>) (NOT . not) (ITE . ite)))

(define (mathsat->esolver ast)
  (let ((var-type-table (collect-var-types ast)))
    (hashtable-for-each var-type-table
                        (lambda (var type)
                          (print "(define " var "::" (cond
                                                      ((eq? type 'bool) 
                                                       "bool")
                                                      (else "real"))
                                 ")")))
    (print "(theorem th1 (not ")
    (let loop ((ast ast))
      (cond
       ((list? ast)
        [assert (ast *esolver-op-table*) (assq (car ast) *esolver-op-table*)]
        (display* "(" (cdr (assq (car ast) *esolver-op-table*)))
        (for-each (lambda (arg)
                    (display " ")
                    (loop arg))
                  (cdr ast))
        (display ")"))
       (else
        (display ast))))
    (print "))")))


(define *simplify-op-table* '((ADD . +) (SUB . -) (LT . <) (LE . <=) (GT . >) (GE . >=) (EQ . EQ)))

(define (mathsat->simplify ast)
  (let ((var-type-table (collect-var-types ast))
        (first-clause? #t))
    (print "(NOT (AND ")
    (for-each (lambda (clause)
                (display "(OR ")
                (for-each (lambda (literal)
                            (let* ((neg? (and (list? literal) (eq? (car literal) 'NOT)))
                                   (atom (if neg? (cadr literal) literal)))
                              (when neg?
                                (display "(NOT "))
                              (match-case atom
                                ((EQ (SUB ?id1 ?id2) (SUB ?id3 ?id4))
                                 (display* "(EQ (- " id1 " " id2 ") (- " id3 " " id4 "))"))
                                ((EQ ?id1 ?id2)
                                 (display* "(EQ " id1 " " id2 ")"))
                                (((or GT LT GE LE) (SUB ?id1 ?id2) ?num)
                                 (display* "(" (cdr (assq (car atom) *simplify-op-table*))
                                           " (- " id1 " " id2 ") " num " )"))
                                (((or GT LT GE LE) ?id1 ?id2)
                                 (display* "(" (cdr (assq (car atom) *simplify-op-table*)) " "
                                           id1 " " id2 ")"))
                                ((? symbol?)
                                 (display atom))
                                (?-
                                 (sign-error "Error in Simplify translation...")))
                              (when neg?
                                (display ")"))
                              (display " ")))
                          (cdr clause))
                (print ")"))
              (cdr ast))
    (print "))")))


(define *svc-op-table* '((ADD . +) (SUB . -) (LT . <) (LE . <=) (GT . >) (GE . >=) (EQ . EQ)))

(define (mathsat->svc ast)
  (let ((var-type-table (collect-var-types ast))
        (first-clause? #t))
    (print "check_valid (NOT (AND ")
    (for-each (lambda (clause)
                (display "(OR ")
                (for-each (lambda (literal)
                            (let* ((neg? (and (list? literal) (eq? (car literal) 'NOT)))
                                   (atom (if neg? (cadr literal) literal)))
                              (when neg?
                                (display "(NOT "))
                              (match-case atom
                                ((EQ (SUB ?id1 ?id2) (SUB ?id3 ?id4))
                                 (display* "(= (+ " id1 " (* -1 " id2 ") ) (+ " id3 " (* -1 " id4 ") ) )"))
                                ((EQ ?id1 ?id2)
                                 (display* "(= " id1 " " id2 ")"))
                                (((or GT LT GE LE) (SUB ?id1 ?id2) ?num)
                                 (display* "(" (cdr (assq (car atom) *simplify-op-table*))
                                           " (+ " id1 " (* -1 " id2 ") ) " num " )"))
                                (((or GT LT GE LE) ?id1 ?id2)
                                 (display* "(" (cdr (assq (car atom) *simplify-op-table*)) " "
                                           id1 " " id2 ")"))
                                ((? symbol?)
                                 (display atom))
                                (?-
                                 (sign-error "Error in Simplify translation...")))
                              (when neg?
                                (display ")"))
                              (display " ")))
                          (cdr clause))
                (print ")"))
              (cdr ast))
    (print ") )")
    (print "quit")))

(define *uclid-op-table* '((ADD . +) (SUB . -) (LT . <) (LE . <=) (GT . >) (GE . >=) (EQ . =)))

(define (mathsat->uclid ast)
  (let ((var-type-table (collect-var-types ast))
        (first-clause? #t))
    (print "MODEL example")
    (print "CONTROL")
    (print "EXTVAR")
    (print "STOREVAR")
    (print "VAR")
    (print "CONST")
    (hashtable-for-each var-type-table
                        (lambda (var type)
                          (print  var " : " (cond
                                             ((eq? type 'bool) 
                                              "TRUTH")
                                             (else "TERM"))
                                  ";")))
    (print "DEFINE")
    (print "EXEC")
    (print "decide(~(")
    (for-each (lambda (clause)
                (if first-clause?
                  (set! first-clause? #f)
                  (print " &"))
                (display "  (")
                (let ((first-literal? #t))
                  (for-each (lambda (literal)
                              (if first-literal?
                                (set! first-literal? #f)
                                (display " | "))
                              (let ((atom (cond
                                           ((and (list? literal) (eq? (car literal) 'NOT))
                                            (display "~")
                                            (cadr literal))
                                           (else
                                            literal))))
                                (match-case atom
                                  ((EQ ?id1 ?id2)
                                   (display* "(" id1 " = " id2 ")"))
                                  (((or GT LT GE LE) (SUB ?id1 ?id2) ?num)
                                   (display* "(" id1 " " (cdr (assq (car atom) *uclid-op-table*)) " " 
                                             (if (< num 0) "pred" "succ") "^" num "(" id2 "))"))
                                  (((or GT LT GE LE) ?id1 ?id2)
                                   (display* "(" id1 " " (cdr (assq (car atom) *uclid-op-table*)) " " id2 ")"))
                                  ((? symbol?)
                                   (display atom))
                                  (?-
                                   (sign-error "Error in UCLID translation...")))))
                          (cdr clause))
                (display ")")))
              (cdr ast))
    (print "));")))

(define *cvc-op-table* '((ADD . +) (SUB . -) (LT . <) (LE . <=) (GT . >) (GE . >=) (EQ . =) (NEQ . /=)
                         (AND . AND) (OR . OR) (XOR . XOR) (IFF . <=>) (IMPLIES . =>) ))

(define (mathsat->cvc ast)
  (let ((var-type-table (collect-var-types ast)))
    (hashtable-for-each var-type-table
                        (lambda (var type)
                          (print var " : " (cond
                                            ((eq? type 'bool) 
                                             "BOOLEAN")
                                            (else "REAL"))
                                 ";")))
    (print "QUERY NOT(")
    (let loop ((ast ast))
      (cond 
       ((list? ast)
        (let ((c (car ast)))
          (cond
           ((memq c '(ADD SUB OR AND IMPLIES IFF LT LE GT GE EQ NEQ))
            (display-infix (cdr (assq c *cvc-op-table*)) (cdr ast) loop))
           ((eq? c 'NOT)
            [assert (ast) (= (length ast) 2)]
            (display "(NOT ")
            (loop (cadr ast))
            (display ")"))
           (else
            (error 'mathsat->cvc "unexpected ast" ast)))))
       (else
        (display ast))))
      (print " );")
      ))

(define (mathsat-info input-file ast)
  (let ((var-type-table (collect-var-types ast))
        (num-bool-vars 0)
        (num-non-bool-vars 0))
    (hashtable-for-each var-type-table
                        (lambda (var type)
                          (if (eq? type 'bool)
                            (set! num-bool-vars (+ num-bool-vars 1))
                            (set! num-non-bool-vars (+ num-non-bool-vars 1)))))
    (print (+ num-bool-vars num-non-bool-vars) " " (prefix input-file) " " num-bool-vars " " num-non-bool-vars)))
  
(define *version* 1.0)

(define (main argv)
  (mathsat/init-lexer!)
  (let ((output-file #f)
        (input-file #f))
    (args-parse (cdr argv)
      (section "Help")
      (("?")
       (args-parse-usage #f)
       (exit 0))
      ((("-h" "--help") (help "?,-h,--help" "This help message"))
       (args-parse-usage #f)
       (exit 0))
      (section "Misc")
      ((("-v" "--version") (help "Version number"))
       (print *version*)
       (exit 0))
      (("--ics" (help "Produce an ICS formula."))
       (set! *mode* 'ics))
      (("--uclid" (help "Produce an UCLID formula."))
       (set! *mode* 'uclid))
      (("--cvc" (help "Produce a CVC formula."))
       (set! *mode* 'cvc))
      (("--svc" (help "Produce a SVC formula."))
       (set! *mode* 'svc))
      (("--esolver" (help "Produce an esolver formula."))
       (set! *mode* 'esolver))
      (("--simplify" (help "Produce an Simplify formula."))
       (set! *mode* 'simplify))
      (("--info" (help "Produce info."))
       (set! *mode* 'info))
      (("-o" ?file (help "The output file"))
       (set! output-file file))
      (else
       (when input-file
         (print-error "Illegal argument `" else "'. Usage:")
         (args-parse-usage #f)
         (exit -1))
       (set! input-file else)))
    (unless input-file
      (sign-error "Usage: mathsat-translator [options] mathsat-file-name\nTry `mathsat-translator --help' for more information.")) 
    (let ((ast (try
                (mathsat/parse-file input-file)
                (lambda (escape proc msg obj)
                  (sign-error "Error at line " *line*)))))
      (with-output-to-port (if output-file
                             (open-output-file output-file)
                             (current-output-port))
        (lambda ()
          (case *mode*
            ((ics)
             (mathsat->ics ast))
            ((uclid)
             (mathsat->uclid ast))
            ((cvc)
             (mathsat->cvc ast))
            ((svc)
             (mathsat->svc ast))
            ((simplify)
             (mathsat->simplify ast))
            ((info)
             (mathsat-info input-file ast))
            (else
             (mathsat->esolver ast))))))))