Add ranked-choice voting math and stuff

3 files changed, 115 insertions, 33 deletions

diff --git a/README.md b/README.md
index 8e78757..73a3356 100644
--- a/README.md
+++ b/README.md
@@ -6,7 +6,7 @@ Imagine this. You're a somehow still functioning student organization of compute
 Anyways, now we can vote online. It's cool, I guess? We have things such as:
  - **Server-side rendering**. That's right, this site (should) (mostly) work without JavaScript.
  - **Server Sent Events** for real-time vote results
- - **(Slightly less) limited voting options**. It's worse than Google Forms! (See To-Dos. All that's left now is ranked choice voting)
+ - **~~Limited~~ voting options**. It's now just as good as Google Forms, but a lot less safe! That's what you get when a bored college student does this in their free time
 
 ## Configuration
 You'll need to set up these values in your environment. Ask an RTP for OIDC credentials. A docker-compose file is provided for convenience. Otherwise, I trust you to figure it out!
@@ -22,5 +22,8 @@ VOTE_STATE=
 ## To-Dos
 - [x] Custom vote options
 - [x] Write-in votes
-- [ ] Ranked choice voting
+- [x] Ranked choice voting
 - [ ] Show options that got no votes
+- [ ] Write-In Fuzzy Matching
+- [ ] Allow results to be hidden until a vote is closed
+- [ ] Don't let the user fuck it up
diff --git a/database/poll.go b/database/poll.go
index 72d493d..1141d86 100644
--- a/database/poll.go
+++ b/database/poll.go
@@ -2,6 +2,7 @@ package database
 
 import (
 	"context"
+	"sort"
 	"time"
 
 	"go.mongodb.org/mongo-driver/bson"
@@ -21,11 +22,6 @@ type Poll struct {
 	AllowWriteIns    bool     `bson:"writeins"`
 }
 
-type Result struct {
-	Option string `bson:"_id"`
-	Count  int    `bson:"count"`
-}
-
 const POLL_TYPE_SIMPLE = "simple"
 const POLL_TYPE_RANKED = "ranked"
 
@@ -149,33 +145,111 @@ func (poll *Poll) GetResult() (map[string]int, error) {
 	defer cancel()
 
 	pollId, _ := primitive.ObjectIDFromHex(poll.Id)
-
-	cursor, err := Client.Database("vote").Collection("votes").Aggregate(ctx, mongo.Pipeline{
-		{{
-			"$match", bson.D{
-				{"pollId", pollId},
-			},
-		}},
-		{{
-			"$group", bson.D{
-				{"_id", "$option"},
-				{"count", bson.D{
-					{"$sum", 1},
-				}},
-			},
-		}},
-	})
-	if err != nil {
-		return nil, err
+	finalResult := make(map[string]int)
+
+	if poll.VoteType == POLL_TYPE_SIMPLE {
+		cursor, err := Client.Database("vote").Collection("votes").Aggregate(ctx, mongo.Pipeline{
+			{{
+				"$match", bson.D{
+					{"pollId", pollId},
+				},
+			}},
+			{{
+				"$group", bson.D{
+					{"_id", "$option"},
+					{"count", bson.D{
+						{"$sum", 1},
+					}},
+				},
+			}},
+		})
+		if err != nil {
+			return nil, err
+		}
+
+		var results []SimpleResult
+		cursor.All(ctx, &results)
+
+		for _, r := range results {
+			finalResult[r.Option] = r.Count
+		}
+		return finalResult, nil
+	} else if poll.VoteType == POLL_TYPE_RANKED {
+		cursor, err := Client.Database("vote").Collection("votes").Aggregate(ctx, mongo.Pipeline{
+			{{
+				"$match", bson.D{
+					{"pollId", pollId},
+				},
+			}},
+		})
+		if err != nil {
+			return nil, err
+		}
+
+		var votes []RankedVote
+		cursor.All(ctx, &votes)
+
+		voteCount := len(votes)
+
+		// ALRIGHT LETS GO INSTANT RUNOFF VOTE LOGIC GOES HERE
+		for {
+			// Create an empty result for counting at this iteration
+			results := make(map[string]int)
+			// Iterate through all cast votes
+			for _, vote := range votes {
+				// Create a list of the options in this vote and sort by preference
+				options := make([]string, 0, len(vote.Options))
+				for key := range vote.Options {
+					options = append(options, key)
+				}
+				sort.SliceStable(options, func(i, j int) bool {
+					return vote.Options[options[i]] < vote.Options[options[j]]
+				})
+
+				// Add a vote for the highest preference option
+				for _, option := range options {
+					// If that option has been eliminated, skip it and go to the next one
+					if containsKey(finalResult, option) {
+						continue
+					}
+					results[option] += 1
+					break
+				}
+			}
+
+			// Once we've gone through all votes, check if we have any options
+			// that have received more than half of the possible votes
+			for _, count := range results {
+				// If so, we're done
+				// This means we won't randomly mess with ties
+				if count*2 >= voteCount {
+					for k, c := range results {
+						finalResult[k] = c
+					}
+					return finalResult, nil
+				}
+			}
+			// If no option has won yet, find the option with the least votes and eliminate
+			// it, noting the number of votes it recieved at the time
+			options := make([]string, 0, len(finalResult))
+			for key := range results {
+				options = append(options, key)
+			}
+			sort.SliceStable(options, func(i, j int) bool {
+				return results[options[i]] < results[options[j]]
+			})
+
+			finalResult[options[len(options)-1]] = results[options[len(options)-1]]
+		}
 	}
+	return nil, nil
+}
 
-	var results []Result
-	cursor.All(ctx, &results)
-
-	result := make(map[string]int)
-	for _, r := range results {
-		result[r.Option] = r.Count
+func containsKey(arr map[string]int, val string) bool {
+	for key, _ := range arr {
+		if key == val {
+			return true
+		}
 	}
-
-	return result, nil
+	return false
 }
diff --git a/database/simple_vote.go b/database/simple_vote.go
index 0adea06..247d0a8 100644
--- a/database/simple_vote.go
+++ b/database/simple_vote.go
@@ -14,6 +14,11 @@ type SimpleVote struct {
 	Option string             `bson:"option"`
 }
 
+type SimpleResult struct {
+	Option string `bson:"_id"`
+	Count  int    `bson:"count"`
+}
+
 func CastSimpleVote(vote *SimpleVote) error {
 	ctx, cancel := context.WithTimeout(context.TODO(), 10*time.Second)
 	defer cancel()