Barak's Charts

charts.py

import os
import re
import csv
import sys

from collections import defaultdict
import matplotlib.pyplot as plt
from math import ceil

MATCH_REGEX = re.compile('(?P<type>PM|QM|QF|SF|F)[ -]?(?P<num>[0-9]+)')a
TRUE_STRING = 'TRUE'

SUCCESS = 'Successful'
FAIL = 'Failed'

class Data(object):
    def __init__(self, data):
        self.full_name = data[0]
        self.team_number = int(data[1])
        stripped_name = data[2].strip().upper()
        parsed_match_id = MATCH_REGEX.match(stripped_name)
        if parsed_match_id is None:
            print "Failed to parse match id: {}".format(stripped_name)
            self._match_id = {'type': None, 'num': None}
        else:
            self._match_id = parsed_match_id.groupdict()
        self.auto_run = data[3] == TRUE_STRING
        self.auto_exchange = data[4] == TRUE_STRING
        self.auto_switch = int(data[5] or 0)
        self.auto_switch_fail = int(data[6] or 0)
        self.auto_scale = int(data[7] or 0)
        self.auto_scale_fail = int(data[8] or 0)
        self.collection = data[9]
        self.tele_switch = int(data[10] or 0)
        self.tele_switch_fail = int(data[11] or 0)
        self.tele_scale = int(data[12] or 0)
        self.tele_scale_fail = int(data[13] or 0)
        self.tele_exchange = int(data[14] or 0)
        self.tele_exchange_fail = int(data[15] or 0)
        self.platform = data[16] == TRUE_STRING
        self.climb = data[17]
        self.partner_climb = data[18]
        self.tech_foul = data[19] == TRUE_STRING
        self.defence_comments = data[20]
        self.comments = data[21]
        
    def __repr__(self):
        return "<Data Team {} | Match {}>".format(self.team_number, self.match_name)
    
    @property
    def match_name(self):
        return "{type} {num}".format(**self._match_id)
            
    @property
    def match_key(self):
        if self._match_id['type'] == 'PM':
            return int(self._match_id['num']) - 100
        elif self._match_id['type'] == 'QM':
            return int(self._match_id['num'])
        elif self._match_id['type'] == 'QF':
            return 1000 + int(self._match_id['num'])
        elif self._match_id['type'] == 'SF':
            return 2000 + int(self._match_id['num'])
        elif self._match_id['type'] == 'F':
            return 3000 + int(self._match_id['num'])
        else:
            return 4000
        
    @property
    def exchange(self):
        return self.tele_exchange
        
    @property
    def switch(self):
        return self.tele_switch + self.auto_switch
        
    @property
    def scale(self):
        return self.tele_scale + self.auto_scale


def get_data():
    with open(sys.argv[1]) as fh:
        rd = csv.reader(fh)
        rd.next() # ignore headers
        return map(Data, rd)


def keep_data(data, top_pct=0.75):
    return int(ceil(len(data) * top_pct))


def average(data, keep=None):
    # Keep all.
    if keep is None:
        keep = len(data)
    
    # Average over that data
    data_sum = sum(sorted(data, reverse=True)[:keep])
    return float(data_sum) / keep

def plot_cube_handling(team_number):
    records = data_by_team[team_number]
    x = range(len(records))
    matches = [d.match_name for d in records]
    exchange = [d.exchange for d in records]
    switch = [d.switch for d in records]
    scale = [d.scale for d in records]
    
    total = map(sum, zip(exchange, switch, scale))
    keep = keep_data(total)
    total_average = average(total, keep)
    if keep < len(total):
        ignored_x = zip(*sorted(enumerate(total), key=lambda d: d[1])[:-keep])[0]
    else:
        ignored_x = []
    
    auto_run = [(1 if d.auto_run else 0) for d in records]
    auto_run_fail = [(1 if not d.auto_run else 0) for d in records]
    auto_switch = [(2 if d.auto_switch > 0 else 0) for d in records]
    auto_switch_fail = [(2 if d.auto_switch_fail > 0 else 0) for d in records]
    auto_scale = [(3 if d.auto_scale > 0 else 0) for d in records]
    auto_scale_fail = [(3 if d.auto_scale_fail > 0 else 0) for d in records]
    platform = [(6 if d.platform else 0) for d in records]
    climb = [(5 if d.climb == SUCCESS else 0) for d in records]
    climb_fail = [(5 if d.climb == FAIL else 0) for d in records]
    partner_climb = [(4 if d.partner_climb == SUCCESS else 0) for d in records]
    partner_climb_fail = [(4 if d.partner_climb == FAIL else 0) for d in records]
    
    plt.subplot(211)
    plt.title('Team {}'.format(team_number))
    lines = plt.stackplot(x, exchange, switch, scale)
    plt.plot([x[0], x[-1]], [total_average, total_average], 'r--', ignored_x, [5 for _ in ignored_x], 'rx')
    plt.xticks([], [])
    plt.xlim(x[0], x[-1])
    plt.yticks(range(13))
    plt.ylabel('Number of Cubes')
    plt.figlegend(lines, ['Exchange', 'Switch', 'Scale'])
    plt.grid(True)
    
    plt.subplot(212)
    plt.plot(x, auto_run, 'gv', x, auto_run_fail, 'rx', 
             x, auto_switch, 'gv', x, auto_switch_fail, 'rx', 
             x, auto_scale, 'gv', x, auto_scale_fail, 'rx',
             x, platform, 'bo',
             x, climb, 'gv', x, climb_fail, 'rx',
             x, partner_climb, 'gv', x, partner_climb_fail, 'rx')
    plt.yticks([1,2,3,4,5,6], ['[A] Run', '[A] Switch', '[A] Scale', 'Partner', 'Climb', 'Platform'])
    plt.ylim(0.5, 7.0)
    plt.xticks(x, matches, rotation=90)
    plt.xlim(x[0], x[-1])
    plt.grid(True)
    if not os.path.isdir('graphs'):
        os.mkdir('graphs')
    plt.savefig('graphs\\{}.png'.format(team_number))
    fig.clear()
    
    with open('graphs\\{}.txt'.format(team_number), 'wb') as fh:
        fh.write("Team {}\r\n".format(team_number))
        for record in records:
            fh.write("Match {}:\r\n".format(record.match_name))
            fh.write("Tech Fouls? {}\r\n".format(record.tech_foul))
            fh.write("{}\r\n".format(record.defence_comments))
            fh.write("{}\r\n\r\n".format(record.comments))
    
data_by_team = defaultdict(list)
    
if __name__ == "__main__":
    all_data = get_data()
    [data_by_team[data.team_number].append(data) for data in all_data]


    sorted_teams = []
    fig = plt.figure()
    fig.set_size_inches(9, 6)
        
    for team in sorted(data_by_team.keys()):
        data_by_team[team].sort(key=lambda data: data.match_key)
        records = data_by_team[team]
        exchange = [d.exchange for d in records]
        switch = [d.switch for d in records]
        scale = [d.scale for d in records]
        total = map(sum, zip(exchange, switch, scale))
        
        keep = keep_data(total)
        exchange = average(exchange, keep)
        switch = average(switch, keep)
        scale = average(scale, keep)
        total = average(total, keep)
        
        sorted_teams.append((team, exchange, switch, scale, total))
        matches = [d.match_name for d in records]
        if len(records) == 1:
            print "WARNING: only one record for team {}! typo?".format(team)
        elif len(set(matches)) < len(records):
            dupes = {match for match in matches if matches.count(match) > 1}
            print "WARNING: Duplicates for team {}: {}".format(team, dupes)
        else:
            print "Parsed team {}.".format(team)
            
        plot_cube_handling(team)
            
    with open("Results.txt", 'w') as fh:
        fh.write("Top Exchange Teams:\n")
        for i, (team, exchange, _, _, _) in enumerate(sorted(sorted_teams, key=lambda x: x[1], reverse=True)[:16]):
            fh.write("{:2}. {} - {:.2f}\n".format(i+1, team, exchange))
        
        fh.write("\nTop Switch Teams:\n")
        for i, (team, _, switch, _, _) in enumerate(sorted(sorted_teams, key=lambda x: x[2], reverse=True)[:16]):
            fh.write("{:2}. {} - {:.2f}\n".format(i+1, team, switch))
        
        fh.write("\nTop Scale Teams:\n")
        for i, (team, _, _, scale, _) in enumerate(sorted(sorted_teams, key=lambda x: x[3], reverse=True)[:16]):
            fh.write("{:2}. {} - {:.2f}\n".format(i+1, team, scale))
            
        fh.write("\nTop Total Teams:\n")
        for i, (team, _, _, _, total) in enumerate(sorted(sorted_teams, key=lambda x: x[4], reverse=True)[:16]):
            fh.write("{:2}. {} - {:.2f}\n".format(i+1, team, total))