;;
;; Neurolucida XML file manipulation.
;;
;; Copyright 2011 Ivan Raikov and the Okinawa Institute of Science and
;; Technology.
;;
;; 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.
;;
;; A full copy of the GPL license can be found at
;; <http://www.gnu.org/licenses/>.
;;

(require-extension files data-structures srfi-1 srfi-13  sxml-transforms sxpath ssax getopt-long )


(define lookup-def 
  (lambda (k lst . rest)
    (let-optionals rest ((default #f))
      (alist-ref k lst eq? default))))


(define (quotewrapped? str)
  (and (string? str) (string-prefix? "\"" str) (string-suffix? "\"" str) ))


(define (quotewrap str)
  (cond ((quotewrapped? str) str)
	((string-any char-whitespace? str)
	 (string-append "\"" str "\""))
	(else str)))


(define (mkdir dir)
  (system* "mkdir -p ~a" (quotewrap dir)))


(define *quiet* #f)


(define (d fstr . args)
  (let ([port (if *quiet* (current-error-port) (current-output-port))])
    (apply fprintf port fstr args)
    (flush-output port) ) )


(define opt-defaults
  `(
    (key             . "color")
    (format          . "swc")
    ))


(define (defopt x)
  (lookup-def x opt-defaults))

(define (symbol-upcase str)
  (string->symbol (string-upcase str)))


(define opt-grammar
  `(
    (data-dir
     "set download directory (default is a randomly generated name in /tmp)"
     (single-char #\d)
     (value (required DIR)))

    (key
     "extraction key"
     (single-char #\k)
     (value       (required "FIELD")
		  (default  ,(defopt 'key))))

    (format
     "output format (swc or asc)"
     (single-char #\f)
     (value       (required "FORMAT")
		  (default  ,(defopt 'format))))

    (help  "Print help"
	    (single-char #\h))
  
  ))


;; Use args:usage to generate a formatted list of options (from OPTS),
;; suitable for embedding into help text.
(define (neurolucida:usage)
  (print "Usage: " (car (argv)) " [options...] operands ")
  (newline)
  (print "Where operands are XML Neurolucida files")
  (newline)
  (print "The following options are recognized: ")
  (newline)
  (width 35)
  (print (parameterize ((indent 5)) (usage opt-grammar)))
  (exit 1))


;; Process arguments and collate options and arguments into OPTIONS
;; alist, and operands (filenames) into OPERANDS.  You can handle
;; options as they are processed, or afterwards.

(define opts    (getopt-long (command-line-arguments) opt-grammar))
(define opt     (make-option-dispatch opts opt-grammar))


(define data-dir (make-parameter #f))

(define (get-data-dir)
  (or (opt 'data-dir)
      (or (data-dir)
	  (let ([dir (create-temporary-directory)])
	    (data-dir dir)
	    dir ) ) ))


(define (create-temporary-directory)
  (let ((dir (or (get-environment-variable "TMPDIR") 
		 (get-environment-variable "TEMP") 
		 (get-environment-variable "TMP") 
		 "/tmp")))
    (let loop ()
      (let* ((n (current-milliseconds))	; current milliseconds
	     (pn (make-pathname dir (string-append "neurolucida-" (number->string n 16)) "tmp")))
	(cond ((file-exists? pn) (loop))
	      (else (mkdir pn) pn))))))


(define neurolucida-xmlns "http://www.mbfbioscience.com/2007/neurolucida")


(define (parse-sxml fpath)
  (with-input-from-file fpath
    (lambda () (cons '*TOP* (ssax:xml->sxml (current-input-port) `((nl . ,neurolucida-xmlns)))))))




(define (extract-element sxml element-name key-attr)
  (let ((element-fullname (string->symbol (conc "nl:" (->string element-name)))))
    (let* ((elements ((sxpath `(// nl:mbf ,element-fullname ))  sxml))
	   (keys     ((sxpath `(// ,element-fullname @ ,key-attr)) `(*TOP* ,@elements))))
      (zip keys elements))))


(define (key=? a b)  (string=? (cadar a) (cadar b)))


(define (partition-by-key a k key=?)
  (partition (lambda (x) (equal? (car x) k)) a))

;;
;; SWC format
;;
;; n T x y z R P
;;
;; n is an integer label that identifies the current point and
;; increments by one from one line to the next.
;;
;; T is an integer representing the type of neuronal segment, such as
;; soma, axon, apical dendrite, etc. The standard accepted integer
;; values are given below.
;;
;;     * 0 = undefined
;;     * 1 = soma
;;     * 2 = axon
;;     * 3 = dendrite
;;     * 4 = apical dendrite
;;     * 5 = fork point
;;     * 6 = end point
;;     * 7 = custom
;;
;; x, y, z gives the cartesian coordinates of each node.
;;
;; R is the radius at that node.
;;
;; P indicates the parent (the integer label) of the current point or
;; -1 to indicate an origin (soma).


(define (print-swc x)
  (define get (compose car alist-ref))

  (let ((sxml `(*TOP* ,@(cdr x))))

    (let* ((offset (make-parameter 0))
           (soma ((sxpath `(nl:contour))  sxml))
	   (soma-points ((sxpath `(nl:point @))  soma))
	   (soma-npoints (length soma-points))
	   (soma-indices (list-tabulate soma-npoints (lambda (x) (+ 1 x)))))

       (for-each (lambda (p n)
		  (let ((data (cdr p)))
		    (let ((T    1) ;; soma
			  (x  (get 'x data))
			  (y  (get 'y data))
			  (z  (get 'z data))
			  (R  (number->string (/ (string->number (get 'd data)) 2)))
			  (P  (- n 1))
			  )
		    (printf "~A ~A ~A ~A ~A ~A ~A~%" n T x y z R (if (zero? P) -1 P)))))
		soma-points soma-indices)
       
       (offset soma-npoints)

       (let recur ((tree ((sxpath `(nl:tree))  sxml))  
		   (parent-index soma-npoints))

	 (let ((points     ((sxpath `(nl:point @))  tree))
	       (subtrees   ((sxpath `(nl:tree))  tree)))
	   
	   (let* ((npoints (length points))
		  (indices (let ((k (+ 1 (offset)))) (list-tabulate npoints (lambda (x) (+ k x)))))
		  (parents (cons parent-index indices)))
	     
	     (for-each (lambda (pt n par)
			 (let ((data (cdr pt)))
			   (let ((T    3) ;; dendrite
				 (x  (get 'x data))
				 (y  (get 'y data))
				 (z  (get 'z data))
				 (R  (number->string (/ (string->number (get 'd data)) 2)))
				 (P  par)
				 )
			     (printf "~A ~A ~A ~A ~A ~A ~A~%" 
                                     n T x y z R P)
			     )))
		       points indices parents)
	     
	     (let ((parent-index1 (+ npoints (offset))))

	       (offset (+ (offset) npoints)) 

	       (if (pair? subtrees)
		   (for-each (lambda (subtree) (recur subtree parent-index1) )
			     subtrees)))
	     )))

    )))


(define (print-pts x)
  (define get (compose car alist-ref))

  (let ((sxml `(*TOP* ,@(cdr x))))

    (let* ((soma ((sxpath `(nl:contour))  sxml))
	   (soma-points ((sxpath `(nl:point @))  soma))
	   (soma-npoints (length soma-points)))

       (for-each (lambda (p)
		  (let ((data (cdr p)))
		    (let ((x  (get 'x data))
			  (y  (get 'y data))
			  (z  (get 'z data))
			  (R  (/ (string->number (get 'd data)) 2))
			  )
		      (printf "~A ~A ~A~%" x y z))))
		soma-points)

       (let recur ((tree ((sxpath `(nl:tree))  sxml))  )

	 (let ((points     ((sxpath `(nl:point @))  tree))
	       (subtrees   ((sxpath `(nl:tree))  tree)))
	   
	   (let ((npoints (length points)))
	     
	     (for-each (lambda (pt n)
			 (let ((data (cdr pt)))
			   (let ((x  (string->number (get 'x data)))
				 (y  (string->number (get 'y data)))
				 (z  (string->number (get 'z data)))
				 (r  (/ (string->number (get 'd data)) 2))
				 )
			     (printf "~A ~A ~A~%" x y z)
			     )))
		       points)
	     
	       (if (pair? subtrees)
		   (for-each (lambda (subtree) (recur subtree) )
			     subtrees))
	     )))

    )))



(define (main)

  (let ((operands          (opt '@)))

    (if (null? operands) (neurolucida:usage))

    (d "data directory is ~s~%" (get-data-dir))

    (let* ((data-map
	    (concatenate
	     (map (lambda (p)
		    (let* ((key             (or (opt 'key) (defopt 'key)))
			   (contents        (parse-sxml p))
			   (contours        (extract-element contents "contour" key))
			   (trees           (extract-element contents "tree" key)))
		      (let recur ((keys (delete-duplicates (map car contours)))
				  (elms (append contours trees)) 
				  (data-map '()))
			(if (null? keys)
			    (if (null? elms) data-map (error "elements with unknown color" elms))
			    (let-values (((lst rest)  (partition-by-key elms (car keys) key=?)))
			      (recur (cdr keys) rest (cons (cons (car keys) (map cadr lst)) data-map)))))
			))
		  operands)))
	   (ddir (get-data-dir))
	   )
      (let ((format (string->symbol (or (opt 'format) (defopt 'format)))))

	(case format
	  ((pts)
	   (for-each (lambda (f x) 
		       (with-output-to-file f
			 (lambda () (print-pts x))))
		     (map (lambda (x) (make-pathname ddir (car (string-split (cadar x) "#")) ".asc")) data-map)
		     data-map))

	  ((swc)
	   (for-each (lambda (f x) 
		       (with-output-to-file f
			 (lambda () (print-swc x))))
		     (map (lambda (x) (make-pathname ddir (car (string-split (cadar x) "#")) ".swc")) data-map)
		     data-map))
	  ))
      )))

(main)