;; clean-channel -- noise reduction

(provide 'snd-clean.scm)
(require snd-dsp.scm snd-generators.scm)

(define goertzel-channel
  (let ((documentation "(goertzel-channel freq beg dur snd (chn 0)) returns the amplitude of the 'freq' spectral component"))
    (lambda* (freq (beg 0) dur snd chn)
      (let* ((sr (srate snd))
	     (rfreq (/ (* 2.0 pi freq) sr))
	     (cs (* 2.0 (cos rfreq))))
	(let ((reader (make-sampler beg snd chn))
	      (len (- (if (number? dur) dur (- (framples snd chn) beg)) 2))
	      (flt (make-two-pole 1.0 (- cs) 1.0)))
	  (do ((i 0 (+ i 1)))
	      ((= i len))
	    (two-pole flt (next-sample reader)))
	  (let ((y1 (two-pole flt (next-sample reader)))
		(y0 (two-pole flt (next-sample reader))))
	    (magnitude (- y0 (* y1 (exp (complex 0.0 (- rfreq))))))))))))


(define* (check-freq freq snd chn)
  (let ((hum 0.0))
    (do ((i 0 (+ i 1))
	 (loc 0.0 (+ loc (round (/ (framples snd chn) 5)))))
	((= i 4))
      (set! hum (+ hum (goertzel-channel freq loc 2048 snd chn))))
    (/ hum 4.0)))


;;; -------- single sample clicks

(define* (remove-single-sample-clicks (jump 8) snd chn)
  (let* ((reader (make-sampler 0 snd chn))
	 (mx (make-moving-average 16)) ; local average of abs difference
	 (samp0 0.0)
	 (samp1 0.0)
	 (samp2 0.0)
	 (fixed 0)
	 (len (framples snd chn))
	 (block-size (min len 1048576)) ; do edits by blocks rather than sample-at-a-time (saves time, memory etc) 1048576=1024*1024
	 (block-ctr 0)
	 (block-beg 0)
	 (block (make-float-vector block-size))
	 (block-changed #f))
     (do ((ctr 0 (+ ctr 1)))
	 ((= ctr len))
       (set! samp0 samp1)
       (set! samp1 samp2)
       (set! samp2 (next-sample reader))
       (set! (block block-ctr) samp2)
       (let* ((df1 (abs (- samp1 samp0)))
	      (df2 (abs (- samp2 samp1)))
	      (df3 (abs (- samp2 samp0)))
	      (local-max (moving-average mx df1)))
	 (if (and (> df1 (* jump local-max))
		  (> df2 (* jump local-max))
		  (or (< df3 local-max)
		      (and (< df3 (* 2 local-max))
			   (< (* (- samp2 samp0)
				 (- samp1 samp0))
			      0.0))))
	     (begin
	       (set! samp1 (* .5 (+ samp0 samp2)))
	       (set! (block (- block-ctr 1)) samp1)
	       (set! block-changed #t)
	       (set! fixed (+ 1 fixed)))))
       (set! block-ctr (+ 1 block-ctr))
       (if (>= block-ctr block-size)
	   (begin
	     (if block-changed
		 (begin
		   (float-vector->channel block block-beg block-size snd chn)
		   (set! block-changed #f)))
	     (set! block-beg (+ block-beg (- block-size 1)))
	     (set! block-ctr 1)
	     (set! (block 0) samp2))))
     (if block-changed
	 (float-vector->channel block block-beg block-ctr snd chn))
    fixed))

(define (test-remove-single-clicks)
  (let ((test (new-sound "test.snd")))
    (let ((data (make-float-vector 1001)))
      (do ((i 2 (+ i 30))
	   (val .9 (- val .05)))
	  ((>= i 1000))
	(float-vector-set! data i val))
      (set! (data 1000) .001)
      (float-vector->channel data)
      (remove-single-sample-clicks)
      (let ((mx (maxamp))
	    (loc (maxamp-position)))
	(if (> mx 0.06)
	    (format () "~%;remove-single-sample-clicks 0: ~A (at ~D)" mx loc)))
      (revert-sound)
      (do ((i 0 (+ i 1))
	   (ang 0.0 (+ ang .01)))
	  ((= i 1000))
	(float-vector-set! data i (+ (float-vector-ref data i) (* .2 (sin ang)))))
      (float-vector->channel data)
      (remove-single-sample-clicks)
      (if (fneq (maxamp) .2)
	  (format () "~%;remove-single-sample-clicks sin max: ~A" (maxamp)))
      (let ((cur-data (channel->float-vector 0))
	    (diff 0.0))
	(do ((i 0 (+ i 1))
	     (ang 0.0 (+ ang .01)))
	    ((= i 1000))
	  (set! diff (max diff (abs (- ( cur-data i) (* .2 (sin ang)))))))
	(if (> diff .2)
	    (format () "~%;remove-single-sample-clicks sine max diff: ~A" diff))))
    (close-sound test)))

;;; -------- pops

(define* (smooth-float-vector data beg dur)
  (let* ((y0 (data beg))
	 (y1 (data (+ beg dur)))
	 (angle (if (> y1 y0) pi 0.0)) 
	 (off (* .5 (+ y0 y1))) 
	 (scale (* 0.5 (abs (- y1 y0))))
	 (incr (/ pi dur)))
    (do ((i 0 (+ i 1)))
	((= i dur))
      (set! (data (+ beg i)) (+ off (* scale (cos (+ angle (* i incr)))))))))

(define* (remove-pops (size 8) snd chn)
  (let* ((reader (make-sampler 0 snd chn))
	 (dly0 (make-delay (* 4 size)))
	 (dly1 (make-delay (* 5 size)))
	 (mx-ahead (make-moving-average (* 4 size))) ; local average of abs difference ahead of current window
	 (mx-behind (make-moving-average (* 4 size))) ; local average of abs difference behind current window
	 (mx (make-moving-average size)) ; local average of abs difference

	 (last-ahead-samp 0.0)
	 (last-dly0-samp 0.0)
	 (last-dly1-samp 0.0)

	 (last-case 0)
	 (fixed 0)
	 (len (framples snd chn))

	 (pad (* 8 size))
	 (block-size (min (+ len pad) 1048576)) ; 1048576=1024*1024
	 (block-ctr 0)
	 (block-beg 0)
	 (block (make-float-vector block-size))
	 (block-changed #f))
    
     (let ((check-val 0.0)
	   (check-start 0)
	   (checker 0)
	   (checked 0)
	   (checking #f)
	   (moving-start #t))
       (do ((ctr 0 (+ ctr 1)))
	   ((= ctr len))
	 (let* ((ahead-samp (next-sample reader))
		(diff-ahead (abs (- ahead-samp last-ahead-samp)))
		(avg-ahead (moving-average mx-ahead diff-ahead))
		(dly0-samp (delay dly0 ahead-samp))
		(cur-diff (abs (- dly0-samp last-dly0-samp)))
		(cur-avg (moving-average mx cur-diff))
		(dly1-samp (delay dly1 ahead-samp))
		(diff-behind (abs (- dly1-samp last-dly1-samp)))
		(avg-behind (moving-average mx-behind diff-behind)))
	   (set! last-ahead-samp ahead-samp)
	   (set! last-dly0-samp dly0-samp)
	   (set! last-dly1-samp dly1-samp)
	   (set! (block block-ctr) ahead-samp)
	   (if checking
	       (begin
		 (set! checked (+ checked 1))
		 (if (or (>= checked (* 2 size))
			 (< cur-avg check-val))
		     (begin
		       (set! fixed (+ fixed 1))
		       (set! checking #f)
		       (smooth-float-vector block (- check-start block-beg) (+ size checker))
		       (set! block-changed #t)))
		 (if moving-start
		     (begin
		       (set! moving-start (< cur-diff avg-behind))
		       (if moving-start
			   (set! check-start (+ check-start 1)))))
		 (if (not moving-start)
		     (set! checker (+ checker 1))))
	       (if (and (> (- ctr last-case) (* 2 size))
			(> cur-avg (* 4 avg-ahead))
			(> cur-avg (* 4 avg-behind)))
		   ;; possible pop
		   (begin
		     (set! check-start (max 0 (- ctr (* 5 size))))
		     (set! moving-start (< cur-diff avg-behind))
		     (if moving-start
			 (set! check-start (+ check-start 1)))
		     (set! checking #t)
		     (set! check-val cur-avg)
		     (set! checker 0)
		     (set! checked 0)
		     (set! last-case ctr))))
	   
	   (set! block-ctr (+ block-ctr 1))
	   (if (>= (+ block-ctr pad) block-size)
	       (begin
		 (if block-changed
		     (begin
		       (float-vector->channel block block-beg (- block-ctr pad) snd chn)
		       (set! block-changed #f)))
		 (set! block-beg (+ block-beg (- block-ctr pad)))
		 (do ((i 0 (+ i 1))
		      (j (- block-ctr pad) (+ j 1)))
		     ((= i pad))
		   (set! (block i) (block j)))
		 (set! block-ctr pad)))))
       
       (if block-changed
	   (float-vector->channel block block-beg block-ctr snd chn)))

    fixed))

(define (test-remove-pops)
  (new-sound "test.snd")
  (let ((data (make-float-vector 4001)))
    (do ((i 100 (+ i 200)))
	((>= i 3800))
      (let ((size (random 8)))
	(do ((k 0 (+ k 1)))
	    ((= k size))
	  (set! (data (+ i k)) (mus-random 1.0)))))
    (float-vector->channel data)
    (remove-pops)
    (let ((mx (maxamp)))
      (if (> mx .01)
	  (format () "~%;test remove-pops 0 case: ~A" mx)))
    (revert-sound)
    (do ((i 0 (+ i 1))
	 (ang 0.0 (+ ang .01)))
	((= i 4000))
      (set! (data i) (+ (data i)
			(* .2 (sin ang)))))
    (float-vector->channel data)
    (remove-pops)
    (let ((mx (maxamp)))
      (if (fneq mx .2)
	  (format () "~%;test remove-pops sine case: ~A" mx)))
    (close-sound)))
	

;;; -------- hum

(define (test-notch-hum)
  (let ((test (with-sound ("test.snd" :srate 22050)
		(let ((osc (make-oscil 60.0))
		      (e (make-env '(0 0 1 .5 9 .5 10 0) :length 44100)))
		   (do ((i 0 (+ i 1)))
		       ((= i 44100))
		     (outa i (* (env e) (oscil osc))))))))
    
    (notch-channel (list 60.0) #f #f #f #f #f #f #t 8)
    (let ((mx (maxamp)))
      (if (> mx .02)
	  (format () "~%;notch hum 0: ~A" mx)))
    (close-sound (find-sound test)))
  (let ((test (with-sound ("test.snd" :srate 22050)
		(let ((p (make-polywave 20.0 (list 2 1 3 1 4 1)))
		      (e (make-env '(0 0 1 .3 9 .3 10 0) :scaler 1/3 :length 44100)))
		   (do ((i 0 (+ i 1)))
		       ((= i 44100))
		     (outa i (* (env e) (polywave p))))))))
    
    (let ((v60 (goertzel 60.0))
	  (v40 (goertzel 40.0))
	  (v80 (goertzel 80.0)))
      (notch-channel (list 60.0) #f #f #f #f #f #f #t 8)
      (let ((e60 (goertzel 60.0))
	    (e40 (goertzel 40.0))
	    (e80 (goertzel 80.0)))
	(if (or (fneq (/ e60 v60) 0.0)
		(fneq (/ e40 v40) 1.0)
		(fneq (/ e80 v80) 1.0))
	    (format () "~%;notch 60: ~A ~A ~A -> ~A ~A ~A" v40 v60 v80 e40 e60 e80))))
    (close-sound (find-sound test)))

  (let ((test (with-sound ("test.snd" :srate 22050)
		(let ((p (make-polywave 5.0 (list 11 1 12 1 13 1)))
		      (e (make-env '(0 0 1 .3 9 .3 10 0) :scaler 1/3 :length 44100)))
		   (do ((i 0 (+ i 1)))
		       ((= i 44100))
		     (outa i (* (env e) (polywave p))))))))
    
    (let ((v60 (goertzel 60.0))
	  (v40 (goertzel 55.0))
	  (v80 (goertzel 65.0)))
      (notch-channel (list 60.0) #f #f #f #f #f #f #t 2)
      (let ((e60 (goertzel 60.0))
	    (e40 (goertzel 55.0))
	    (e80 (goertzel 65.0)))
	(if (or (> (/ e60 v60) 0.01)
		(< (/ e40 v40) 0.99)
		(< (/ e80 v80) 0.99))
	    (format () "~%;notch 60 tight: ~A ~A ~A -> ~A ~A ~A" v40 v60 v80 e40 e60 e80))))
    (close-sound (find-sound test))))


;;; -------- DC

(define (test-remove-DC)
  (let ((test (new-sound "test.snd"))
	(data (make-float-vector 4001)))
    (do ((i 0 (+ i 1))
	 (ang 0.0 (+ ang .01)))
	((= i 4000))
      (float-vector-set! data i (+ .1 (mus-random 0.1) (* .2 (sin ang)))))
    (float-vector->channel data)
    (let ((dc (goertzel 0.0))
	  (sig (goertzel 35.0)))
      (let ((dcflt (make-filter 2 (float-vector 1 -1) (float-vector 0 -0.99))))
	(map-channel (lambda (y) (filter dcflt y)))
	(let ((ndc (goertzel 0.0))
	      (nsig (goertzel 35.0)))
	  (if (or (> (/ ndc dc) .1)
		  (< (/ nsig sig) .4))
	      (format () "~%;remove-DC: ~A -> ~A (~A), ~A -> ~A (~A)" dc ndc (/ ndc dc) sig nsig (/ nsig sig))))))
    (close-sound test)))


(define* (tvf-channel snd chn)
  (let* ((size (framples snd chn))
	 (avg-data (make-float-vector size))
	 (ctr 0)
	 (mx (maxamp snd chn))
	 (avg-size 8192)
	 (rd0 (make-sampler 0 snd chn))
	 (xhat 0.0)
	 (frm (make-formant :radius (- 1.0 (/ 500.0 (srate snd))) :frequency 600))
	 (del (make-moving-sum avg-size))
	 (K 0.0)

	 (maxg 0.0)
	 (ming 1000.0)
	 (maxK 0.0)
	 (minK 1000.0)
	 )

    (do ((i 0 (+ i 1)))
	((= i avg-size))
      (moving-sum del (formant frm (rd0))))

    (map-channel
     (lambda (datum)
       (let ((xhatminus xhat)
	     (avg (moving-sum del (formant frm (rd0)))))

	 (set! K (min 1.0 (+ .1 (/ avg 100.0))))
;	 (set! K .5)

	 (set! (avg-data ctr) K)
	 (set! ctr (+ ctr 1))

	 (set! maxg (max maxg avg))
	 (set! ming (min ming avg))
	 (set! maxK (max maxK K))
	 (set! minK (min minK K))

	 (set! xhat (+ xhatminus
		       (* K (- datum xhatminus))))
	 xhatminus))
     0 size snd chn)

    (let ((mx1 (maxamp snd chn)))
      (scale-channel (/ mx mx1) snd chn))

;    (format () ";K ~A to ~A, avg ~A to ~A" minK maxK ming maxg)
;    avg-data
    ))



(define* (clean-channel snd chn)

  ;; look for obvious simple clicks
  (let ((clicks (as-one-edit (lambda () (remove-single-sample-clicks 8 snd chn)))))
    (if (> clicks 0)
	(format () "~%; fixed ~D single sample clicks" clicks)
	(format () "~%; no single-sample clicks found")))

  ;; look for obvious clipping and try to reconstruct
  (let ((mx (maxamp snd chn)))
    (if (>= mx 1.0)
	(let ((clips (unclip-channel snd chn)))
	  (if (eq? clips 'no-clips)
	      (format () "~%; no clipped portions found")
	      (format () "~%; reconstructed ~D clipped portions" (list-ref clips 3))))
	(format () "~%; no obvious clipping (max amp: ~A)" mx)))

  ;; look for pops
  (let ((total-pops 0))
    (call-with-exit
     (lambda (quit)
       (for-each
	(lambda (size)
	  (let ((pops (as-one-edit (lambda () (remove-pops size snd chn)))))
	    (set! total-pops (+ total-pops pops))
	    (if (> pops 0)
		(format () "~%; fixed ~D ~D-sample ~A" pops size (if (= pops 1) "pop" "pops"))
		(quit))))
	(list 4 8 16 32))))
    (if (= total-pops 0)
	(format () "~%; no pops found")))

  ;; look for hum
  (let* ((hum60 (check-freq 60.0 snd chn))
	 (hum55 (check-freq 55.0 snd chn))
	 (hum (max hum60 hum55)))
    (if (> hum 30.0)
	(let ((humf (if (> hum60 hum55) 60.0 55.0)))
	  (notch-channel (list humf) 4096 0 (framples snd chn) snd chn #f #t 4)
	  (format () "~%; notch out ~D cycle hum: ~A -> ~A" (floor humf) hum (check-freq humf snd chn)))))

  ;; look for DC
  (let ((dc (check-freq 0.0 snd chn)))
    (if (> dc 30.0)
	(let ((dcflt (make-filter 2 (float-vector 1 -1) (float-vector 0 -0.99))))
	  (map-channel (lambda (y) (filter dcflt y)) 0 (framples snd chn) snd chn)
	  (format () "~%; block DC: ~A -> ~A" dc (check-freq 0.0 snd chn)))))

  ;; time-varying low-pass filter
  (tvf-channel snd chn)
  )


(define* (clean-sound snd)
  (let ((index (or snd (selected-sound) (car (sounds)))))
    (if (not (sound? index))
	(error 'no-such-sound (list "clean-sound" snd))
	(let ((chns (channels index)))
	  (do ((chn 0 (+ chn 1)))
	      ((= chn chns))
	    (clean-channel index chn))))))
