#!/usr/bin/env python3
import argparse
import csv
from functools import partial
import itertools
import matplotlib.pyplot as plt
import numpy
from operator import itemgetter
import re

def readcsv(filename):
    """Read a (BlackBoard) CSV file into a header and data"""
    with open(filename, 'r') as csvfile:
        reader = csv.reader(csvfile, delimiter='\t', quotechar='"')
        data = []
        header = next(reader)
        for row in reader:
            data.append(row)
        return header, numpy.array(data)

def parse_float(f):
    """Parse a string into a float and return 0. when that fails"""
    try:
        return float(f)
    except ValueError:
        return 0.

def parse_floats(matrix):
    """Apply parse_float on all elements of a matrix"""
    m_new = []
    for row in matrix:
        m_new.append([parse_float(v) for v in row])
    return m_new

def remove_zeros(matrix):
    """Remove zeros from a matrix

    Typically, BlackBoard grade lists will contain many zeros that we don't
    want to distort the boxplots. This function removes zeros on a per-column
    basis."""
    m_new = []
    for col in matrix:
        m_new.append([v for v in col if v != 0.0])
    return m_new

def normalise(matrix):
    """Normalise grades to a 0-100 range
    
    Some grades are given in a 0-10 range, others in a 0-100. This function
    normalises *but these two* into a 0-100 range"""
    m_new = []
    for row in matrix:
        m_new.append(row if max(row) > 10.0 else [10 * i for i in row])
    return m_new

def strip_header(h):
    """Strip common BlackBoard additions in the CSV header"""
    return h.split('[Total Pts:')[0]

def remove_empty_lists(headers, data):
    """Remove empty columns and the corresponding headers from a matrix"""
    new_headers, new_data = [], []
    for h, col in zip(headers, numpy.transpose(data)):
        if not all([x == 0.0 for x in col]):
            new_headers.append(h)
            new_data.append(col)

    return new_headers, new_data

def header_regex_callback(header, regex='', invert=False):
    """Check whether a regex occurs in a header

    This is an example of a possible callback function."""
    match = re.compile(regex).search(header) != None
    return match != invert

def participant_wherehas_callback(headers, participant, regex='', invert=False):
    """Check whether a participant has a grade for some assignment"""
    regex = re.compile(regex)
    mh = [regex.search(h) != None for h in headers[6:]]
    for h, d in zip(mh, participant[6:]):
        if h and d != '':
            return (not invert)
    return invert

def check_header_callback(headers, data, header_callback):
    """For each header, check that we want to show it, and remove data if not"""
    new_headers, new_data = [], []
    for h, d in zip(headers, data):
        if header_callback(h):
            new_headers.append(h)
            new_data.append(d)
    return new_headers, new_data

def check_participant_callback(headers, data, participant_callback):
    """For each participant, check that we want to show it, and remove if not"""
    new_data = []
    for d in data:
        if participant_callback(headers, d):
            new_data.append(d)
    return new_data

def make_groups(headers, data, column):
    i = 0
    for h in headers:
        if column in h:
            break
        i += 1

    data = list(numpy.transpose(data))
    data.sort(key=itemgetter(i))
    data = [list(x) for _, x in itertools.groupby(data, itemgetter(i))]
    new_headers = []
    new_data = []
    j = 0
    for h in headers:
        for group in data:
            if not all([x[j] == 0.0 for x in group]):
                new_headers.append(h + ' (group: ' + str(group[0][i]) + ')')
                new_data.append([x[j] for x in group])
        j += 1

    return new_headers, new_data

def plotgrades(headers, data, skip=0,
        participant_callback=lambda x:True, header_callback=lambda x:True,
        group_participants_column=None):
    """Plot grades corresponding to headers in a boxplot"""
    data = check_participant_callback(headers, data, participant_callback)
    headers, data = map(strip_header, headers[skip:]), data[skip:]
    data = parse_floats(data)
    headers, data = remove_empty_lists(headers, data)
    if group_participants_column is not None:
        headers, data = make_groups(headers, data, group_participants_column)
    data = remove_zeros(data)
    headers, data = check_header_callback(headers, data, header_callback)
    data = normalise(data)

    if len(data) > 0:
        plt.boxplot(data)
    
        ax = plt.gca()
        plt.xticks(range(1, len(data) + 1), headers, rotation=90)
        ax.set_ylim([0,100])
    
        plt.show()

def parse_args():
    """Parse command line arguments"""
    pars = argparse.ArgumentParser(description='Plot BlackBoard grades')
    
    pars.add_argument('-s', '--skip', metavar='n', type=int, default=0,
            help='Skip the first n columns')
    pars.add_argument('-w', '--where', metavar='regex', default='',
            help='Restrict what grades are shown with a regex')
    pars.add_argument('-wh', '--where-has', metavar='regex', default='',
            help='Only count participants with a grade for a matching item')
    pars.add_argument('-g', '--group-by', metavar='column', default=None,
            help='Group participants by this column')
    pars.add_argument('-iw', '--invert-where', action='store_true',
            help='Invert --where regex')
    pars.add_argument('-iwh', '--invert-where-has', action='store_true',
            help='Invert --where-has regex')
    pars.add_argument('filename', metavar='file',
            help='The CSV file as exported by BlackBoard')

    return pars.parse_args()

def main():
    """Plot BlackBoard grades from a CSV file in boxplots"""
    args = parse_args()
    headers, data = readcsv(args.filename)

    plotgrades(headers, data, skip=args.skip,
            header_callback=partial(header_regex_callback,
                regex=args.where, invert=args.invert_where),
            participant_callback=partial(participant_wherehas_callback,
                regex=args.where_has, invert=args.invert_where_has),
            group_participants_column=args.group_by)

if __name__ == '__main__':
    main()