;; AIscm - Guile extension for numerical arrays and tensors.
;; Copyright (C) 2013, 2014, 2015, 2016, 2017, 2018, 2019 Jan Wedekind <jan@wedesoft.de>
;;
;; This program is free software: you can redistribute it and/or modify
;; it under the terms of the GNU General Public License as published by
;; the Free Software Foundation, either version 3 of the License, or
;; (at your option) any later version.
;;
;; This program is distributed in the hope that it will be useful,
;; but WITHOUT ANY WARRANTY; without even the implied warranty of
;; MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
;; GNU General Public License for more details.
;;
;; You should have received a copy of the GNU General Public License
;; along with this program.  If not, see <http://www.gnu.org/licenses/>.
;;
(define-module (aiscm tensorflow)
  #:use-module (oop goops)
  #:use-module (ice-9 optargs)
  #:use-module (aiscm core)
  #:use-module (aiscm util)
  #:export (to-tensor from-tensor make-session close-session run make-tensor tf-from-tensor
            with-session graph make-graph make-description tf-finish-operation tf-add-input
            tf-run tf-add-gradient tf-output? tf-graph-operation-by-name tf-reset-graph
            tf-graph-export tf-graph-import tf-operation-names
            ; Code generated using aiscm/tensorflow.scm.in
<?cs each:op = Op ?>            <?cs var:op.name ?>
<?cs /each ?>            ))

(load-extension "libguile-aiscm-tensorflow" "init_tensorflow")

(define graph (make-graph))

(define (tf-graph-export file-name)
  (tf-graph-export_ graph file-name))

(define (tf-graph-import file-name)
  (tf-graph-import_ graph file-name))

(define (tf-reset-graph)
  (set! graph (make-graph)))

(define typemap
  (list (cons <bool>            TF_BOOL      )
        (cons <ubyte>           TF_UINT8     )
        (cons <byte>            TF_INT8      )
        (cons <usint>           TF_UINT16    )
        (cons <sint>            TF_INT16     )
        (cons <uint>            TF_UINT32    )
        (cons <int>             TF_INT32     )
        (cons <ulong>           TF_UINT64    )
        (cons <long>            TF_INT64     )
        (cons <float>           TF_FLOAT     )
        (cons <double>          TF_DOUBLE    )
        (cons <complex<float>>  TF_COMPLEX64 )
        (cons <complex<double>> TF_COMPLEX128)
        (cons <obj>             TF_STRING    )))

(define inverse-typemap (alist-invert typemap))

(define (to-zero-d-array value)
  (let [(arr (make (multiarray (upcast-integer (native-type value)) 0) #:shape '()))]
    (set arr value)
    arr))

(define (to-tensor value)
  (if (is-a? value <multiarray<>>)
    (make-tensor (assq-ref typemap (typecode value)) (shape value) (size-of value) (memory (ensure-default-strides value)))
    (to-tensor (to-zero-d-array value))))

(define (from-tensor tensor)
  (let [(info (tf-from-tensor tensor))]
    (get (make (multiarray (assq-ref inverse-typemap (car info)) (length (cadr info))) #:shape (cadr info) #:memory (caddr info)))))

(define (make-session)
  (make-tf-session graph))

(define (close-session s)
  (close-tf-session s)
  #t)

(define-syntax with-session
  (syntax-rules ()
    ((_ (placeholder) body ...)
     (let ((placeholder #f))
       (dynamic-wind
	 (lambda () (set! placeholder (make-session)))
	 (lambda () body ...)
	 (lambda ()
	   ;; session was managed by GC, so just close it.
	   (close-session placeholder)))))))

(define (run session alist expression)
  (let* [(tensors (map (lambda (entry) (cons (car entry) (to-tensor (cdr entry)))) alist))
         (result  (tf-run session tensors expression))]
    (if (list? result) (map from-tensor result) (from-tensor result))))

(define (tf-add-gradient expression variable)
  (tf-add-gradient_ graph expression variable))

(define (to-tf-output value)
  (if (tf-output? value)
    value
    (if (is-a? value <multiarray<>>)
      (tf-const #:dtype (typecode value) #:value value)
      (to-tf-output (to-zero-d-array value)))))

(define (tf-graph-operation-by-name name)
  (tf-graph-operation-by-name_ graph name))

(define (tf-operation-names)
  (tf-operation-names_ graph))

; Code generated using aiscm/tensorflow.scm.in
<?cs each:op = Op ?>(define (<?cs var:op.name ?><?cs each:in = op.input_arg ?> <?cs name:in ?><?cs /each ?> . keyword-args)
  (let-keywords keyword-args #f (name <?cs each:attr = op.attr ?> <?cs if:attr == "bool" ?>(<?cs name:attr ?> '())<?cs else ?><?cs name:attr ?><?cs /if ?><?cs /each ?>)
    (let [(description (make-description graph "<?cs name:op ?>" (or name (symbol->string (gensym "x")))))]
<?cs each:in = op.input_arg ?><?cs if:in == "list" ?>      (tf-add-input-list description (map to-tf-output <?cs name:in ?>))
<?cs /if ?><?cs if:in == "single" ?>      (tf-add-input description (to-tf-output <?cs name:in ?>))
<?cs /if ?><?cs /each ?><?cs each:attr = op.attr ?><?cs if:attr == "string" ?>      (if <?cs name:attr ?>
        (tf-set-attr-string description "<?cs name:attr ?>" (symbol->string <?cs name:attr ?>)))
<?cs /if ?><?cs if:attr == "int" ?>      (if <?cs name:attr ?>
        (tf-set-attr-int description "<?cs name:attr ?>" <?cs name:attr ?>))
<?cs /if ?><?cs if:attr == "list(int)" ?>      (if <?cs name:attr ?>
        (tf-set-attr-int-list description "<?cs name:attr ?>" <?cs name:attr ?>))
<?cs /if ?><?cs if:attr == "bool" ?>      (if (not (null? <?cs name:attr ?>))
        (tf-set-attr-bool description "<?cs name:attr ?>" <?cs name:attr ?>))
<?cs /if ?><?cs if:attr == "float" ?>      (if <?cs name:attr ?>
        (tf-set-attr-float description "<?cs name:attr ?>" <?cs name:attr ?>))
<?cs /if ?><?cs if:attr == "list(float)" ?>      (if <?cs name:attr ?>
        (tf-set-attr-float-list description "<?cs name:attr ?>" <?cs name:attr ?>))
<?cs /if ?><?cs if:attr == "type" ?>      (if <?cs name:attr ?>
        (tf-set-attr-type description "<?cs name:attr ?>" (assq-ref typemap <?cs name:attr ?>)))
<?cs /if ?><?cs if:attr == "tensor" ?>      (if <?cs name:attr ?>
        (tf-set-attr-tensor description "<?cs name:attr ?>" (to-tensor <?cs name:attr ?>)))
<?cs /if ?><?cs if:attr == "shape" ?>      (if <?cs name:attr ?>
        (tf-set-attr-shape description "<?cs name:attr ?>" <?cs name:attr ?>))
<?cs /if ?><?cs /each ?>      (tf-finish-operation description <?cs var:op.n_output ?>))))

<?cs /each ?>
