; (c) 2008 Dmitri Hrapof <hrapof@common-lisp.ru>
; This file is distributed under the terms of the LLGPL.

(defpackage oci-rest
   (:use common-lisp oci cffi))
(in-package oci-rest)

(define-foreign-library kissfft
  (:unix "libkissfft.so")
  (t (:default "libkissfft")))

(use-foreign-library kissfft)

(defcfun "kiss_fft_alloc" :pointer (nfft :int) (inverse :int) (mem :pointer) (lenmem :pointer))
(defcfun "kiss_fft" :void (cfg :pointer) (fin :pointer) (fout :pointer))

(defcfun "kiss_fftr_alloc" :pointer (nfft :int) (inverse :int) (mem :pointer) (lenmem :pointer))
(defcfun "kiss_fftr" :void (cfg :pointer) (fin :pointer) (fout :pointer))

(defcfun "kiss_fftnd_alloc" :pointer (dims :pointer) (ndims :int) (inverse :int) (mem :pointer) (lenmem :pointer))
(defcfun "kiss_fftnd" :void (cfg :pointer) (fin :pointer) (fout :pointer))

(defcfun "kiss_fftndr_alloc" :pointer (dims :pointer) (ndims :int) (inverse :int) (mem :pointer) (lenmem :pointer))
(defcfun "kiss_fftndr" :void (cfg :pointer) (fin :pointer) (fout :pointer))

(defcfun "hr_mag_inv_phase" :void (height :int) (width :int) (ft1 :pointer) (ft2 :pointer) (mags :pointer))
(defcfun "hr_mag_fft" :void (height :int) (width :int) (inverse :int) (ft1 :pointer) (ft2 :pointer) (mags :pointer))
(defcfun "hr_rotate" :void (height :int) (width :int) (angle :float) (in :pointer) (out :pointer))


(defcstruct kiss-fft-cpx
  (r :float)
  (i :float))

(defmacro with-kiss ((cfg-var init-form) &body body)
  `(let ((,cfg-var ,init-form))
     (unwind-protect
	  (progn ,@body)
       (foreign-free ,cfg-var))))

(defun c-fft! (complex-vector)
  (declare (optimize (speed 3) (safety 0)))
  (let ((n (length complex-vector)))
    (with-kiss (kfg (kiss-fft-alloc n 0 (null-pointer) (null-pointer)))
      (with-foreign-object (fin 'kiss-fft-cpx n)
	(dotimes (k n)
	  (with-foreign-slots 
	      ((r i) (mem-aref fin 'kiss-fft-cpx k) kiss-fft-cpx)
	    (setf r (coerce (realpart (aref complex-vector k)) 'single-float)
		  i (coerce (imagpart (aref complex-vector k)) 'single-float))))
	(with-foreign-object (fout 'kiss-fft-cpx n)
	  (kiss-fft kfg fin fout)
	  (dotimes (k n complex-vector)
	    (with-foreign-slots 
		((r i) (mem-aref fout 'kiss-fft-cpx k) kiss-fft-cpx)
	      (setf (aref complex-vector k) (complex r i)))))))))

(defun c-fft-2d! (image &optional inverse)
  (declare (optimize (speed 3) (safety 0)))
  (let ((h (array-dimension image 0))
	(w (array-dimension image 1))
	(s (array-total-size image)))
    (with-foreign-object (dims :int 2)
      (setf (mem-aref dims :int 0) h (mem-aref dims :int 1) w)
      (with-kiss (kfg (kiss-fftnd-alloc 
		       dims 2 (if inverse 1 0) (null-pointer) (null-pointer)))
	(with-foreign-object (fin 'kiss-fft-cpx s)
	  (dotimes (k h)
	    (dotimes (j w)
	      (with-foreign-slots 
		  ((r i) (mem-aref fin 'kiss-fft-cpx (+ (* k w) j)) 
		   kiss-fft-cpx)
		(setf r (coerce (realpart (aref image k j)) 'single-float)
		      i (coerce (imagpart (aref image k j)) 'single-float)))))
	  (with-foreign-object (fout 'kiss-fft-cpx s)
	    (kiss-fftnd kfg fin fout)
	    (dotimes (k h image)
	      (dotimes (j w)
		(with-foreign-slots 
		    ((r i) (mem-aref fout 'kiss-fft-cpx (+ (* k w) j))
		     kiss-fft-cpx)
		  (setf (aref image k j) (complex r i)))))))))))

(defun c-inverse-fft-2d! (image) (c-fft-2d! image t))

(defun c-fft-2d&mags! (image)
  (declare (optimize (speed 3) (safety 0)))
  (let ((h (array-dimension image 0))
	(w (array-dimension image 1))
	(s (array-total-size image))
	(lmage (make-array (array-dimensions image))))
    (with-foreign-object (fin 'kiss-fft-cpx s)
      (dotimes (k h)
	(dotimes (j w)
	  (with-foreign-slots 
	      ((r i) (mem-aref fin 'kiss-fft-cpx (+ (* k w) j)) 
	       kiss-fft-cpx)
	    (setf r (coerce (realpart (aref image k j)) 'single-float)
		  i (coerce (imagpart (aref image k j)) 'single-float)))))
      (with-foreign-object (fout 'kiss-fft-cpx s)
	(with-foreign-object (mag :float s)
	  (hr-mag-fft h w 0 fin fout mag)
	  (dotimes (k h (values lmage image))
	    (dotimes (j w)
	      (let ((l (+ (* k w) j)))
		(with-foreign-slots 
		    ((r i) (mem-aref fout 'kiss-fft-cpx l)
		     kiss-fft-cpx)
		  (setf (aref image k j) (complex r i)))
		(setf (aref lmage k j) (mem-aref mag :float l))))))))))
	    

(defun c-real-fft-2d! (image)
  (declare (optimize (speed 3) (safety 0)))
  (let ((h (array-dimension image 0))
	(w (array-dimension image 1))
	(s (array-total-size image)))
    (with-foreign-object (dims :int 2)
      (setf (mem-aref dims :int 0) h (mem-aref dims :int 1) w)
      (with-kiss (kfg (kiss-fftndr-alloc 
		       dims 2 0 (null-pointer) (null-pointer)))
	(with-foreign-object (fin :float s)
	  (dotimes (k h)
	    (dotimes (j w)
	      (setf (mem-aref fin :float (+ (* k w) j)) 
		    (coerce (aref image k j) 'single-float))))
	  (with-foreign-object (fout 'kiss-fft-cpx s)
	    (kiss-fftndr kfg fin fout)
	    (let ((h/2 (/ h 2)) (w/2 (/ w 2)))
	      (dotimes (k h)
		(dotimes (j (1+ w/2))
		  (with-foreign-slots 
		      ((r i) (mem-aref fout 'kiss-fft-cpx 
				       (+ (* k (1+ w/2)) j))
		       kiss-fft-cpx)
		    (setf (aref image k j) (complex r i)))))

	      (loop for j from 1 below w/2 do
		   (loop for k from 0 below h by h/2 do
			(setf (aref image k (+ j w/2))
			      (conjugate (aref image k (- w/2 j))))))

	      (loop for k from 1 below h/2 do
		   (loop for j from 1 below w/2 do
			(setf (aref image (- h/2 k) (+ j w/2))
			      (conjugate (aref image (+ k h/2) (- w/2 j))))))

	      (loop for k from 1 below h/2 do
		   (loop for j from 1 below w/2 do
			(setf (aref image (+ k h/2) (+ j w/2))
			      (conjugate (aref image (- h/2 k) (- w/2 j))))))
	      ))))))
  image)

(defun c-fourier-correlation! (f1 f2)
  (declare (optimize (speed 3) (safety 0)))
  (let ((h (array-dimension f1 0))
	(w (array-dimension f2 1))
	(s (array-total-size f1)))
    (with-foreign-object (ft1 'kiss-fft-cpx s)
      (dotimes (k h)
	(dotimes (j w)
	  (with-foreign-slots 
	      ((r i) (mem-aref ft1 'kiss-fft-cpx (+ (* k w) j)) 
	       kiss-fft-cpx)
	    (setf r (coerce (realpart (aref f2 k j)) 'single-float)
		  i (coerce (imagpart (aref f2 k j)) 'single-float)))))
      (with-foreign-object (ft2 'kiss-fft-cpx s)
	(dotimes (k h)
	  (dotimes (j w)
	    (with-foreign-slots 
		((r i) (mem-aref ft2 'kiss-fft-cpx (+ (* k w) j)) 
		 kiss-fft-cpx)
	      (setf r (coerce (realpart (aref f1 k j)) 'single-float)
		    i (coerce (imagpart (aref f1 k j)) 'single-float)))))
	(with-foreign-object (mag :float s)
	  (hr-mag-inv-phase h w ft1 ft2 mag)
	  (dotimes (k h f1)
	    (dotimes (j w)
	      (setf (aref f1 k j) (mem-aref mag :float (+ (* k w) j))))))))))

(defun c-rotate (image a)
  (declare (optimize (speed 3) (safety 0)))
  (let ((h (array-dimension image 0))
	(w (array-dimension image 1))
	(s (array-total-size image))
	(lmage (make-array (array-dimensions image))))
    (with-foreign-object (in :unsigned-char s)
      (dotimes (k h)
	(dotimes (j w)
	  (setf (mem-aref in :unsigned-char (+ (* k w) j)) (aref image k j))))
      (with-foreign-object (out :unsigned-char s)
	(hr-rotate h w (coerce a 'single-float) in out)
	(dotimes (k h lmage)
	  (dotimes (j w)
	    (setf (aref lmage k j) (mem-aref out :unsigned-char (+ (* k w) j)))))))))

(defun commit-adultery ()
  (setf 
   (fdefinition 'oci::fft-2d!) (function oci-rest::c-fft-2d!)
   (fdefinition 'oci::real-fft-2d!) (function oci-rest::c-real-fft-2d!)
   (fdefinition 'oci::inverse-fft-2d!) (function oci-rest::c-inverse-fft-2d!)
   (fdefinition 'oci::fourier-correlation!) (function oci-rest::c-fourier-correlation!)
   (fdefinition 'oci::fft-n-mags!) (function oci-rest::c-fft-2d&mags!);))
   (fdefinition 'oci::rotate-image) (function oci-rest::c-rotate))
   
(defun jpeg->gray (filename &optional (channel 1))
  (multiple-value-bind (a h w) (jpeg:decode-image filename)
    (let ((image (make-array (list h w) :element-type '(unsigned-byte 8)))
	  (stride (/ (array-total-size a) h w)))
      (dotimes (i h image)
	(dotimes (j w) 
	  (setf (aref image i j) 
		(aref a (+ (* i w stride) (* j stride) channel))))))))

(defparameter *figures* nil)
    
(defun slurp-numbers! (&optional (suffix ""))
  (let (f)
    (dotimes (i 10)
      (let ((mask (read-netpbm (format nil "~~/Desktop/~a~a.pnm" i suffix))))
	(push 
	 (list i (treshold-fourier! 0.1 (pad-image 128 mask 256)))
	 f)))
    (setf *figures* f))
  t)

(defun find-8s (image)
  (let ((ht (grayscale-image (hough+sobel-circle-transform 6 image))) circles)
    (dotimes (i (array-dimension image 0))
      (dotimes (j (array-dimension image 1))
	(when (< 150 (aref ht i j)) (push (cons i j) circles))))
;    (format t "~a~%" circles)
    (remove-if 
     #'null
     (mapcar
      #'(lambda (g)
	  (if (> (length g) 1)
	      (cons
	       (round (apply #'average (mapcar #'car g)))
	       (round (apply #'average (mapcar #'cdr g))))))
      (group #'(lambda (a b) (and (> 15 (abs (- (car a) (car b))))
				  (> 6 (abs (- (cdr a) (cdr b))))))
	     circles :every t)))))

(defun unite-centres (image lst)
  (let (eights eights-p)
    (mapcar
     #'(lambda (g)
;       (format t "GG ~a~%" g)
	 (let ((w (car (sort (copy-list g) #'> 
			     :key #'(lambda (e) (third (second e)))))))
	   (when (and (member (car w) '(3 5)) (member 8 g :key #'car))
	     (if (null eights-p) 
		 (setf eights-p t eights (find-8s image))) ;not very functional 
;	     (format t " ~a ~a ~a ~a~%" (car w) (mapcar #'car g) (cons (car (second w)) (cadr (second w))) eights)
	     (if (member 
		  (cons (car (second w)) (cadr (second w))) eights
		  :test #'(lambda (a b) 
			    (> *rec-num-hor* (euclidean-distance a b))))
		 (setf (car w) 8)))
	   (list (car w) (cons (car (second w)) (cadr (second w)))
		 (third (second w)))))
	 (group 
	  #'(lambda (a b)
	      (and (> (* 2 *rec-num-hor*) (abs (- (first a) (first b))))
		   (> (* 1.5 *rec-num-hor*) (abs (- (second a) (second b))))))
	  lst :every t :key #'second))))

(defun recognize-number (image)
  (apply
   #'values
   (sort
    (mapcar
     #'(lambda (n)
	 (let ((s 
		(format nil "~{~a~}" 
			(mapcar #'car (sort n #'oci::hor-point< :key #'cadr)))))
	   (list s (apply #'+ (mapcar #'third n)) (length s))))
     (group 
      #'(lambda (a b)
	  (and (< *rec-num-hor* (abs (- (cdr a) (cdr b))) *join-centres-pixels*)
	       (> *rec-num-ver* (abs (- (car a) (car b))))))
      (unite-centres 
       image
       (let (number)
	 (dolist (f *figures* number)
	   (let ((n (car f)) (m (cadr f)))
	     (dolist (p (fourier-search m image :t1 0.7 :mask-cooked-p t))
	       (if (> (third p) 400) (push (list n p) number)))))))
      :key #'second))
    #'> :key #'third)))

(defparameter *rec-num-hor* 10)
(defparameter *rec-num-ver* 7)
(defparameter *join-centres-angles* 5)
(defparameter *join-centres-pixels* 25)
(defparameter *number-rects* '((90 -120 210 120))) ;(60 -120 180 120)))

(defun switch-resolution (r)
  (case r
    (600
     (setf *join-centres-pixels* 25
	   *rec-num-hor* 10
	   *rec-num-ver* 7
	   *number-rects* '((90 -120 210 120)))) ;(60 -120 180 120))))
    (300
     (setf *join-centres-pixels* 18
	   *rec-num-hor* 7
	   *rec-num-ver* 5
	   *number-rects* '((66 -88 154 88)))))) ;(41 -82 123 82))))))

(defun get-centres (sc blackmaskname jpegname &optional debug-p)
  (apply
   #'append
   (let* ((mask (logarithm-filter (read-netpbm blackmaskname)))
	  (jpeg (negate-grayscale-image (jpeg->gray jpegname)))
	  (blacksmalljpeg (logarithm-filter (shrink-image sc jpeg)))
	  (cs (join-centres
	       *join-centres-angles* *join-centres-pixels*
	       (filter-centres 
		60 95 blacksmalljpeg 
		(fourier-mellin-search mask blacksmalljpeg)))))
     (if debug-p
	 (write-netpbm "~/im/0-0-0-0.pnm" (mark-image blacksmalljpeg cs)))
     (mapcar
      #'(lambda (c)
	  (interleave
	   #'(lambda (r a o)
	       (let ((k (restore-rect r a o jpeg)))
		 (list c k (if k (round (image-mean k))) a
		       (eq r (car *number-rects*)))))
	   nil
	   *number-rects*
	   (list (fourth c))
	   (list (cons (* sc (first c)) (* sc (second c))))))
      cs))))

(defun ignore-nulls (centres &optional debug-p)
  (remove-if 
   #'(lambda (e)
       (when (and debug-p (second e))
	 (let ((c (first e)))
	   (write-netpbm
	    (format 
	     nil "~~/im/all/~a-~a-~a-~a-~a-~a.pnm" 'x
	     (first c) (second c) (round (fourth e)) (third e) (fifth e))
	    (second e))))
       (or
	(null (second e))
	(not (< 10 (third e) 35))))
   centres))

(defun find-numbers (lst &optional debug-p)
  (mapcar
   #'(lambda (i)
       (let ((c (first i)) (n (recognize-number (second i))))
	 (if debug-p
	     (if (> (length (car n)) 6)
		 (write-netpbm
		  (format 
		   nil "~~/im/~a-~a-~a-~a-~a-~a.pnm" (car n)
		   (first c) (second c) (round (fourth i)) (third i) (fifth i))
		  (second i))
		 (write-netpbm
		  (format 
		   nil "~~/im/bad/~a-~a-~a-~a-~a-~a.pnm" (car n)
		   (first c) (second c) (round (fourth i)) (third i) (fifth i))
		  (second i))))
	 (append n (list (first c) (second c) (round (fourth i))))))
   lst))

(defun unite-numbers (n)
  (sort
;   (remove-duplicates
    (mapcar
     #'(lambda (g) 
	 (car 
	  (sort 
	   g 
	   #'(lambda (a b) 
	       (if (= (third a) (third b)) 
		   (> (cadr a) (cadr b))
		   (> (third a) (third b)))))))
     (oci:group 
      #'(lambda (a b) 
	  (and (> *join-centres-pixels* (abs (- (fifth a) (fifth b))))
	       (> *join-centres-pixels* (abs (- (fourth a) (fourth b))))))
;	       (> 0.35 (abs (- (sin (* pi (/ (sixth a) 180))) 
;			       (sin (* pi (/ (sixth b) 180))))))))
      n :every t))
 ;   :test #'string= :key #'car)
   #'string< :key #'car))

(defun fm-jpeg-search (sc blackmaskname jpegname &optional debug-p)
  (remove-if
    #'(lambda (x) (or (not (stringp x)) (> 7 (length x))))
    (find-numbers 
     (ignore-nulls (get-centres sc blackmaskname jpegname debug-p) debug-p)
     debug-p)
    :key #'car))