#lang racket/base

(require racket/contract
         racket/list
         racket/vector
         "../base.rkt"
         "../data.rkt"
         "bits.rkt"
         "bmp.rkt"
         "constants.rkt"
         "galois.rkt"
         "layout.rkt"
         "version.rkt")

(provide/contract
 [qr-encode-data (-> qr-spec? bytes?)]
 [qr-encode/bmp (-> qr-spec? bmp?)])

(define (encode-data/pad data version edc-level)
  (let* ([max-words (qr-data-word-count version edc-level)]
         [max-bits (* max-words 8)]
         [bits (qr-data->bits data version)]
         [bit-count (length bits)])
    (cond
      [(> bit-count max-bits) (error "too many data bits" bit-count)]
      [(>= bit-count (- max-bits 8)) (bits->bytes bits)]
      [else
       (let* ([words (bits->bytes
                      (append bits '(#f #f #f #f)))]
              [word-count (bytes-length words)]
              [pad-count (- max-words word-count)])
         (bytes-append words
                       (list->bytes
                        (build-list pad-count
                                    (lambda (i)
                                      (if (even? i)
                                          #b11101100
                                          #b00010001))))))])))

(define (encode-data/ecd-blocks data version edc-level)
  (let-values ([(data-words) (encode-data/pad data version edc-level)]
               [(data-counts ecd-counts) (qr-data/edc-block-sizes version edc-level)])
    (let* ([data-blocks (let loop ([counts data-counts]
                                   [start 0])
                          (if (null? counts)
                              '()
                              (let* ([end (+ start (car counts))]
                                     [block (subbytes data-words start end)])
                                (cons block (loop (cdr counts) end)))))]
           [ecd-blocks (map make-ecd-block data-blocks ecd-counts)])
      (values data-blocks ecd-blocks))))

(define (interleave . lsts)
  (let loop ([lsts lsts])
    (let ([lsts (filter (lambda (l)
                          (not (null? l)))
                        lsts)])
      (if (null? lsts)
          '()
          (append (map car lsts)
                  (loop (map cdr lsts)))))))

(define (make-ecd-block data len)
  (let ([gs (make-rs-generator len)])
    (do ([ds (bytes->list data) (cdr ds)]
         [bs (make-vector len 0)
             (let ([x (gf+ (vector-ref bs (sub1 len))
                           (car ds))])
               (vector-map gf+
                 (vector-append #(0) (vector-drop-right bs 1))
                 (vector-map (lambda (g)
                               (gf* g x))
                             gs)))])
      ((null? ds) (list->bytes (reverse (vector->list bs)))))))

(define (qr-encode-data spec)
  (let-values
      ([(data-blocks edc-blocks)
        (encode-data/ecd-blocks (qr-spec-data spec)
                                (qr-spec-version spec)
                                (qr-spec-edc-level spec))])
    (bytes-append
     (list->bytes 
      (apply interleave
             (map bytes->list data-blocks)))
     (list->bytes
      (apply interleave
             (map bytes->list edc-blocks))))))

(define (apply-data-masks sym mask)
  (for/list ([n (in-range 8)])
    (qr-apply-data-mask n sym mask)))

(define (score-masked sym)
  (define (count-runs lst [is-equal? eqv?])
    (if (null? lst)
        '()
        (let next ([x (car lst)]
                   [n 1]
                   [lst (cdr lst)])
          (cond
            [(null? lst) (list n)]
            [(is-equal? x (car lst))
             (next x (add1 n) (cdr lst))]
            [else
             (cons n
                   (next (car lst) 1 (cdr lst)))]))))
  (define (all-eq? x . xs)
    (andmap (lambda (y)
              (eq? x y))
            xs))
  (define (count-locator-patterns bits)
    (let ([bits (append '(#f #f #f #f) bits '(#f #f #f #f))])
      (for/sum ([i (in-range (- (length bits) 11))])
        (let ([bits (take (drop bits i) 11)])
          (if (or (equal? bits '(#f #f #f #f #t #f #t #t #t #f #t))
                  (equal? bits '(#t #f #t #t #t #f #t #f #f #f #f)))
              1
              0)))))
  (let ([size (bmp-width sym)])
    (let ([runs (for/sum ([i (in-range size)])
                  (apply +
                         (map (lambda (len)
                                (if (>= len 5)
                                    (+ 3 (- len 5))
                                    0))
                              (append
                               (count-runs
                                (for/list ([j (in-range size)])
                                  (get-px sym i j)))
                               (count-runs
                                (for/list ([j (in-range size)])
                                  (get-px sym j i)))))))]
          [blocks (for*/sum ([y (in-range (sub1 size))]
                             [x (in-range (sub1 size))])
                    (if (all-eq? (get-px sym x y)
                                 (get-px sym (add1 x) y)
                                 (get-px sym x (add1 y))
                                 (get-px sym (add1 x) (add1 y)))
                        3
                        0))]
          [locators (for/sum ([i (in-range size)])
                      (* (+ (count-locator-patterns
                             (for/list ([j (in-range size)])
                               (get-px sym i j)))
                            (count-locator-patterns
                             (for/list ([j (in-range size)])
                               (get-px sym j i))))
                         40))]
          [proportion (let ([total-pixels (* size size)]
                            [set-pixels (for*/sum ([y (in-range size)]
                                                   [x (in-range size)]
                                                   #:when (get-px sym x y))
                                          1)])
                        (* 10 (quotient (abs (- 50 (quotient (* 100 set-pixels) total-pixels)))
                                        5)))])
      (+ runs blocks locators proportion))))

(define (write-format-info! version edc-level mask bmp)
  (let* ([bits (append (integer->bits (vector-memv edc-level #(M L H Q)) 2)
                       (integer->bits mask 3))]
         [bch (gf-mod (arithmetic-shift (bits->integer bits) 10)
                      #b10100110111)]
         [bits (append bits (integer->bits bch 10))]
         [bits (map (lambda (b1 b2)
                      (not (eq? b1 b2)))
                    bits
                    (integer->bits #b101010000010010 15))])
    (let ([side (qr-size version)])
      (for ([i (in-naturals)]
            [b (in-list (reverse bits))]
            #:when b)
        (cond
          [(<= i 5) (set-px! bmp 8 i)
                    (set-px! bmp (- side 1 i) 8)]
          [(<= i 7) (set-px! bmp 8 (+ i 1))
                    (set-px! bmp (- side 1 i) 8)]
          [(<= i 8) (set-px! bmp 7 8)
                    (set-px! bmp 8 (- side 6))]
          [else (set-px! bmp (- 14 i) 8)
                (set-px! bmp 8 (+ (- side 15) i))]))
      bmp)))

(define (write-version-info! version bmp)
  (let* ([bch (vector-ref qr-version-edc-bits-table (- version 7))]
         [bits (append (integer->bits version 6)
                       (integer->bits bch 12))])
    (let ([size (bmp-width bmp)])
      (for ([b (in-list (reverse bits))]
            [i (in-naturals)]
            #:when b)
        (let-values ([(q r) (quotient/remainder i 3)])
          (set-px! bmp (+ (- size 11) r) q)
          (set-px! bmp q (+ (- size 11) r)))))))

(define (qr-encode/bmp spec)
  (let* ([version (qr-spec-version spec)]
         [bits (qr-encode-data spec)]
         [sym (qr-template version)]
         [mask (qr-data-mask version)])
    (qr-layout-data! sym mask bits)
    (let ([masked (apply-data-masks sym mask)])
      (for ([sym (in-list masked)]
            [n (in-naturals)])
        (when (qr-has-format-information? version)
          (write-format-info! version (qr-spec-edc-level spec) n sym))
        (when (qr-has-version-information? version)
          (write-version-info! version sym)))
      (car
       (sort masked
             <
             #:key score-masked
             #:cache-keys? #t)))))