(module gl-utils-ply (load-ply
                      load-ply-mesh)

(import
  scheme
  (chicken base)
  (chicken io)
  (chicken pathname)
  (chicken platform)
  (chicken port)
  (chicken keyword)
  (prefix epoxy #:gl)
  gl-utils-core
  gl-utils-bytevector
  gl-utils-mesh
  z3
  matchable
  srfi-42
  miscmacros
  srfi-1
  srfi-13)

;;; Ply
(define data-format (make-parameter #f))

(define (type->unsigned-type type)
  (case type
    ((char: int8: byte: uchar: uint8: unsigned-byte:) unsigned-byte:)
    ((short: int16: ushort: uint16: unsigned-short:) unsigned-short:)
    ((int: int32: integer: integer32: uint: uint32: unsigned-int: unsigned-int32:
           unsigned-integer: unsigned-integer32:)
     unsigned-int:)))

(define (type->setter type)
  (case type
    ((char: int8: byte:) bytevector-u8-set!)
    ((uchar: uint8: unsigned-byte:) bytevector-s8-set!)
    ((short: int16:) bytevector-s16-set!)
    ((ushort: uint16: unsigned-short:) bytevector-u16-set!)
    ((int: int32: integer: integer32:) bytevector-s32-set!)
    ((uint: uint32: unsigned-int: unsigned-int32:
	    unsigned-integer: unsigned-integer32:)
     bytevector-u32-set!)
    ((float: float32:) bytevector-f32-set!)
    ((double: float64:) bytevector-f64-set!)))

(define (bytes-per-element els)
  (foldl (lambda (count el)
           (+ count (type->bytes (second el))))
         0 els))

;; For ((name type) ...), return ((name b0 b1 b2) ... )
(define (get-bytes els)
  (let loop ((els els) (result '()) (n 0))
    (if (null? els)
        result
        (let ((bytes (type->bytes (cadar els))))
          (loop (cdr els)
                (cons (cons (caar els)
                            (list-ec (: i bytes) (+ n i)))
                      result)
                (+ n bytes))))))

(define (endian-swap?)
  (or (and (member little-endian: (features))
        (equal? (data-format) binary-big-endian:))
     (and (member big-endian: (features))
        (equal? (data-format) binary-little-endian:))))

(define (binary el spec)
  (match-let (((name n . vars) el))
    (if* (assoc name spec)
         (let* ((var-bytes (get-bytes vars))
                (mapping (concatenate
                          (list-ec (: el (second it))
                                   ((if (endian-swap?)
                                        reverse
                                        identity)
                                    (cdr (assoc el var-bytes))))))
                (stride (length mapping))
                (buffer (make-bytevector (* n stride))))
           (dotimes (vertex n)
             (let ((bytes (list-ec (: _ (bytes-per-element vars))
                                   (read-byte))))
               (for-each (lambda (m i)
                           (bytevector-u8-set! buffer (+ (* vertex stride) i)
                                               (list-ref bytes m)))
                         mapping
                         (iota stride))))
           buffer)
         (dotimes (_ (* n (bytes-per-element vars)))
           (read-byte)))))

(define (ascii el spec)
  (match-let (((name n . vars) el))
    (if* (assoc name spec)
         (let* ((var/type/bytes/i (map (lambda (v i)
                                         (list (car v)
                                               (second v)
                                               (type->bytes (second v))
                                               i))
                                       vars
                                       (iota (length vars))))
                (elements (map (lambda (el) (assoc el var/type/bytes/i))
                               (second it)))
                (mapping (map fourth elements))
                (offsets (reverse
                          (let loop ((lst (map third elements))
                                     (accum '())
                                     (count 0))
                            (if (null? lst)
                                accum
                                (loop (cdr lst)
                                      (cons count accum)
                                      (+ count (car lst)))))))
                (setters (map (lambda (el) (type->setter (second el)))
                              elements))
                (stride (fold + 0 (map third elements)))
                (buffer (make-bytevector (* n stride))))
           (dotimes (vertex n)
             (let ((values (list-ec (: _ (length vars)) (read))))
               (for-each (lambda (mapping offset setter)
                           (setter buffer 
                                   (+ (* vertex stride) offset)
                                   (list-ref values mapping)))
                         mapping
                         offsets
                         setters)))
           buffer)
         (dotimes (_ (* n (length vars)))
           (read)))))

(define n-face-vertices (make-parameter 0))

(define (binary-list el spec)
  (match-let (((name n (var ('list: list-type type))) el))
    (when (> (type->bytes list-type) 1)
      (error 'load-ply "Face list type must be one byte" list-type))
    (if* (assoc name spec)
         (let* ((buffer #f)
                (n-bytes (type->bytes type))
                (mapping ((if (endian-swap?)
                              reverse
                              identity)
                          (iota n-bytes))))
           (dotimes (face n)
             (let ((n-verts (read-byte)))
               (if (= (n-face-vertices) 0)
                   (begin
                     (n-face-vertices n-verts)
                     (set! buffer (make-bytevector (* n n-verts n-bytes))))
                   (when (not (= (n-face-vertices) n-verts))
                     (error 'load-ply "Number of elements must be constant in face list")))
               (dotimes (vertex n-verts)
                 (let ((bytes (list-ec (: b n-bytes)
                                       (read-byte))))
                   (for-each
                    (lambda (byte i)
                      (bytevector-u8-set! buffer 
                                          (+ (* vertex n-bytes)
                                             (* face n-verts n-bytes)
                                             i)
                                          byte))
                    bytes
                    mapping)))))
           buffer)
         (dotimes (_ n)
           (let ((n-verts (read-byte)))
             (if (= (n-face-vertices) 0)
                 (n-face-vertices n-verts)
                 (when (not (= (n-face-vertices) n-verts))
                   (error 'load-ply "Number of elements must be constant in face list")))
             (dotimes (_ (* n-verts (type->bytes type)))
               (read-byte)))))))

(define (ascii-list el spec)
  (match-let (((name n (var ('list: list-type type))) el))
    (if* (assoc name spec)
         (let* ((buffer #f)
                (n-bytes (type->bytes type))
                (setter (type->setter type)))
           (dotimes (face n)
             (let ((n-verts (read)))
               (if (= (n-face-vertices) 0)
                   (begin
                     (n-face-vertices n-verts)
                     (set! buffer (make-bytevector (* n n-verts n-bytes))))
                   (when (not (= (n-face-vertices) n-verts))
                     (error 'load-ply "Number of elements must be constant in face list")))
               (dotimes (vertex n-verts)
                 (setter buffer (+ (* vertex n-bytes)
                                   (* face n-verts n-bytes))
                         (read)))))
           buffer)
         (dotimes (_ n)
           (let ((n-verts (read)))
             (if (= (n-face-vertices) 0)
                 (n-face-vertices n-verts)
                 (when (not (= (n-face-vertices) n-verts))
                   (error 'load-ply "Number of elements must be constant in face list")))
             (dotimes (_ n-verts) (read)))))))

(define (get-buffers elements spec)
  (let ((buffers (list-ec (: el elements)
                   (cons (first el)
                         (if (list? (car (cdaddr el)))
                             (if (equal? (data-format) ascii:)
                                 (ascii-list el spec)
                                 (binary-list el spec))
                             (if (equal? (data-format) ascii:)
                                 (ascii el spec)
                                 (binary el spec)))))))
    ; Return buffers in order of spec
    (list-ec (: s spec)
             (cdr (assoc (car s) buffers)))))

(define (check-spec spec elements)
  (define (err el) (error 'load-ply "buffer-spec does not match ply" el))
  (for-each (lambda (s)
              (if* (assoc (first s) elements)
                   (for-each (lambda (var)
                               (unless (assoc var (cddr it))
                                 (err var)))
                             (if (list? (second s))
                                 (second s)
                                 (list (second s))))
                   (err (first s))))
            spec))

(define (load-ply file buffer-spec)
  (define (err) (error 'load-ply "Not a valid PLY file:" file))
  (define (parse-elements)
    (reverse
     (map reverse
          (let loop ((elements '()))
            (match (string-tokenize (read-line))
              (("comment" . rest) (loop elements))
              (("element" name n)
               (loop (cons (list (string->number n) (string->keyword name))
                           elements)))
              (("property" "list" list-type type name)
               (loop (cons (cons (list (string->symbol name)
                                       (list list:
                                             (string->keyword list-type)
                                             (string->keyword type)))
                                 (car elements))
                           (cdr elements))))
              (("property" type name)
               (loop (cons (cons (list (string->symbol name) (string->keyword type))
                                 (car elements))
                           (cdr elements))))
              (("end_header") elements)
              (other (err)))))))
  (with-input-from-port (if (equal? (pathname-extension file) "gz")
                            (z3:open-compressed-input-file file)
                            (open-input-file file))
    (lambda ()
      (match (read-line)
        ("ply" #t)
        (else (err)))
       (match (string-tokenize (read-line))
         (("format" "ascii" version) (data-format #:ascii))
         (("format" "binary_little_endian" version)
          (data-format #:binary-little-endian))
         (("format" "binary_big_endian" version)
          (data-format #:binary-big-endian))
         (else (err)))
       (let ((elements (parse-elements)))
         (check-spec buffer-spec elements)
         (values (get-buffers elements buffer-spec)
                 elements)))))

(define (verts->primitive verts)
  (case verts
    ((1) points:)
    ((2) lines:)
    ((3) triangles:)
    (else (error 'load-ply-mesh "Unsupported number of vertices" verts))))

(define (load-ply-mesh ply #!key vertex face)
  (define (attribute-vars attribute elements)
    (let* ((vertices (cddr (assoc vertex: elements)))
           (type (let loop ((els (cddr attribute))
                            (type (second (assoc (cadr attribute) vertices))))
                   (if (null? els)
                       type
                       (if (equal? type (second (assoc (car els) vertices)))
                           (loop (cdr els) type)
                           (error 'load-ply-mesh "Properties of the same attribute must have the same type"
                                  attribute))))))
      (list (car attribute)
            type
            (length (cdr attribute))
            normalized: #t)))
  (let ((buffer-spec `((vertex: ,(flatten (map cdr vertex)))
                       (face: ,face))))
    (let-values (((buffers elements) (load-ply ply buffer-spec)))
      (unless (and (assoc face: elements)
                 (assoc vertex: elements))
        (error 'load-ply-mesh "Ply must contain vertex and face elements" ply))
      (let* ((face (cdr (assoc face: elements)))
             (primitive-type (verts->primitive (n-face-vertices)))
             (element-type (type->unsigned-type (third (cadadr face))))
             (attributes (map (lambda (v) (attribute-vars v elements))
                              vertex)))
        (make-mesh vertices: `(attributes: ,attributes
                               initial-elements: ,(car buffers))
                   indices: `(type: ,element-type
                              initial-elements: ,(cadr buffers))
                   mode: primitive-type)))))

)