Add an option to print out elected people; add aliases.

main.go

@@ -1,6 +1,7 @@
 package main
 
 import (
+	"flag"
 	"fmt"
 	"io"
 	"net/http"
@@ -11,6 +12,18 @@ import (
 	"google.golang.org/protobuf/proto"
 )
 
+var (
+	printNames = flag.Bool("print_names", false, "Whether to print the names of candidates who got in.")
+)
+
+var aliases = map[string]string{
+	"KOMITET WYBORCZY PRAWO I SPRAWIEDLIWOŚĆ":                                                            "PiS",
+	"KOALICYJNY KOMITET WYBORCZY KOALICJA OBYWATELSKA PO .N IPL ZIELONI":                                 "KO",
+	"KOALICYJNY KOMITET WYBORCZY TRZECIA DROGA POLSKA 2050 SZYMONA HOŁOWNI - POLSKIE STRONNICTWO LUDOWE": "TD/PSL",
+	"KOMITET WYBORCZY NOWA LEWICA":                                                                       "Lewica",
+	"KOMITET WYBORCZY KONFEDERACJA WOLNOŚĆ I NIEPODLEGŁOŚĆ":                                              "Konfederacja",
+}
+
 var mandatesPerRegion = []int{12, 8, 14, 12, 13, 15, 12, 12, 10, 9, 12, 8, 14, 10, 9, 10, 9, 12, 20, 12, 12, 11, 15, 14, 12, 14, 9, 7, 9, 9, 12, 9, 16, 8, 10, 12, 9, 9, 10, 8, 12}
 
 type mandateCount struct {
@@ -19,6 +32,8 @@ type mandateCount struct {
 }
 
 func main() {
+	flag.Parse()
+
 	resp, err := http.Get(`https://static2.wybory.gov.pl/sejmsenat2023/data/komitety.blob`)
 	if err != nil {
 		fmt.Fprintf(os.Stderr, "%v\n", err)
@@ -44,6 +59,9 @@ func main() {
 		for _, cs := range p.Candidates {
 			for _, c := range cs.Candidate {
 				membership[c.Id] = p.Meta.Name
+				if aliases[p.Meta.Name] != "" {
+					membership[c.Id] = aliases[p.Meta.Name]
+				}
 				name[c.Id] = c.Name
 			}
 		}
@@ -52,7 +70,7 @@ func main() {
 	totals := map[string]int{}
 	for i, numMandates := range mandatesPerRegion {
 		fmt.Printf("Region %d:\n", i+1)
-		regionMandates, err := getRegionMandates(i+1, numMandates, membership, name)
+		regionMandates, candidates, err := getRegionMandates(i+1, numMandates, membership, name)
 		if err != nil {
 			fmt.Fprintf(os.Stderr, "%v\n", err)
 			continue
@@ -64,6 +82,11 @@ func main() {
 		sort.Slice(sorted, func(i, j int) bool { return sorted[i].mandates > sorted[j].mandates })
 		for _, x := range sorted {
 			fmt.Printf("\t%s: %d\n", x.party, x.mandates)
+			if *printNames {
+				for _, c := range candidates[x.party] {
+					fmt.Printf("\t\t%s\n", c)
+				}
+			}
 		}
 		for p, m := range regionMandates {
 			totals[p] += m
@@ -82,34 +105,42 @@ func main() {
 }
 
 type weightedVote struct {
-	party  string
+	entity string
 	weight int
 }
 
-func getRegionMandates(id int, numMandates int, membership map[int32]string, name map[int32]string) (map[string]int, error) {
+func getRegionMandates(id int, numMandates int, membership map[int32]string, name map[int32]string) (map[string]int, map[string][]string, error) {
 	resp, err := http.Get(fmt.Sprintf(`https://static.wybory.gov.pl/sejmsenat2023/data/obkw/okr_sejm/%d.blob`, id))
 	if err != nil {
-		return nil, err
+		return nil, nil, err
 	}
 	defer resp.Body.Close()
 	body, err := io.ReadAll(resp.Body)
 	if err != nil {
-		return nil, err
+		return nil, nil, err
 	}
 
 	r := &pb.Region{}
 	opts := proto.UnmarshalOptions{DiscardUnknown: true}
 	if err := opts.Unmarshal(body, r); err != nil {
-		return nil, err
+		return nil, nil, err
 	}
 
 	votes := map[string]int32{}
+	votesPerParty := map[string][]weightedVote{}
 
 	for _, p := range r.WhoKnows[0].Person {
-		votes[membership[p.Id]] += p.VoteCount.Cities
-		votes[membership[p.Id]] += p.VoteCount.Villages
-		votes[membership[p.Id]] += p.VoteCount.Ships
-		votes[membership[p.Id]] += p.VoteCount.Foreign
+		v := int32(0)
+		v += p.VoteCount.Cities
+		v += p.VoteCount.Villages
+		v += p.VoteCount.Ships
+		v += p.VoteCount.Foreign
+
+		votes[membership[p.Id]] += v
+		votesPerParty[membership[p.Id]] = append(votesPerParty[membership[p.Id]], weightedVote{name[p.Id], int(v)})
+	}
+	for p := range votesPerParty {
+		sort.Slice(votesPerParty[p], func(i, j int) bool { return votesPerParty[p][i].weight > votesPerParty[p][j].weight })
 	}
 
 	weighted := []weightedVote{}
@@ -123,7 +154,14 @@ func getRegionMandates(id int, numMandates int, membership map[int32]string, nam
 
 	mandates := map[string]int{}
 	for i := 0; i < numMandates; i++ {
-		mandates[weighted[i].party]++
+		mandates[weighted[i].entity]++
+	}
+
+	candidates := map[string][]string{}
+	for p := range votesPerParty {
+		for i := 0; i < mandates[p]; i++ {
+			candidates[p] = append(candidates[p], votesPerParty[p][i].entity)
+		}
 	}
-	return mandates, nil
+	return mandates, candidates, nil
 }