;;; -*- mode: lisp; syntax: common-lisp -*-
;;;
;;; GRIMP - The Grayscale Image Manipulation Program
;;;
;;; by Peter Salvi, 2007

(defpackage grimp
  (:use :cl :iterate)
  (:export :make-binary-image :make-gray-image
           :read-pbm          :read-pgm
           :blur
           :copy
           :enhance-contrast
           :histogram
           :horizontal-flip
           :monochrome
           :write-image))

(in-package :grimp)


;;; Classes

(defclass image ()
  ((width  :initarg :width  :reader   width)
   (height :initarg :height :reader   height)
   (data   :initarg :data   :accessor data)))

(defmethod print-object ((object image) stream)
  (print-unreadable-object (object stream :type t)
    (with-slots (width height) object
      (format stream "~dx~d" width height))))

(defclass binary-image (image)
  ())

(defun make-binary-image (width height &optional data)
  (let ((data (or data (make-array (* (ceiling width 8) height)
                                   :initial-element 0))))
    (make-instance 'binary-image :width width :height height :data data)))

(defclass gray-image (image)
  ((grays :initarg :grays :reader grays)))

(defun make-gray-image (width height grays &optional data)
  (let ((data (or data (make-array (* width height) :initial-element 0))))
    (make-instance 'gray-image
                   :width width :height height :grays grays :data data)))


;;; Input/Output

(defun read-without-comment (s)
  "Reads the next data from stream S, skipping any comments."
  (if (eq (peek-char nil s) #\#)
      (progn
        (read-line s)
        (read-without-comment s))
      (read s)))

(defun read-pbm (filename)
  "Read a binary PBM file and returns a BINARY-IMAGE."
  (let (width height pos)
    (with-open-file (s filename)
      (when (not (equal (symbol-name (read s)) "P4"))
        (error "not a binary PBM file" ))
      (setf width  (read-without-comment s)
            height (read-without-comment s)
            pos    (file-position s)))
    (let ((data (make-array (* (ceiling width 8) height)
                            :element-type '(unsigned-byte 8))))
      (with-open-file (s filename :element-type '(unsigned-byte 8))
        (file-position s pos)
        (read-sequence data s))
      (make-binary-image width height data))))

(defun read-pgm (filename)
  "Reads a binary PGM file and returns a GRAY-IMAGE."
  (let (width height grays pos)
    (with-open-file (s filename)
      (when (not (equal (symbol-name (read s)) "P5"))
        (error "not a binary PGM file" ))
      (setf width  (read-without-comment s)
            height (read-without-comment s)
            grays  (read-without-comment s)
            pos    (file-position s)))
    (let ((data (make-array (* width height)
                            :element-type '(unsigned-byte 8))))
      (with-open-file (s filename :element-type '(unsigned-byte 8))
        (file-position s pos)
        (read-sequence data s))
      (make-gray-image width height grays data))))

(defgeneric write-image (image filename)
  (:documentation "Writes IMAGE to FILENAME."))

(defmethod write-image ((image binary-image) filename)
  (with-open-file (s filename :direction :output :if-exists :supersede)
    (format s "P4~%# Generated by GRIMP~%~d ~d~%"
            (width image) (height image)))
  (with-open-file (s filename :direction :output :if-exists :append
                     :element-type '(unsigned-byte 8))
    (write-sequence (data image) s))
  t)

(defmethod write-image ((image gray-image) filename)
  (with-open-file (s filename :direction :output :if-exists :supersede)
    (format s "P5~%# Generated by GRIMP~%~d ~d~%~d~%"
            (width image) (height image) (grays image)))
  (with-open-file (s filename :direction :output :if-exists :append
                     :element-type '(unsigned-byte 8))
    (write-sequence (data image) s))
  t)


;;; Low-level image manipulation

(defgeneric copy (image)
  (:documentation "Creates a copy of IMAGE."))

(defmethod copy ((image binary-image))
  (make-binary-image (width image) (height image) (copy-seq (data image))))

(defmethod copy ((image gray-image))
  (make-gray-image (width image) (height image) (grays image)
                   (copy-seq (data image))))

(defgeneric pos (image i j)
  (:documentation "Returns the data at position (I, J) of IMAGE."))

(defmethod pos ((image binary-image) i j)
  (let ((byte-width (ceiling (width image) 8)))
    (= (ldb (byte 1 (- 7 (mod i 8)))
            (elt (data image) (+ (* j byte-width) (floor i 8))))
       1)))

(defmethod pos ((image gray-image) i j)
  (elt (data image) (+ (* j (width image)) i)))

(defgeneric (setf pos) (value image i j)
  (:documentation "Sets the data at position (I, J) of IMAGE to VALUE."))

(defmethod (setf pos) (value (image binary-image) i j)
  (let ((byte-width (ceiling (width image) 8)))
    (setf (ldb (byte 1 (- 7 (mod i 8)))
               (elt (data image) (+ (* j byte-width) (floor i 8))))
          (if value 1 0))))

(defmethod (setf pos) (value (image gray-image) i j)
  (setf (elt (data image) (+ (* j (width image)) i)) value))


;;; High-level functions

(defun horizontal-flip (image)
  "Flips the image horizontally."
  (let ((result (copy image)))
    (with-slots (width height)
        result
      (iter (for j from 0 below height)
            (iter (for i from 0 below (floor width 2))
                  (rotatef (pos result i j) (pos result (- width i 1) j)))))
    result))

(defun blur (gray-image &optional (iterations 5))
  "Blurs the image with a box filter ITERATIONS times."
  (let ((result (copy gray-image)))
    (with-slots (width height data)
        result
      (iter (repeat iterations)
            (with tmp = (copy-seq data))
            (iter (for j from 1 below (1- height))
                  (iter (for i from 1 below (1- width))
                        (setf (elt tmp (+ (* j width) i))
                              (floor (+ (elt data (+ (* j width) (1- i)))
                                        (elt data (+ (* j width) (1+ i)))
                                        (elt data (+ (* (1- j) width) i))
                                        (elt data (+ (* (1+ j) width) i)))
                                     4))))
            (map-into data #'identity tmp)))
    result))

(defun monochrome (gray-image &optional (threshold 0.5))
  "Convert GRAY-IMAGE to a BINARY-IMAGE, using 0 < THRESHOLD < 1."
  (with-slots (width height grays)
      gray-image
    (let ((thresh (floor (* threshold grays)))
          (result (make-binary-image width height)))
      (iter (for j from 0 below height)
            (iter (for i from 0 below width)
                  (setf (pos result i j) (<= (pos gray-image i j) thresh))))
      result)))

(defun frequency (gray-image)
  "Returns a vector that contains the number of pixels of each gray value."
  (with-slots (width height grays)
      gray-image
    (let ((result (make-array (1+ grays) :initial-element 0)))
      (iter (for j from 0 below height)
            (iter (for i from 0 below width)
                  (incf (elt result (pos gray-image i j)))))
      result)))

(defun histogram (gray-image)
  "Creates a BINARY-IMAGE from the histogram of GRAY-IMAGE."
  (let* ((histogram (frequency gray-image))
         (width     (length histogram))
         (max-value (iter (for x in-vector histogram) (maximize x)))
         (result    (make-binary-image width 101)))
    (iter (for i from 0 below width)
          (for val = (- 100 (floor (* (elt histogram i) 100) max-value)))
          (iter (for j from 0 below (1- val))
                (setf (pos result i j) nil))
          (iter (for j from val to 100)
                (setf (pos result i j) t)))
    result))

(defun enhance-contrast (gray-image &key low high (percent 10))
  "Enhances contrast by pulling the LOW grays to black and HIGHs to white.
LOW and HIGH are given in the range [0, 1]. If LOW or HIGH are not supplied,
they are the first place where the frequency of grayscale points reaches
PERCENT% of the maximum."
  (with-slots (width height grays)
      gray-image
    (let* ((histogram (frequency gray-image))
           (max-value (iter (for x in-vector histogram) (maximize x)))
           (low       (or (and low (floor (* low grays)))
                          (position (floor (* percent max-value) 100)
                                    histogram :test #'<=)))
           (high      (or (and high (floor (* high grays)))
                          (position (floor (* percent max-value) 100)
                                    histogram :test #'<= :from-end t)))
           (result    (make-gray-image width height grays)))
      (map-into (data result)
                #'(lambda (x)
                    (floor (* (- (max low (min high x)) low) grays)
                           (- high low)))
                (data gray-image))
      result)))