#lang scheme/base

(require
 "../tools.ss"
 scheme/match
 "output-process.ss")
;; (provide
;;  read-pgm  load-pgm
;;  write-pgm save-pgm
;;  make-yuvplay
;;  make-yuv4mpeg-writer
;;  )

(provide (all-defined-out))

;; PGM files

;; (define (read-pgm-line)
;;   (let ((line (read-bytes-line)))
;;     (if (and (not eof? line)
;;              (eq? 35 (car (bytes->list line))))
;;         (read-pgm-line)
;;         line)))


;; Image = 3D array encoded in a 1D vector.
(define-struct image (width height colors data))

;; Because shifts (up,down,left,right) are common in filter
;; algorithms, it's not a good idea to represent an image as a list of
;; lists. Random access vectors accomodate needs better.

(define (new-image w h [c 1] [v 0])
  (make-image w h c (make-vector (* w h c) v)))
(define clone-image
  (match-lambda
   ((struct image (w h c d))
    (new-image w h c))))

(define (image-size img)
  (* (image-width img)
     (image-height img)
     (image-colors img)))


    
;; HOFs

;; The idea is to push the separation of concerns to the max, and
;; ignore the efficiency concern completely. This should lead to a
;; small number of primitives in the 'algebra of programs' for image
;; processors, which then can be used to build a specialized compiler
;; which moves some of the parameterization to compile time.


;; MAP: The unary case is trivial, but the binary case has 3
;; essentially different primitive forms:
;;
;;   2->1
;;   1->1 X shift
;;   1->1 Y shift

(define (unary-loop! i fn r a )
  (vector-set! r i
               (fn (vector-ref a i))))

(define (binary-loop! i i->j fn r a b)
  (vector-set! r i
               (fn (vector-ref a i)
                   (vector-ref b (i->j i)))))




;; image-data and image-clone binding red tape
(define-syntax-rule (let-images ((r img-r)  ;; result image
                                 (a img-a)  ;; first one = prototype of result
                                 (p img-p) ...) . body)
  (lambda (img-a img-p ...)
    (let ((img-r (clone-image img-a)))
      (let ((r (image-data img-r))
            (a (image-data img-a))
            (p (image-data img-p)) ...)
        . body)
      img-r)))

(define (image-map-u fn)
  (let-images ((r img-r)
               (a img-a))
    (for ((i (in-range (image-size img-r))))
      (unary-loop! i fn r a))))
                     
(define (image-map-b fn)
  (let-images ((r img-r)
               (a img-a)
               (b img-b))
    (for ((i (in-range (image-size img-r))))
      (binary-loop! i id fn r a b))))

(define (make-image-map-shifted image->stride)
  (lambda (fn)
    (let-images ((r img-r)
                 (a img-a))
      (define w (image-width img-r))
      (define size (image-size img-r))
      (define stride (image->stride img-r))
      (define (neighbour i) (modulo (+ stride i) size))
      ;; (printf "stride: ~s\n" stride)
      (for ((i (in-range size)))
           (binary-loop! i neighbour fn r a a)))))


(define image-map-y (make-image-map-shifted image-width))
(define image-map-x (make-image-map-shifted (lambda _ 1)))

;; FOLD: Several folds are possible. For now, only one is implemented:
;; fold over a single image with x,y coordinates passed in (i.e. to
;; compute moments).

(define (image-fold/moment fn)
  (lambda (init-state img-a)
    (let ((a (image-data img-a))
          (w (image-width img-a))
          (h (image-height img-a)))
        (for/fold  ((state init-state))
                   ((y (in-range h)))
          (for/fold  ((state state))
                     ((x (in-range w)))
            (fn state
                (vector-ref a (+ x (* y w)))
                x y))))))

  
    
;; Shortcuts:

;; Different map HOFs bring binary/unary scalar operators to
;; binary/unary image operators. These are all curried: they return a
;; lifted operation.

(define U image-map-u)
(define B image-map-b)
(define X image-map-x)
(define Y image-map-y)

;; Fold style operator transformers.

;; FM: (init-state fn) -> init-state
;; fn: (state piksel x y) -> state
(define FM image-fold/moment)  ;; For implementing moment-style integrals

    
        
;; INPUT/OUTPUT       


(define (image-bytes img)
  (list->bytes
   (map (lambda (x)
          (inexact->exact
           (floor
            (cond
             ((> x 255) 255)
             ((< x 0) 0)
             (else x)))))
        (vector->list (image-data img)))))

(define (write-image img)
  (write-bytes (image-bytes img)))

(define (image-out fn)
  (lambda (img [p (current-output-port)])
    (parameterize ((current-output-port p))
      (fn img
          (image-width img)
          (image-height img)
          255))))

;; PGM

(define write-pgm
  (image-out
   (lambda (img width height levels)
     (printf "P5\n# CREATOR: image-io.ss\n~a ~a\n~a\n" width height levels)
     (write-image img)
     (flush-output))))

(define (image->bytes img)
  (let ((port (open-output-bytes)))
    (write-pgm img port)
    (get-output-bytes port)))


;; This is a hack to read just the format I need, which has a comment
;; on the 2nd line

(define (read-pgm [p (current-input-port)])
  (parameterize ((current-input-port p))
    (unless (eq? 'P5 (read))
      (error 'not-a-pgm-file))
    (read-line) ;; newline
    (read-line) ;; comment line
    (let* ((width (read))
           (height (read))
           (levels (read))
           (out (new-image width height))
           (data (image-data out)))
      
      (read-line) ;; newline
      ;; (printf "pgm: ~a ~a ~a\n" width height levels)
      ;; The rest is binary data
      (let ((i 0))
        (for ((y (in-range height)))
          (for ((b (read-bytes width)))
            (vector-set! data i b)
            (inc! i))))
      out)))


(define (load-pgm file)
  (with-input-from-file file
    (lambda () (read-pgm))))

(define (save-pgm file img)
  (with-output-to-file file
    (lambda () (write-pgm img))
    #:exists 'replace))



;; YUV4MPEG

(define (make-yuv4mpeg-writer)
  (define plane #f)
  (define write-frame
    (image-out
     (lambda (img . _)
       (printf "FRAME\n")
       (write-image img)
       (plane)
       (plane)
       (flush-output))))

  (define write-first-frame
    (image-out
     (lambda (img width height levels)
       (printf "YUV4MPEG2 W~a H~a F25:1\n" width height)
       (let ((black (make-bytes (/ (* width height) 4) #x80)))
         (set! plane (lambda () (write-bytes black))))
       (write-frame img))))
    
  (image-out
   (lambda (img . _)
     (if plane 
         (write-frame img)
         (write-first-frame img)))))

(define (make-yuv4mpeg-out make-port)
  (define port #f)
  (define w -1)
  (define h -1)
  (define write-frame #f)
  (lambda (img)
    (let ((_w (image-width img))
          (_h (image-height img)))
      (unless
          (and (= w _w)
               (= h _h))
        (and port (close-output-port port))
        (set! port (make-port))
        (set! write-frame (make-yuv4mpeg-writer))
        (set! w _w)
        (set! h _h)))
    (write-frame img port)))

(define (make-yuv4mpeg-process . cmdline)
  (make-yuv4mpeg-out (lambda () (apply open-output-process cmdline))))

(define (make-yuvplay) (make-yuv4mpeg-process "yuvplay"))

;; debug