main.ss
;; gnuplot.plt: mzscheme interface for gnuplot
;;
;; (C) Copyright 2008 Dimitris Vyzovitis <vyzo at media.mit.edu>
;;
;; gnuplot.plt 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.
;;
;; mzsocket 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 gnuplot.plt.  If not, see <http://www.gnu.org/licenses/>.
;;
;; requires gnuplot-4.x
;; For more info on gnuplot see: http://www.gnuplot.info
#lang scheme/base

(require scheme/match
         scheme/string
         scheme/file
         scheme/system)

(provide
  gnuplot?
  (rename-out (data? gnuplot-data?) (gnuplot-in gnuplot-input-port))
  gnuplot-data
  gnuplot-data/file
  gnuplot-item
  gnuplot-plot
  gnuplot-splot
  gnuplot-replot
  gnuplot-multiplot
  gnuplot-hardcopy
  gnuplot-set
  gnuplot-set*
  gnuplot-unset
  gnuplot-reset
  gnuplot-spawn
  gnuplot-kill
  gnuplot-program
  gnuplot-set-program!
  gnuplot-writeln)

(define-struct data (source temp? bin?))
(define-struct gnuplot (pid out in ctl))

;; data: [[number? ...] ...] ; generated data
;; #:file path-string? ; for persistent data
;; #:comments [string? ...] ; optional comments for generated file
;; #:binary bool? ; use binary data in matrix form
(define (gnuplot-data data 
                      #:file (fname #f) 
                      #:comments (comments #f)
                      #:tmpdir (tmpdir #f)
                      #:binary (bin? #f))
  (let ((temp? (not fname))
        (fname (or fname (make-temporary-file "plot.~a" #f tmpdir))))
    (call-with-output-file fname
      (lambda (port)
        (when (and (not bin?) comments)
          (for ((x comments)) (fprintf port "# ~a~n" x)))
        (let-values 
            (((writev nextln)
              (if bin?
                (values
                  (lambda (x) (write-bytes ; XXX endianness
                          (real->floating-point-bytes x 4)
                          port))
                  void)
                (values
                  (lambda (x) (fprintf port "~a " x))
                  (lambda () (fprintf port "~n"))))))
          (for ((row data))
            (for ((x row)) (writev x))
            (nextln)))
        (flush-output port)
        (make-data* fname temp? bin?))
      #:exists 'truncate)))

(define (make-data* source temp? bin?)
  (let ((data (make-data source temp? bin?)))
    (when temp? (will-register *gc* data gc-data))
    data))

;; fname: path-string?
(define (gnuplot-data/file fname (bin? #f))
  (make-data fname #f bin?))

;; plot items: ties data to plot options
;; data: data? or string?
;; opts: list?
(define (gnuplot-item data (opts null))
  `(item ,data . ,opts))

;; plotting
(define (gnuplot-plot gplot
                      #:range (range #f)
                      . items)
  (gnuplot-plot* gplot 'plot range items))

(define (gnuplot-splot gplot
                       #:range (range #f)
                       . items)
  (gnuplot-plot* gplot 'splot range items))

(define (gnuplot-replot gplot  . items)
  (gnuplot-plot* gplot 'replot #f items))

;; range: [<range> ...]/#f
(define (gnuplot-plot* gplot cmd range items)
  (gnuplot-writeln gplot
    "~a ~a ~a" 
    cmd 
    (if range (/string `(seq* . ,(map make-range range))) "")
    (/string `(seq . ,items))))

(define (gnuplot-multiplot gplot 
                           #:layout (layout null)
                           . plots)
  (gnuplot-set gplot `(multiplot . ,layout))
  (for ((doplot plots)) (doplot gplot))
  (gnuplot-unset gplot 'multiplot))

;; 'hardcopy' plot
(define (gnuplot-hardcopy gplot outp (plot gnuplot-replot)
                          #:term (term '(postscript color)))
  ;; save+restore terminal
  (gnuplot-set gplot '(terminal push))
  (gnuplot-set gplot `(terminal . ,term))
  (gnuplot-set gplot `(output (path ,outp)))
  (plot gplot)
  (gnuplot-unset gplot 'output)
  (gnuplot-set gplot '(terminal pop)))

(define (gnuplot-set gplot opt)
  (gnuplot-writeln gplot "set ~a" (make-option opt)))

(define (gnuplot-set* gplot opts)
  (for ((x opts)) (gnuplot-set gplot x)))

(define (gnuplot-unset gplot opt)
  (gnuplot-writeln gplot "unset ~a" opt))

(define (gnuplot-reset gplot)
  (gnuplot-writeln gplot "reset"))

;; option formatting
;;   deal with gnuplot's idiosyncratic syntax...
(define (/string x)
  (match x
    (`(seq . ,vs) (string-join (map /string vs) ","))
    (`(seq: . ,vs) (string-join (map /string vs) ":"))
    (`(seq* . ,vs) (string-join (map /string vs) " "))
    (`(range . ,rg) (make-range rg))
    (`(= ,x ,v) (format "~a = ~a" x (/string v)))
    (`(path ,v) (format "'~a'" (if (path? v) (path->string v) v)))
    (`(str ,v) (format "\"~a\"" v))
    (`(item ,data . ,opts)
     (string-join 
      (cons
       (/string (if (data? data) `(path ,(data-source data)) `(str ,data)))
       (map /string 
         (if (and (data? data) (data-bin? data) (not (memq 'binary opts)))
           (cons 'binary opts)
           opts)))
      " "))
    (#f "")
    (else (format "~a" x))))

(define (make-option opt)
  (/string `(seq* . ,opt)))

;; range formatting
;;  <range> :: [] ; turn off auto scale
;;             [min  max] ; #f for omitting
;;             [dummy min max] ; for parametric plots
(define (make-range range)
  (match range
    ('() "[]")
    (`(,min ,max) (format "[~a:~a]" (or min "") (or max "")))
    (`(,dummy ,min ,max) 
     (format "[~a = ~a:~a]" dummy (or min "") (or max "")))
    ((? string?) range)))

;; gnuplot process
(define *gnuplot*
  (case (system-type)
    ((unix macosx) "gnuplot")
    ((windows) "pgnuplot.exe")))

(define (gnuplot-set-program! prog)
  (set! *gnuplot* prog))
(define (gnuplot-program) *gnuplot*)

;; read? ; keep gnuplot's error port open for reading
;;  Note: gnuplot is very idiosyncratic when it comes to interface
;;   uses stdout for printing its prompt and other garbage
;;   uses stderr for printing the useful things (with a completely inconsistent
;;    output set...)
(define (gnuplot-spawn (read? #f))
  ;; gc register
  (match (process *gnuplot*)
    ((list out in pid err ctl)
     (close-input-port out)
     (unless read? 
       (close-input-port err))
     (let ((gplot (make-gnuplot pid in (if read? err #f) ctl)))
       (will-register *gc* gplot gnuplot-kill)
       gplot))))

(define (gnuplot-kill gplot)
  ((gnuplot-ctl gplot) 'kill)
  (when (gnuplot-in gplot)
    (close-input-port (gnuplot-in gplot)))
  (close-output-port (gnuplot-out gplot)))

(define (gnuplot-writeln gplot fmt . args)
  (let ((port (gnuplot-out gplot)))
    (apply fprintf port fmt args)
    (newline port)
    (flush-output port)))

;; gcing of temp files + processes
(define *gc* (make-will-executor))
(define *gc-thread*
  (thread
   (lambda ()
     (let lp () (will-execute *gc*) (lp)))))
(define (gc-data data)
  (delete-file (data-source data)))