#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 (solve-nnls fertilizers))
  (for/list ([fertilizer (in-list fertilizers)]
             [quantity (in-array solution-array)])
    (cons fertilizer quantity)))

(define (solve-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
                        [(false? latest-target) 0]
                        [(or (false? latest-measurement) (zero? latest-measurement)) latest-target]
                        [(and (number? latest-measurement) (number? latest-target))
                         (* 100 (/ (- latest-target latest-measurement) latest-measurement))]))
                    deficit)))
  (define error-threshold 10e-4)
  (lawson-hanson-1974 fertilizer-product-matrix deficits error-threshold))

;; Algorithm lifted from the Wikipedia article on NNLS
(define (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.
  (define (max-w-in-R w)
    (if (set-empty? R)
        (values -inf.0 #f)
        (let ([max-j (argmax (λ (j) (colv-ref w j)) (set->list R))])
          (values (colv-ref w max-j) 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))
            (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-scale (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))
                  (hash-ref (fertilizer-product-values product) selected-nutrient 0))))

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

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

  (run-tests (test-suite "NNLS"
               (test-case "Build fertilizer product matrix"
                 (connect! #:path 'memory)
                 (migrate-all!)

                 (define n1 (create-nutrient! "N1" "" "N1"))
                 (define n2 (create-nutrient! "N2" "" "N2"))
                 (define nutrients (list n1 n2))

                 (define f1 (create-fertilizer-product! "F1" "F1" (hash n1 10 n2 20)))
                 (define f2 (create-fertilizer-product! "F2" "F2" (hash n1 30 n2 5)))
                 (define fertilizers (list f1 f2))

                 (define matrix (get-fertilizer-product-matrix nutrients fertilizers))

                 (check-= (matrix-ref matrix 0 0) 10 0 "N1 in F1")
                 (check-= (matrix-ref matrix 0 1) 30 0 "N1 in F2")
                 (check-= (matrix-ref matrix 1 0) 20 0 "N2 in F1")
                 (check-= (matrix-ref matrix 1 1) 5 0 "N2 in F2")

                 (disconnect!))

               (test-case "Single nutrient, single fertilizer"
                 (define A (matrix [[2]]))
                 (define y (col-matrix [10]))
                 (define ε 1e-6)

                 (define result (lawson-hanson-1974 A y ε))

                 (check-= (matrix-ref result 0 0) 5.0 ε "Should give x = 5 since 2*5 = 10"))

               (test-case "Two variables, known solution"
                 (define A (matrix [[1 0] [0 1]]))
                 (define y (col-matrix [3 4]))
                 (define ε 1e-6)

                 (define result (lawson-hanson-1974 A y ε))

                 (check-= (matrix-ref result 0 0) 3.0 ε "x1 should be 3")
                 (check-= (matrix-ref result 1 0) 4.0 ε "x2 should be 4"))

               (test-case "Overdetermined system"
                 (define A (matrix [[1 1] [2 1] [1 2]]))
                 (define y (col-matrix [3 5 5]))
                 (define ε 1e-4)

                 (define result (lawson-hanson-1974 A y ε))

                 ;; Solution should be approximately [1.636, 1.636] (least squares fit)
                 (check-= (matrix-ref result 0 0) 1.636 0.01 "x1 approximately 1.636")
                 (check-= (matrix-ref result 1 0) 1.636 0.01 "x2 approximately 1.636"))

               (test-case "Non-negativity enforcement"
                 (define A (matrix [[1 -1] [1 1]]))
                 (define y (col-matrix [1 3]))
                 (define ε 1e-6)

                 (define result (lawson-hanson-1974 A y ε))

                 ;; All results should be non-negative
                 (check-true (>= (matrix-ref result 0 0) 0) "x1 should be non-negative")
                 (check-true (>= (matrix-ref result 1 0) 0) "x2 should be non-negative"))

               (test-case "Zero target"
                 (define A (matrix [[1 2] [3 4]]))
                 (define y (col-matrix [0 0]))
                 (define ε 1e-6)

                 (define result (lawson-hanson-1974 A y ε))

                 (check-= (matrix-ref result 0 0) 0.0 ε "x1 should be 0")
                 (check-= (matrix-ref result 1 0) 0.0 ε "x2 should be 0"))

               (test-case "Ferti recipe"
                 (connect! #:path 'memory)
                 (migrate-all!)

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

                 (create-nutrient-measurement! test-date (hash nitrogen 0 phosphorus 0))
                 (create-nutrient-target! test-date (hash nitrogen 100 phosphorus 50))

                 (create-fertilizer-product! "Nitrogen" "King Nitrogen" (hash nitrogen 100))
                 (create-fertilizer-product! "Phosphorus"
                                             "Phosphorescent Baboon"
                                             (hash nitrogen 10 phosphorus 100))
                 (create-fertilizer-product! "Diluted phosphorus"
                                             "John's Phosphorus"
                                             (hash nitrogen 3 phosphorus 30))

                 (define recipe (find-ferti-recipe))

                 (check-equal? (length recipe) 3 "Should have 3 fertilizer products")

                 (for ([pair recipe])
                   (check-true (>= (cdr pair) 0) "Fertilizer quantity should be non-negative"))

                 (disconnect!))

               ;; Test deficit calculation edge cases
               (test-case "Deficit calculation with missing data"
                 (connect! #:path 'memory)
                 (migrate-all!)

                 (define n (create-nutrient! "TestNutrient" "" "TN"))

                 ;; No measurement, no target
                 (check-false (get-latest-nutrient-measurement-value n))
                 (check-false (get-latest-nutrient-target-value n))

                 ;; Add only target
                 (create-nutrient-target! test-date (hash n 100))
                 (check-= (get-latest-nutrient-target-value n) 100 0)

                 ;; Add measurement
                 (create-nutrient-measurement! test-date (hash n 50))
                 (define measured (get-latest-nutrient-measurement-value n))
                 (define targeted (get-latest-nutrient-target-value n))

                 ;; Deficit should be 100% (from 50 to 100)
                 (define deficit (* 100 (/ (- targeted measured) measured)))
                 (check-= deficit 100.0 0.01 "Deficit should be 100%")

                 (disconnect!))

               ;; Test recipe with realistic constraints
               (test-case "Recipe calculation with real-world scenario"
                 (connect! #:path 'memory)
                 (migrate-all!)

                 (define n (create-nutrient! "N" "" "N"))
                 (define p (create-nutrient! "P" "" "P"))
                 (define k (create-nutrient! "K" "" "K"))

                 ;; Current levels
                 (create-nutrient-measurement! test-date (hash n 50 p 10 k 100))

                 ;; Target levels
                 (create-nutrient-target! test-date (hash n 150 p 30 k 200))

                 ;; Fertilizers with different NPK ratios
                 (create-fertilizer-product! "" "Balanced" (hash n 100 p 100 k 100))
                 (create-fertilizer-product! "Nitrogen blend" "High-N" (hash n 200 p 50 k 50))
                 (create-fertilizer-product! "Phosphorus blend" "High-P" (hash n 50 p 200 k 50))
                 (create-fertilizer-product! "Potassium blend" "High-K" (hash n 50 p 50 k 200))

                 (define recipe (find-ferti-recipe))

                 (check-equal? (length recipe) 4 "Should have 4 fertilizer options")

                 ;; Verify solution is non-negative
                 (for ([pair recipe])
                   (check-true (>= (cdr pair) 0)
                               (format "~a quantity must be non-negative"
                                       (fertilizer-name (car pair)))))

                 (disconnect!)))))