fmos.py

import requests
import os
import csv
from collections import Counter
from math import ceil
from tabulate import tabulate

TBA_API_BASE = "https://www.thebluealliance.com/api/v3"

def get_event_matches(event_key, api_key):
    url = f"{TBA_API_BASE}/event/{event_key}/matches/simple"
    headers = {"X-TBA-Auth-Key": api_key}
    r = requests.get(url, headers=headers)
    r.raise_for_status()
    matches = [m for m in r.json() if m["comp_level"] == "qm"]  # qualification only
    matches.sort(key=lambda x: (x["match_number"], x["key"]))
    return matches

def get_unique_teams(matches):
    teams = set()
    for m in matches:
        for side in ["red", "blue"]:
            teams.update(m["alliances"][side]["team_keys"])
    return teams

def choose_balanced_schedules(matches):
    watch_counts = Counter()
    red_schedule = []
    blue_schedule = []

    # Gather all teams for target calculation
    all_teams = get_unique_teams(matches)
    total_slots = 2 * len(matches)  # 2 watchers per match (red + blue)
    ideal_avg = total_slots / len(all_teams)
    ideal_target = ceil(ideal_avg)

    print(f"\nIdeal average watch count per team: {ideal_avg:.2f} (target {ideal_target})")
    print(f"Total unique teams: {len(all_teams)} | Total matches: {len(matches)}\n")

    for m in matches:
        red_teams = m["alliances"]["red"]["team_keys"]
        blue_teams = m["alliances"]["blue"]["team_keys"]

        # pick least-watched team from each side
        red_pick = min(red_teams, key=lambda t: watch_counts[t])
        watch_counts[red_pick] += 1
        red_schedule.append({"Match #": m["match_number"], "Team to Watch": red_pick})

        blue_pick = min(blue_teams, key=lambda t: watch_counts[t])
        watch_counts[blue_pick] += 1
        blue_schedule.append({"Match #": m["match_number"], "Team to Watch": blue_pick})

    # mark teams watched more than target
    for sched in (red_schedule, blue_schedule):
        for entry in sched:
            team = entry["Team to Watch"]
            if watch_counts[team] > ideal_target:
                entry["Team to Watch"] = f"{team} *extra*"

    return red_schedule, blue_schedule, watch_counts, ideal_target

def print_and_export_schedule(event_key, side, schedule, counts, ideal_target):
    print(f"\n===== {side.upper()} SIDE WATCH SCHEDULE =====")
    print(tabulate(schedule, headers="keys", tablefmt="pretty"))

    print(f"\nTeam appearance counts ({side} side cumulative):")
    for team, count in sorted(counts.items()):
        extra = " *extra*" if count > ideal_target else ""
        print(f"{team}: {count}{extra}")

    filename = f"{event_key}_{side}_balanced_watch_schedule.csv"
    with open(filename, "w", newline="", encoding="utf-8") as f:
        writer = csv.DictWriter(f, fieldnames=["Match #", "Team to Watch"])
        writer.writeheader()
        writer.writerows(schedule)
    print(f"✅ CSV exported: {filename}")

def main():
    api_key = "4KL2Rp6vKD1ghUczwu8uzzRXIjwlCjosTIdcofi0kcOxsXu69sDuIbTgk7OszEku"
    if not api_key:
        raise SystemExit("Please set environment variable TBA_API_KEY")

    event_key = input("Enter event key (e.g., 2024ilch): ").strip()

    matches = get_event_matches(event_key, api_key)
    if not matches:
        raise SystemExit("No qualification matches found for this event.")

    red_schedule, blue_schedule, watch_counts, ideal_target = choose_balanced_schedules(matches)

    print_and_export_schedule(event_key, "red", red_schedule, watch_counts, ideal_target)
    print_and_export_schedule(event_key, "blue", blue_schedule, watch_counts, ideal_target)

if __name__ == "__main__":
    main()