Assign week to quarters.

weeks_to_quarters.py

"""
Map the next 1000 weeks to quarters in a sensible way so that
most quarters are 13 weeks long... and occasionally we introduce
a 14 week quater, so as to maximise the amount of working days that
fall in their correct quarter
"""
from datetime import date, timedelta
from collections import defaultdict

import networkx as nx

JAN, FEB, MAR = 1, 2, 3
APR, MAY, JUN = 4, 5, 6
JUL, AUG, SEP = 7, 8, 9
OCT, NOV, DEC = 10, 11, 12

Q1 = 'Q1'
Q2 = 'Q2'
Q3 = 'Q3'
Q4 = 'Q4'

NEXT_YEAR = +1
THIS_YEAR = 0
LAST_YEAR = -1

MO2QTR = {SEP: (Q1, NEXT_YEAR),
          OCT: (Q1, NEXT_YEAR),
          NOV: (Q1, NEXT_YEAR),
          DEC: (Q2, NEXT_YEAR),
          JAN: (Q2, THIS_YEAR),
          FEB: (Q2, THIS_YEAR),
          MAR: (Q3, THIS_YEAR),
          APR: (Q3, THIS_YEAR),
          MAY: (Q3, THIS_YEAR),
          JUN: (Q4, THIS_YEAR),
          JUL: (Q4, THIS_YEAR),
          AUG: (Q4, THIS_YEAR)}


class Week:
    """ Represents a single week """

    def __init__(self, monday):
        """ Initialize a week object """
        if monday.weekday() != 0:
            raise ValueError('%s is not a Monday' % monday)
        self.monday = monday
        self.weekdays = [self.monday+timedelta(i) for i in range(5)]
        # Establish the candidate quarters for this week.
        self.quarters = defaultdict(int)
        for weekday in self.weekdays:
            year, month = weekday.year, weekday.month
            qtr, ydelta = MO2QTR[month]
            quarter = '%s-%s' % (year+ydelta, qtr)
            self.quarters[quarter] += 1
        ordered_quarters = sorted(self.quarters.items(),
                                  key=lambda x: x[1],
                                  reverse=True)
        self.pref_qtr = ordered_quarters[0][0]
        self.alt_qtr = None
        if len(ordered_quarters) > 1:
            self.alt_qtr = ordered_quarters[1][0]
        self.selected_qtr = None

    def min_qtr(self):
        """ Return the earliest quarter """
        return min(self.pref_qtr, self.alt_qtr)

    def max_qtr(self):
        """ Return the latest quarter """
        return max(self.pref_qtr, self.alt_qtr)

    def __lt__(self, other):
        return self.monday < other.monday

    def __repr__(self):
        if self.alt_qtr:
            return '<Week %s *%s* %s>' % (self.monday, self.pref_qtr, self.alt_qtr)
        return '<Week %s %s>' % (self.monday, self.pref_qtr)


def next_weekday(dtime, weekday=0):
    """ Get the next weekday after date d
    0 = Monday, 1 = Tuesday, ... 6=Sunday """
    days_ahead = weekday - dtime.weekday()
    if days_ahead <= 0:  # Target day already happened this week
        days_ahead += 7
    nextwd = dtime + timedelta(days_ahead)
    return date(year=nextwd.year, month=nextwd.month, day=nextwd.day)


def gen_weeks(num=1000, start_date=date.today()):
    """ Create n weeks from the monday before start_date, onwards """
    last_monday = next_weekday(start_date) - timedelta(days=7)
    weeks = []
    for _ in range(num):
        week = Week(last_monday)
        weeks.append(week)
        last_monday += timedelta(7)
    return weeks


def make_graph(undecided):
    """ Make a graph of the weeks listed in undecided. Edges are matching qtrs """
    graph = nx.Graph()
    common_qtrs = defaultdict(set)
    for week in undecided:
        graph.add_node(week)
        common_qtrs[week.pref_qtr].add(week)
        common_qtrs[week.alt_qtr].add(week)
    for nodes in common_qtrs.values():
        if len(nodes) == 1:
            continue
        left, right = list(nodes)
        graph.add_edge(left, right)
    return graph


def select_qtrs(weeks):
    """ given a sequence of weeks, select the quarters they'll go into """
    qtow = defaultdict(set)
    undecided = []
    for week in weeks:
        if len(week.quarters) == 1:
            # Must pick preferred quarter
            qtow[week.pref_qtr].add(week)
        else:
            undecided.append(week)

    trivial_progress = True
    while trivial_progress:
        trivial_progress = False
        decided = []
        for week in undecided:
            if len(qtow[week.pref_qtr]) == 12 and len(qtow[week.alt_qtr]) == 13:
                # Trivial
                qtow[week.pref_qtr].add(week)
                decided.append(week)
                trivial_progress = True
        for week in decided:
            undecided.remove(week)

    print("Need to decide where %s of %s weeks fall" % (len(undecided), len(weeks)))
    # Lets make a graph of the undecided weeks to see which are connected runs.
    graph = make_graph(undecided)

    runs = sorted(nx.connected_component_subgraphs(graph), key=lambda x: -len(x.nodes()))
    print("Need to decide where %s runs of weeks fall" % len(runs))

    decided_cnt = 0
    for run in runs:
        # Each connected run of qtr spanning weeks needs to be scheduled early or late.
        # Pick whichever leads to the largest number of correclty aligned days.
        early, late = 0, 0
        for week in run.nodes():
            # NB this breaks if you name the quartes something that doesn't sort alphabetically.
            early += week.quarters[week.min_qtr()]
            late += week.quarters[week.max_qtr()]
        for week in run.nodes():
            if early > late:
                qtow[week.min_qtr()].add(week)
            elif late > early:
                qtow[week.max_qtr()].add(week)
            decided_cnt += 1

    return qtow


def print_assignments(qtow):
    """ Pretty print the quarter to week mapping. """
    for qtr in sorted(qtow.keys()):
        weeks = qtow[qtr]
        if not weeks:
            continue
        first_week = min(weeks)
        last_week = max(weeks)
        num = len(weeks)
        print('%s\tWeek 1 starting %s\tWeek %s starting %s' % (qtr,
                                                               first_week.monday,
                                                               num,
                                                               last_week.monday))


def calculate():
    """ Create loads of week objects and select quarters for them. """
    weeks = gen_weeks(num=3000, start_date=date(2010, 9, 10))
    assignments = select_qtrs(weeks)
    print_assignments(assignments)

if __name__ == '__main__':
    calculate()