master_script.py

#!/usr/bin/env python3
import csv
import sys
import random

# -------------------------------
# Data Structures
# -------------------------------

class Ballot:
    def __init__(self, preferences, weight=1.0):
        # Each ballot stores the ordered list of candidate preferences.
        self.preferences = [pref.strip() for pref in preferences]
        self.weight = round(weight, 5)
        self.current_index = 0  # pointer to current active preference

    def get_current_candidate(self, candidates_status):
        # Return the next candidate on the ballot who is still active.
        while self.current_index < len(self.preferences):
            candidate = self.preferences[self.current_index]
            if candidates_status.get(candidate, 'active') == 'active':
                return candidate
            else:
                self.current_index += 1
        return None  # ballot exhausted

# -------------------------------
# STV Election Function (Revised)
# -------------------------------

def stv_election(ballots, seats):
    # Build the set of candidates from all ballots.
    candidate_set = set()
    for ballot in ballots:
        for pref in ballot.preferences:
            candidate_set.add(pref)
    # Initialize candidate state.
    candidates = {
        cand: {
            'status': 'active',   # 'active', 'elected', or 'eliminated'
            'votes': 0.0,         # current round vote total
            'ballots': []         # ballots allocated to this candidate
        }
        for cand in candidate_set
    }
    # For tie‑breaking, record vote totals from previous rounds.
    candidate_history = {cand: [] for cand in candidate_set}

    total_ballots = len(ballots)
    # Compute the threshold (quota) using ballots / (seats + 1)
    threshold = round(total_ballots / (seats + 1), 5)
    print("Total ballots:", total_ballots)
    print("Seats:", seats)
    print("Threshold:", threshold)

    elected_candidates = []    # List to record elected candidates in order.
    eliminated_candidates = [] # List for eliminated candidates.
    round_num = 0

    # All ballots are initially active.
    active_ballots = ballots[:]

    while len(elected_candidates) < seats and any(candidates[c]['status'] == 'active' for c in candidates):
        round_num += 1
        print("\n----- Round", round_num, "-----")
        # Reset vote counts and ballot lists for active candidates.
        for cand in candidates:
            if candidates[cand]['status'] == 'active':
                candidates[cand]['votes'] = 0.0
                candidates[cand]['ballots'] = []

        # Count votes by allocating each active ballot to its current active candidate.
        for ballot in active_ballots:
            current = ballot.get_current_candidate({c: candidates[c]['status'] for c in candidates})
            if current:
                candidates[current]['votes'] = round(candidates[current]['votes'] + ballot.weight, 5)
                candidates[current]['ballots'].append(ballot)

        # Record round vote totals for tie‑breaking.
        for cand in candidates:
            if candidates[cand]['status'] == 'active':
                candidate_history[cand].append(candidates[cand]['votes'])
            else:
                candidate_history[cand].append(0.0)

        # -------------------------------
        # Display round vote totals (only show candidates that are active or elected)
        # -------------------------------
        for cand in candidates:
            if candidates[cand]['status'] != 'eliminated':
                print(f"{cand}: {candidates[cand]['votes']} ({candidates[cand]['status']})")

        # -------------------------------
        # Elect candidates who meet/exceed the threshold.
        # Process them in descending order of votes.
        # -------------------------------
        to_elect = [cand for cand in candidates if candidates[cand]['status'] == 'active' and candidates[cand]['votes'] >= threshold]
        to_elect.sort(key=lambda c: candidates[c]['votes'], reverse=True)

        for cand in to_elect:
            if len(elected_candidates) < seats:
                candidates[cand]['status'] = 'elected'
                elected_candidates.append(cand)
                print(f"\nCandidate '{cand}' elected with {candidates[cand]['votes']} votes.")
                # Process surplus transfer if the candidate's votes exceed the threshold.
                if candidates[cand]['votes'] > threshold:
                    surplus = round(candidates[cand]['votes'] - threshold, 5)
                    transfer_value = round(surplus / candidates[cand]['votes'], 5)
                    print(f"Transferring surplus of {surplus} votes from '{cand}' with transfer value {transfer_value}.")
                    # Copy the ballots and then remove them from active_ballots
                    transferred_ballots = candidates[cand]['ballots'][:]
                    new_active = []
                    for ballot in transferred_ballots:
                        new_weight = round(ballot.weight * transfer_value, 5)
                        if new_weight > 0:
                            ballot.weight = new_weight
                            ballot.current_index += 1  # move to next preference
                            new_active.append(ballot)
                    # Remove transferred ballots from active_ballots so they're not double counted.
                    active_ballots = [b for b in active_ballots if b not in transferred_ballots]
                    # Add the reweighted ballots to the active pool.
                    active_ballots.extend(new_active)
            if len(elected_candidates) == seats:
                break  # Stop if we have filled all seats.

        if len(elected_candidates) == seats:
            break

        # If the number of remaining active candidates equals the number of seats left, elect them automatically.
        remaining_seats = seats - len(elected_candidates)
        active_cands = [c for c in candidates if candidates[c]['status'] == 'active']
        if remaining_seats == len(active_cands):
            for cand in active_cands:
                candidates[cand]['status'] = 'elected'
                elected_candidates.append(cand)
                print(f"\nCandidate '{cand}' elected by default (only remaining candidate).")
            break

        # -------------------------------
        # If no candidate was elected in this round, eliminate the candidate with the fewest votes.
        # -------------------------------
        if not to_elect:
            active_cands = [c for c in candidates if candidates[c]['status'] == 'active']
            if not active_cands:
                break
            min_votes = min(candidates[c]['votes'] for c in active_cands)
            tied = [c for c in active_cands if candidates[c]['votes'] == min_votes]
            if len(tied) > 1:
                # Use previous round vote totals to break tie.
                for prev in range(round_num - 2, -1, -1):
                    prev_votes = [(c, candidate_history[c][prev]) for c in tied]
                    min_prev = min(v for c, v in prev_votes)
                    tied = [c for c, v in prev_votes if v == min_prev]
                    if len(tied) == 1:
                        break
                if len(tied) > 1:
                    eliminated = random.choice(tied)
                    print(f"\nTie in elimination among {tied}; randomly eliminating '{eliminated}'.")
                else:
                    eliminated = tied[0]
            else:
                eliminated = tied[0]
            candidates[eliminated]['status'] = 'eliminated'
            eliminated_candidates.append(eliminated)
            print(f"\nCandidate '{eliminated}' eliminated with {candidates[eliminated]['votes']} votes.")
            # Remove ballots allocated to the eliminated candidate and transfer them.
            transferred_ballots = candidates[eliminated]['ballots'][:]
            new_active = []
            for ballot in transferred_ballots:
                ballot.current_index += 1  # move to next preference
                new_active.append(ballot)
            # Remove those ballots from active_ballots.
            active_ballots = [b for b in active_ballots if b not in transferred_ballots]
            active_ballots.extend(new_active)

        # Prepare for next round: keep only ballots that still have an active candidate.
        active_ballots = [b for b in active_ballots if b.get_current_candidate({c: candidates[c]['status'] for c in candidates}) is not None]

    # -------------------------------
    # Final Results
    # -------------------------------
    print("\n===== Final Results =====")
    print("Elected candidates (in order):")
    for cand in elected_candidates:
        print(" -", cand)
    if eliminated_candidates:
        print("\nEliminated candidates:")
        for cand in eliminated_candidates:
            print(" -", cand)
    return elected_candidates

# -------------------------------
# CSV Ballot Reader
# -------------------------------

def read_ballots_from_csv(file_path):
    ballots = []
    with open(file_path, newline='') as csvfile:
        reader = csv.DictReader(csvfile)
        for row in reader:
            # The "Candidates" column is expected to be a comma‐separated list.
            prefs = row["Candidates"].split(",")
            if len(prefs) < 5:
                print("Warning: ballot with fewer than five preferences:", prefs)
            ballots.append(Ballot(preferences=prefs))
    return ballots

# -------------------------------
# Main Entry Point
# -------------------------------

def main():
    if len(sys.argv) < 3:
        print("Usage: python stv.py <ballots_csv_file> <number_of_seats>")
        sys.exit(1)
    file_path = sys.argv[1]
    try:
        seats = int(sys.argv[2])
    except ValueError:
        print("Number of seats must be an integer.")
        sys.exit(1)
    ballots = read_ballots_from_csv(file_path)
    stv_election(ballots, seats)

if __name__ == "__main__":
    main()