#lang racket

(provide find-ferti-recipe)

(require math/array
         math/matrix
         "../models/nutrient.rkt"
         "../models/nutrient-measurement.rkt"
         "../models/nutrient-target.rkt"
         "../models/fertilizer-product.rkt")

(define (find-ferti-recipe)
  (define fertilizers (get-fertilizer-products))
  (define solution-array (find-nnls fertilizers))
  (for/list ([i (length fertilizers)])
    (cons (list-ref fertilizers i)
          (array-ref solution-array (vector i 0)))))

(define (find-nnls fertilizers)
  (define nutrients (get-nutrients))
  (define fertilizer-product-matrix (get-fertilizer-product-matrix nutrients fertilizers))
  (define deficits
    (->col-matrix
     (for/list ([n nutrients])
       (define latest-measurement (get-latest-nutrient-measurement-value n))
       (define latest-target (get-latest-nutrient-target-value n))
       (define deficit
         (cond
           [(or (false? latest-measurement)
                (zero? latest-measurement))
            latest-target]
           [(false? latest-target)
            0]
           [(and (number? latest-measurement)
                 (number? latest-target))
            (* 100
               (/ (- latest-target latest-measurement)
                  latest-measurement))]
           [else (error "either the target or measurement are not numbers")]))
       deficit)))
  (define error-threshold 10e-4)
  (lawson-hanson-1974 fertilizer-product-matrix deficits error-threshold))

;; Algorithm lifted from the Wikipedia article on NNLS
(define/contract (lawson-hanson-1974 A y ε)
  ;;;;;;;;;
  ;; Inputs
  ;;;;;;;;;

  (-> matrix?     ; Real-valued matrix A of dimension m × n
      col-matrix? ; Real-valued column matrix (vector) y of dimension m
      real?       ; Real-value ɛ, tolerance for the stopping criterion
      col-matrix?)              ; Real-valued solution column matrix x
  (define-values
    (m                                 ; Number of nutrients
     n)                                ; Number of fertilizer products
    (matrix-shape A))


  ;;;;;;;;;;;;;
  ;; Initialize
  ;;;;;;;;;;;;;

  ;; The passive set P is initially empty.
  (define P (mutable-set))
  ;; The active set R initially contains the indexes to the nutrients allowed to be...
  (define R (list->mutable-set (range n)))

  (define (colv-ref v i)
    (matrix-ref v i 0))

  ;; Gradient-like vector for residual error.
  (define (compute-w x)
    (matrix* (matrix-transpose A)
             (matrix- y (matrix* A x))))

  ;; max over j in R of w_j, returning (values max-val j*)
  (define (max-w-in-R w)
    (for/fold ([max-val -inf.0] [max-j #f])
              ([j (in-set R)])
      (define v (colv-ref w j))
      (if (> v max-val)
          (values v j)
          (values max-val max-j))))

  ;; Build full candidate vector s from current P:
  ;; s_P = (A_Pᵀ A_P)⁻¹ A_Pᵀ y, s_R = 0
  (define (make-s-from-P)
    (if (set-empty? P)
        (make-matrix n 1 0)
        (let* ([idxs (sort (set->list P) <)]
               [AP   (submatrix A (::) idxs)]
               [sP   (matrix* (matrix-inverse
                               (matrix* (matrix-transpose AP) AP))
                              (matrix-transpose AP)
                              y)])
          ;; map: column index j in P -> corresponding sP entry
          (define mapping
            (for/list ([j idxs] [k (in-naturals)])
              (cons j (colv-ref sP k))))
          (define (s-at i)
            (define p (assoc i mapping))
            (if p (cdr p) 0))
          (build-matrix n 1 (λ (i j) (s-at i))))))

  ;; The "first try" x represents no addition of any fertilizer.
  (define x (make-matrix n 1 0))


  ;;;;;;;;;;;;;
  ;; Outer loop
  ;;;;;;;;;;;;;

  (let outer-loop ()
    (define w (compute-w x))

    (cond
      ;; If no remaining candidates in R, we're done.
      [(set-empty? R)
       x]

      [else
       (define-values (max-val j*) (max-w-in-R w))

       ;; Stopping criterion: max(w_R) <= ε
       (cond
         [(or (not j*) (<= max-val ε))
          x]

         [else
          ;; Add j* to P, remove from R
          (set-remove! R j*)
          (set-add! P j*)

          ;; Inner loop: adjust until s_P > 0
          (let inner-loop ()
            (define s (make-s-from-P))

            ;; min(s_P)
            (define min-sP
              (if (set-empty? P)
                  +inf.0
                  (for/fold ([mn +inf.0])
                            ([j (in-set P)])
                    (min mn (colv-ref s j)))))

            (cond
              ;; If all s_P > 0 (or P empty), accept s as new x and go back to outer loop
              [(or (set-empty? P) (> min-sP 0))
               (set! x s)
               (outer-loop)]

              [else
               ;; Compute α = min_{i in P, s_i <= 0} x_i / (x_i - s_i)
               (define α
                 (for/fold ([a +inf.0])
                           ([j (in-set P)])
                   (define sj (colv-ref s j))
                   (if (<= sj 0)
                       (let* ([xj  (colv-ref x j)]
                              [den (- xj sj)])
                         (if (> den 0)
                             (min a (/ xj den))
                             a))
                       a)))

               (when (or (equal? α +inf.0) (<= α 0))
                 (error 'lawson-hanson-1974 "no valid α in inner loop"))

               ;; x ← x + α (s − x)
               (define new-x
                 (matrix+ x (matrix* α (matrix- s x))))

               ;; Move to R all indices j in P with x_j <= 0
               (define to-remove '())
               (for ([j (in-set P)])
                 (when (<= (colv-ref new-x j) 0)
                   (set! to-remove (cons j to-remove))))
               (for ([j to-remove])
                 (set-remove! P j)
                 (set-add! R j))

               (set! x new-x)
               (inner-loop)]))])])))

(define (get-fertilizer-product-matrix nutrients fertilizers)
  ;; Lines are nutrients, columns are fertilizers
  (build-matrix (length nutrients)
                (length fertilizers)
                (λ (i j)
                  (define selected-nutrient (list-ref nutrients i))
                  (define product (list-ref fertilizers j))
                  (define pair (assoc selected-nutrient
                                      (fertilizer-product-values product)))
                  (if pair (cdr pair) 0))))

(module+ test
  (require rackunit
           rackunit/text-ui
           "../db/conn.rkt"
           "../db/migrations.rkt")

  (define test-date "2025-01-01")

  (run-tests
   (test-suite
    "Nutrient measurement model"
    #:before (λ ()
               (connect! #:path 'memory)
               ;; (connect! #:path "test.sqlite3")
               (migrate-all!)

               (define nitrogen (create-nutrient! "Nitrogen" "N"))
               (define phosphorus (create-nutrient! "Phosphorus" "P"))

               (create-nutrient-measurement! test-date
                                             `((,nitrogen . 0)
                                               (,phosphorus . 0)))
               (create-nutrient-target! test-date
                                        `((,nitrogen . 100)
                                          (,phosphorus . 50)))

               (create-fertilizer-product! "King Nitrogen"
                                           `((,nitrogen . 100)))
               (create-fertilizer-product! "Phosphorescent Baboon"
                                           `((,nitrogen . 10)
                                             (,phosphorus . 100)))
               (create-fertilizer-product! "John's Phosphorus"
                                           `((,nitrogen . 3)
                                             (,phosphorus . 30))))
    #:after (λ ()
              (disconnect!))

    (test-case "Solve for NNLS"
      (displayln
       (format "Final solution for fertilizers is combination ~a"
               (find-ferti-recipe)))))))