#lang scheme/base
(require scheme/class
         scheme/contract)

;; This is an implementation of the Octree Quantization algorithm.  This implementation 
;; follows the sketch in:
;;
;; Dean Clark.  Color Quantization using Octrees.  Dr. Dobbs Portal, January 1, 1996.
;; http://www.ddj.com/184409805


;; quantize: bytes (integer-in 1 255) -> (values bytes gif-colormap color)
;; Adaptively quantizes 24-bit image data to fit within, at most, 256 colors
;; (including the transparent color.)  The signature is meant to match the
;; function provided by file/gif.
(define (quantize argb #:palette [palette 255])
  (let* ([an-octree (make-octree-from-argb argb palette)]
         [len (bytes-length argb)]
         [dest (make-bytes (/ len 4))])
    (for ([i (in-range 0 len 4)])
      (let ([a (bytes-ref argb i)]
            [r (bytes-ref argb (+ i 1))]
            [g (bytes-ref argb (+ i 2))]
            [b (bytes-ref argb (+ i 3))]
            [n (arithmetic-shift i -2)])
        (cond
          [(alpha-opaque? a)
           (let ([palette-index (octree-lookup-index an-octree r g b)])
             (bytes-set! dest n palette-index))]
          [else
           (bytes-set! dest n TRANSPARENT-INDEX)])))
    (values dest (vector->list (octree-palette an-octree)) TRANSPARENT-INDEX)))


;; make-octree-from-argb: bytes positive-number -> octree
;; Constructs an octree ready to quantize the colors from an-argb.
(define (make-octree-from-argb an-argb n)
  (unless (> n 0)
    (raise-type-error 'make-octree-from-argb "positive number" n))
  (let ([an-octree (new-octree)]
        [len (bytes-length an-argb)])
    (let make-octree-loop ([i 0])
      (when (< i len)
        (let ([a (bytes-ref an-argb i)]
              [r (bytes-ref an-argb (+ i 1))]
              [g (bytes-ref an-argb (+ i 2))]
              [b (bytes-ref an-argb (+ i 3))])
          (when (alpha-opaque? a)
            (octree-insert-color! an-octree r g b)
            (let reduction-loop ()
              (when (> (octree-leaf-count an-octree) n)
                (octree-reduce! an-octree)
                (reduction-loop)))))
        (make-octree-loop (+ i 4))))
    (octree-finalize! an-octree)
    an-octree))


;; alpha-opaque? byte -> boolean
;; Returns true if the alpha value is considered opaque.
(define (alpha-opaque? a)
  (>= a 128))

(define TRANSPARENT-INDEX 0)

;; quantize-bitmap: bitmap positive-number -> bitmap
;; Given a bitmap, returns a new bitmap quantized to, at most, n colors.
#;
(define (quantize-bitmap bm n)
  (let* ([width (send bm get-width)]
         [height (send bm get-height)]
         [len (* width height 4)]
         [source-buffer (make-bytes len)]
         [_ (send bm get-argb-pixels 0 0 width height source-buffer)]
         [an-octree (make-octree-from-argb source-buffer n)]
         [dest-buffer (make-bytes len)])
    (let quantize-bitmap-loop ([i 0])
      (when (< i len)
        (let* ([i+1 (+ i 1)]
               [i+2 (+ i 2)]
               [i+3 (+ i 3)]
               [a (bytes-ref source-buffer i)]
               [r (bytes-ref source-buffer i+1)]
               [g (bytes-ref source-buffer i+2)]
               [b (bytes-ref source-buffer i+3)])
          (cond
            [(alpha-opaque? a)
             (let-values ([(new-r new-g new-b)
                           (octree-lookup an-octree r g b)])
               (bytes-set! dest-buffer i 255)
               (bytes-set! dest-buffer i+1 new-r)
               (bytes-set! dest-buffer i+2 new-g)
               (bytes-set! dest-buffer i+3 new-b))]
            [else
             (bytes-set! dest-buffer i 0)
             (bytes-set! dest-buffer i+1 0)
             (bytes-set! dest-buffer i+2 0)
             (bytes-set! dest-buffer i+3 0)]))
        (quantize-bitmap-loop (+ i 4))))
    (let* ([new-bm (make-object bitmap% width height)]
           [dc (make-object bitmap-dc% new-bm)])
      (send dc set-argb-pixels 0 0 width height dest-buffer)
      (send dc set-bitmap #f)
      new-bm)))

;; The maximum level height of an octree.
(define MAX-LEVEL 7)

(define octet/c (integer-in 0 7))


;; A color is a (vector byte byte byte)

;; An octree is a:
(define-struct octree (root            ; node
                       leaf-count      ; number
                       reduction-heads ; (vectorof (or/c node #f))
                       palette)        ; (vectorof (or/c color #f))
  #:mutable)
;; reduction-heads is used to accelerate the search for a reduction candidate.


;; A subtree node is a:
(define-struct node (leaf?          ; bool
                     npixels        ; number  -- number of pixels this subtree node represents
                     redsum         ; number
                     greensum       ; number
                     bluesum        ; number
                     children       ; (vectorof (or/c #f node))
                     next           ; (or/c #f node)
                     palette-index) ; (or/c #f byte?)
  #:mutable)
;; node-next is used to accelerate the search for a reduction candidate.


;; new-octree: -> octree
(define (new-octree)
  (let* ([root-node (make-node #f ;; not a leaf
                               0  ;; no pixels under us yet
                               0  ;; red sum
                               0  ;; green sum
                               0  ;; blue sum
                               (make-vector 8 #f) ;; no children so far
                               #f ;; next
                               #f ;; palette-index
                               )]
         [an-octree
          (make-octree root-node
                       0 ; no leaves so far
                       (make-vector (add1 MAX-LEVEL) #f) ; no reductions so far
                       (make-vector 256 #(0 0 0)))])        ; the palette
    ;; Although we'll almost never reduce to this level, initialize the first
    ;; reducible node to the root, for completeness sake.
    (vector-set! (octree-reduction-heads an-octree) 0 root-node)
    an-octree))


;; rgb->index: natural-number byte byte byte -> octet
;; Given a level and an (r,g,b) triplet, returns an octet that can be used
;; as an index into our octree structure.
(define (rgb->index level r g b)
  (bitwise-ior (bitwise-and 4 (arithmetic-shift r (- level 5)))
               (bitwise-and 2 (arithmetic-shift g (- level 6)))
               (bitwise-and 1 (arithmetic-shift b (- level 7)))))


;; octree-insert-color!: octree byte byte byte -> void
;; Accumulates a new r,g,b triplet into the octree.
(define (octree-insert-color! an-octree r g b)
  (node-insert-color! (octree-root an-octree) an-octree r g b 0))


;; node-insert-color!: node octree byte byte byte natural-number -> void
;; Adds a color to the node subtree.  While we hit #f, we create new nodes.
;; If we hit an existing leaf, we accumulate our color into it.
(define (node-insert-color! a-node an-octree r g b level)
  (let insert-color-loop ([a-node a-node]
                          [level level])
    (cond [(node-leaf? a-node)
           ;; update the leaf with the new color
           (set-node-npixels! a-node (add1 (node-npixels a-node)))
           (set-node-redsum! a-node (+ (node-redsum a-node) r))
           (set-node-greensum! a-node (+ (node-greensum a-node) g))
           (set-node-bluesum! a-node (+ (node-bluesum a-node) b))]
          [else
           ;; create the child node if necessary
           (let ([index (rgb->index level r g b)])
             (unless (vector-ref (node-children a-node) index)
               (let ([new-node (make-node (= level MAX-LEVEL) ; leaf?
                                          0  ; npixels
                                          0  ; redsum
                                          0  ; greensum
                                          0  ; bluesum
                                          (make-vector 8 #f) ; no children yet
                                          #f ; and no next node yet
                                          #f ; or palette index
                                          )])
                 (vector-set! (node-children a-node) index new-node)
                 (cond
                   [(= level MAX-LEVEL)
                    ;; If we added a leaf, mark it in the octree.
                    (set-octree-leaf-count! an-octree
                                            (add1 (octree-leaf-count an-octree)))]
                   [else
                    ;; Attach the node as a reducible node if it's interior.
                    (set-node-next!
                     new-node (vector-ref (octree-reduction-heads an-octree)
                                          (add1 level)))
                    (vector-set! (octree-reduction-heads an-octree) 
                                 (add1 level)
                                 new-node)])))
             ;; and recur on the child node.
             (insert-color-loop (vector-ref (node-children a-node) index) 
                                (add1 level)))])))


;; octree-reduce!: octree -> void
;; Reduces one of the subtrees, collapsing the children into a single node.
(define (octree-reduce! an-octree)
  (node-reduce! (pop-reduction-candidate! an-octree) an-octree))


;; node-reduce!: node octree -> void
;; Reduces the interior node.
(define (node-reduce! a-node an-octree)
  (for ([child (in-vector (node-children a-node))]
        #:when child)
    (set-node-npixels! a-node (+ (node-npixels a-node)
                                 (node-npixels child)))
    (set-node-redsum! a-node (+ (node-redsum a-node)
                                (node-redsum child)))
    (set-node-greensum! a-node (+ (node-greensum a-node)
                                  (node-greensum child)))
    (set-node-bluesum! a-node (+ (node-bluesum a-node)
                                 (node-bluesum child)))
    (set-octree-leaf-count! an-octree (sub1 (octree-leaf-count an-octree))))
  (set-node-leaf?! a-node #t)
  (set-octree-leaf-count! an-octree (add1 (octree-leaf-count an-octree))))


;; find-reduction-candidate!: octree -> node
;; Returns a bottom-level interior node for reduction.  Also takes the
;; candidate out of the conceptual queue of reduction candidates.
(define (pop-reduction-candidate! an-octree)
  (let loop ([i MAX-LEVEL])
    (cond
      [(vector-ref (octree-reduction-heads an-octree) i)
       =>
       (lambda (candidate-node)
         (when (> i 0)
           (vector-set! (octree-reduction-heads an-octree) i
                        (node-next candidate-node)))
         candidate-node)]
      [else
       (loop (sub1 i))])))


;; octree-finalize!: octree -> void
;; Finalization does a few things:
;; * Walks through the octree and reduces any interior nodes with just one leaf child.
;;   Optimizes future lookups.
;; * Fills in the palette of the octree and the palette indexes of the leaf nodes.
;; * Note: palette index 0 is always reserved for the transparent color.
(define (octree-finalize! an-octree)
  ;; Collapse one-leaf interior nodes.
  (let loop ([a-node (octree-root an-octree)])
    (for ([child (in-vector (node-children a-node))]
          #:when (and child (not (node-leaf? child))))
      (loop child)
      (when (interior-node-one-leaf-child? a-node)
        (node-reduce! a-node an-octree))))

  ;; Attach palette entries.
  (let ([current-palette-index 1])
    (let loop ([a-node (octree-root an-octree)])
      (cond [(node-leaf? a-node)
             (let ([n (node-npixels a-node)])
               (vector-set! (octree-palette an-octree) current-palette-index
                            (vector (quotient (node-redsum a-node) n)
                                    (quotient (node-greensum a-node) n)
                                    (quotient (node-bluesum a-node) n)))
               (set-node-palette-index! a-node current-palette-index)
               (set! current-palette-index (add1 current-palette-index)))]
            [else
             (for ([child (in-vector (node-children a-node))]
                   #:when child)
               (loop child))]))))


;; interior-node-one-leaf-child?: node -> boolean
(define (interior-node-one-leaf-child? a-node)
  (let ([child-list (filter values (vector->list (node-children a-node)))])
    (and (= (length child-list) 1)
         (node-leaf? (car child-list)))))


;; octree-lookup: octree byte byte byte -> (values byte byte byte)
;; Returns the palettized color.
(define (octree-lookup an-octree r g b)
  (let* ([index (node-lookup-index (octree-root an-octree) an-octree r g b 0)]
         [vec (vector-ref (octree-palette an-octree) index)])
    (values (vector-ref vec 0)
            (vector-ref vec 1)
            (vector-ref vec 2))))


;; octree-lookup-index: octree byte byte byte -> byte
;; Returns the paletized color as an index into the octree-palette vector.
(define (octree-lookup-index an-octree r g b)
  (node-lookup-index (octree-root an-octree) an-octree r g b 0))


;; node-lookup-index: node byte byte byte natural-number -> byte
;; Returns the palettized color index.
(define (node-lookup-index a-node an-octree r g b level)
  (let loop ([a-node a-node]
             [level level])
    (if (node-leaf? a-node)
      (node-palette-index a-node)
      (let ([child (vector-ref (node-children a-node) (rgb->index level r g b))])
        (unless child
          (error 'node-lookup-index
                 "color (~a, ~a, ~a) not previously inserted"
                 r g b))
        (loop child (add1 level))))))


(define color? byte?)

(define (gif-colormap? l)
  (and (list? l)
       (member (length l) '(2 4 8 16 32 64 128 256))
       (andmap (lambda (c)
                 (and (vector? c)
                      (= 3 (vector-length c))
                      (color? (vector-ref c 0))
                      (color? (vector-ref c 1))
                      (color? (vector-ref c 2))))
               l)))

(define (argb-bytes? b)
  (and (bytes? b)
       (zero? (remainder (bytes-length b) 4))))


(provide/contract
 [quantize ((argb-bytes?) (#:palette [integer-in 1 255])
            . ->* . (values bytes? gif-colormap? color?))]
 #;
 [quantize-bitmap
  ([is-a?/c bitmap%] natural-number/c . -> . [is-a?/c bitmap%])]

 [struct octree
         ([root node?]
          [leaf-count natural-number/c]
          [reduction-heads (vectorof (or/c node? false/c))]
          [palette (vectorof (vector/c byte? byte? byte?))])]
 [struct node
         ([leaf? boolean?]
          [npixels natural-number/c]
          [redsum natural-number/c]
          [greensum natural-number/c]
          [bluesum natural-number/c]
          [children (vectorof (or/c false/c node?))]
          [next (or/c node? false/c)]
          [palette-index byte?])]
 [new-octree (-> octree?)]
 [octree-insert-color! (octree? byte? byte? byte? . -> . any)]
 [octree-lookup
  (octree? byte? byte? byte? . -> . (values byte? byte? byte?))]
 [octree-reduce! (octree? . -> . any)]
 [octree-finalize! (octree? . -> . any)]
 [rgb->index (natural-number/c byte? byte? byte? . -> . octet/c)])