(define-module (deliberate irv)
  #:use-module (ice-9 receive)
  #:use-module (srfi srfi-1)
  #:use-module (deliberate preference)
  #:use-module (deliberate utils)
  #:export (irv-counts->first-counts irv-first-counts->scores social-choice-irv social-preference-irv)
  #:export (delete-lowest irv-round lowest-count)
  )

;; Find the social prefrence ranking using instant-runoff voting..
;;
;; Parameters:
;;   counts: List of cells of the form (`count` . `ranking`)
;;
;; Returns:
;;   A list of lists. The outer list represents the social ranking from
;;   most to least preferred. The inner lists represent sets of alternatives
;;   tied for the same place.
;;
(define (social-preference-irv counts)
  (reverse (social-preference-irv-rev counts)))

;; Recursive helper function generates ordering from least to most preferred
;;
(define (social-preference-irv-rev counts)
  (if (nil? counts)
      '()
      (let* ((delete-counts (irv-round counts))
	     (delete-alts (car delete-counts))
	     (next-counts (cdr delete-counts)))
	(cons delete-alts (social-preference-irv-rev next-counts))
	)))

;; Find the winner of a vote using instant-runoff voting.
;;
;; Parameters:
;;   counts: List of cells of the form (`count` . `ranking`)
;;
;; Returns:
;;   A list contianing alternatives tied for first place.
;;
(define (social-choice-irv counts)
  (car (social-preference-irv counts)))


;; INTERNAL UTILITIES ;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;

;; Extract first-place alternatives from preference counts
;;
;; Parameters:
;;   counts: List of cells of the form (`count` . `ranking`)
;;
;; Returns:
;;   List of cells of the form (`count` . `alternative`), where
;;   `alternative` is the highest ranked alternative in the corresponding
;;   input pair..
;;
(define (irv-counts->first-counts counts)
  (map
   (lambda (count-ranked)
     (cons (pref-count count-ranked)
	   (car (pref-ranked count-ranked))))
   counts))

;; Calculate total first place counts for each alternative
;;
;; Parameters:
;;  first-counts: Output of `irv-counts->first-counts`
;;
;; Returns:
;;   A list of pairs of the form (`total` . `alternative`) where `total`
;;   is the total count for `alternative` over all input pairs.
;;
(define (irv-first-counts->scores first-counts)
  (let ((groups (group-by cdr first-counts)))
    (map (lambda (group)
	   (cons
	    (reduce + 0 (map car (cdr group)))
	    (car group)))
	 groups)))

;; Finds the lowest total score.
;;
;; Parameters:
;;   scores: Output of `irv-first-counts->scores`
;;
;; Returns:
;;   The lowest count
;;
(define (lowest-count scores)
  (apply min (map car scores)))

;; Removes the alternative with the lowest total first-place count.
;;
;; Parameters:
;;   scores: List of pairs of the form (`total` . `alternative`) representing
;;     the total number of first-ranked votes for `alternative`.
;;
;; Returns:
;;   A pair of lists: (`to-keep` . `to-delete`)
;;     to-keep: A list of alternatives to keep into the next voting round.
;;     to-delete: A list of alternatives to remove from voting.
;;
(define (delete-lowest scores)
  (let ((lowest (lowest-count scores)))
    (receive (to-delete to-keep)
	(partition
	 (lambda (elt) (= lowest (car elt)))
	 scores)
      (cons
       (map cdr to-keep)
       (map cdr to-delete))
      )))

;; Calculates one round of instant-runoff voting.
;;
;; Paramters:
;;   counts: List of pairs of the form (`count` . `ranking`)
;;
;; Returns:
;;   A list of the same format as `counts` with the lowest ranked
;;   alternative(s) removed. Any pairs with no alternatives remaining are
;;   removed entirely.
;;
(define (irv-round counts)
  (let* ((scores
	 (irv-first-counts->scores
	  (irv-counts->first-counts counts)))
	 (keep-delete (delete-lowest scores))
	 (keep-alts (car keep-delete))
	 (delete-alts (cdr keep-delete)))
    (cons
     delete-alts
     (filter-empty
      (map (lambda (count-elt)
	     (cons (car count-elt)
		   (filter (lambda (alt)
			     (memv alt keep-alts))
			   (cdr count-elt))
		   ))
	   counts)))
    ))

;; Remove pairs from `counts` having no remaining alternatives.
;;
(define (filter-empty counts)
  (filter (lambda (count-elt)
	    (not (nil? (cdr count-elt))))
	  counts))