#!/usr/bin/env python

"""

Recurrence rule calculation.

Copyright (C) 2014, 2015 Paul Boddie <paul@boddie.org.uk>

This program is free software; you can redistribute it and/or modify it under

the terms of the GNU General Public License as published by the Free Software

Foundation; either version 3 of the License, or (at your option) any later

version.

This program is distributed in the hope that it will be useful, but WITHOUT

ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS

FOR A PARTICULAR PURPOSE.  See the GNU General Public License for more

details.

You should have received a copy of the GNU General Public License along with

this program.  If not, see <http://www.gnu.org/licenses/>.

----

References:

RFC 5545: Internet Calendaring and Scheduling Core Object Specification

          (iCalendar)

          http://tools.ietf.org/html/rfc5545

----

FREQ defines the selection resolution.

DTSTART defines the start of the selection.

INTERVAL defines the step of the selection.

COUNT defines a number of instances

UNTIL defines a limit to the selection.

BY... qualifiers select instances within each outer selection instance according

to the recurrence of instances of the next highest resolution. For example,

BYDAY selects days in weeks. Thus, if no explicit week recurrence is indicated,

all weeks are selected within the selection of the next highest explicitly

specified resolution, whether this is months or years.

BYSETPOS in conjunction with BY... qualifiers permit the selection of specific

instances.

Within the FREQ resolution, BY... qualifiers refine selected instances.

Outside the FREQ resolution, BY... qualifiers select instances at the resolution

concerned.

Thus, FREQ and BY... qualifiers need to be ordered in terms of increasing

resolution (or decreasing scope).

"""

from calendar import monthrange

from datetime import date, datetime, timedelta

import operator

# Frequency levels, specified by FREQ in iCalendar.

freq_levels = (

    "YEARLY",

    "MONTHLY",

    "WEEKLY",

    None,

    None,

    "DAILY",

    "HOURLY",

    "MINUTELY",

    "SECONDLY"

    )

# Enumeration levels, employed by BY... qualifiers.

enum_levels = (

    None,

    "BYMONTH",

    "BYWEEKNO",

    "BYYEARDAY",

    "BYMONTHDAY",

    "BYDAY",

    "BYHOUR",

    "BYMINUTE",

    "BYSECOND"

    )

# Map from levels to lengths of datetime tuples.

lengths = [1, 2, 3, 3, 3, 3, 4, 5, 6]

positions = [l-1 for l in lengths]

# Map from qualifiers to interval units. Here, weeks are defined as 7 days.

units = {"WEEKLY" : 7}

# Make dictionaries mapping qualifiers to levels.

freq = dict([(level, i) for (i, level) in enumerate(freq_levels) if level])

enum = dict([(level, i) for (i, level) in enumerate(enum_levels) if level])

weekdays = dict([(weekday, i+1) for (i, weekday) in enumerate(["MO", "TU", "WE", "TH", "FR", "SA", "SU"])])

# Functions for structuring the recurrences.

def get_next(it):

    try:

        return it.next()

    except StopIteration:

        return None

def get_parameters(values):

    "Return parameters from the given list of 'values'."

    d = {}

    for value in values:

        parts = value.split("=", 1)

        if len(parts) < 2:

            continue

        key, value = parts

        if key in ("COUNT", "BYSETPOS"):

            d[key] = int(value)

        else:

            d[key] = value

    return d

def get_qualifiers(values):

    """

    Process the list of 'values' of the form "key=value", returning a list of

    qualifiers of the form (qualifier name, args).

    """

    qualifiers = []

    frequency = None

    interval = 1

    for value in values:

        parts = value.split("=", 1)

        if len(parts) < 2:

            continue

        key, value = parts

        if key == "FREQ" and freq.has_key(value):

            qualifier = frequency = (value, {})

        elif key == "INTERVAL":

            interval = int(value)

            continue

        elif enum.has_key(key):

            qualifier = (key, {"values" : get_qualifier_values(key, value)})

        else:

            continue

        qualifiers.append(qualifier)

    if frequency:

        frequency[1]["interval"] = interval

    return qualifiers

def get_qualifier_values(qualifier, value):

    """

    For the given 'qualifier', process the 'value' string, returning a list of

    suitable values.

    """

    if qualifier != "BYDAY":

        return map(int, value.split(","))

    values = []

    for part in value.split(","):

        weekday = weekdays.get(part[-2:])

        if not weekday:

            continue

        index = part[:-2]

        if index:

            index = int(index)

        else:

            index = None

        values.append((weekday, index))

    return values

def order_qualifiers(qualifiers):

    "Return the 'qualifiers' in order of increasing resolution."

    l = []

    for qualifier, args in qualifiers:

        if enum.has_key(qualifier):

            level = enum[qualifier]

            if special_enum_levels.has_key(qualifier):

                args["interval"] = 1

                selector = special_enum_levels[qualifier]

            else:

                selector = Enum

        else:

            level = freq[qualifier]

            selector = Pattern

        l.append(selector(level, args, qualifier))

    l.sort(key=lambda x: x.level)

    return l

def get_datetime_structure(datetime):

    "Return the structure of 'datetime' for recurrence production."

    l = []

    offset = 0

    for level, value in enumerate(datetime):

        if level == 2:

            offset = 3

        l.append(Enum(level + offset, {"values" : [value]}, "DT"))

    return l

def combine_datetime_with_qualifiers(datetime, qualifiers):

    """

    Combine 'datetime' with 'qualifiers' to produce a structure for recurrence

    production.

    """

    iter_dt = iter(get_datetime_structure(datetime))

    iter_q = iter(order_qualifiers(qualifiers))

    l = []

    from_dt = get_next(iter_dt)

    from_q = get_next(iter_q)

    have_q = False

    context = []

    context.append(from_dt.args["values"][0])

    # Consume from both lists, merging entries.

    while from_dt and from_q:

        _level = from_dt.level

        level = from_q.level

        # Datetime value at wider resolution.

        if _level < level:

            from_dt = get_next(iter_dt)

            context.append(from_dt.args["values"][0])

        # Qualifier at wider or same resolution as datetime value.

        else:

            if not have_q:

                if isinstance(from_q, Enum) and level > 0:

                    repeat = Pattern(level - 1, {"interval" : 1}, None)

                    repeat.context = tuple(context)

                    l.append(repeat)

                have_q = True

            from_q.context = tuple(context)

            l.append(from_q)

            from_q = get_next(iter_q)

            if _level == level:

                context.append(from_dt.args["values"][0])

                from_dt = get_next(iter_dt)

    # Complete the list.

    while from_dt:

        l.append(from_dt)

        from_dt = get_next(iter_dt)

    while from_q:

        if not have_q:

            if isinstance(from_q, Enum) and level > 0:

                repeat = Pattern(level - 1, {"interval" : 1}, None)

                repeat.context = tuple(context)

                l.append(repeat)

            have_q = True

        from_q.context = tuple(context)

        l.append(from_q)

        from_q = get_next(iter_q)

    return l

# Datetime arithmetic.

def combine(t1, t2):

    """

    Combine tuples 't1' and 't2', returning a copy of 't1' filled with values

    from 't2' in positions where 0 appeared in 't1'.

    """

    return tuple(map(lambda x, y: x or y, t1, t2))

def scale(interval, pos):

    """

    Scale the given 'interval' value to the indicated position 'pos', returning

    a tuple with leading zero elements and 'interval' at the stated position.

    """

    return (0,) * pos + (interval,)

def get_seconds(t):

    "Convert the sub-day part of 't' into seconds."

    t = t + (0,) * (6 - len(t))

    return (t[3] * 60 + t[4]) * 60 + t[5]

def update(t, step):

    "Update 't' by 'step' at the resolution of 'step'."

    i = len(step)

    # Years only.

    if i == 1:

        return (t[0] + step[0],) + tuple(t[1:])

    # Years and months.

    elif i == 2:

        month = t[1] + step[1]

        return (t[0] + step[0] + (month - 1) / 12, (month - 1) % 12 + 1) + tuple(t[2:])

    # Dates and datetimes.

    else:

        updated_for_months = update(t, step[:2])

        # Dates only.

        if i == 3:

            d = datetime(*updated_for_months)

            s = timedelta(step[2])

        # Datetimes.

        else:

            d = datetime(*updated_for_months)

            s = timedelta(step[2], get_seconds(step))

        return to_tuple(d + s, len(t))

def to_tuple(d, n=None):

    "Return 'd' as a tuple, optionally trimming the result to 'n' positions."

    if not isinstance(d, date):

        return d

    if n is None:

        if isinstance(d, datetime):

            n = 6

        else:

            n = 3

    return d.timetuple()[:n]

def get_first_day(first_day, weekday):

    "Return the first occurrence at or after 'first_day' of 'weekday'."

    first_day = date(*first_day)

    first_weekday = first_day.isoweekday()

    if first_weekday > weekday:

        return first_day + timedelta(7 - first_weekday + weekday)

    else:

        return first_day + timedelta(weekday - first_weekday)

def get_last_day(last_day, weekday):

    "Return the last occurrence at or before 'last_day' of 'weekday'."

    last_day = date(*last_day)

    last_weekday = last_day.isoweekday()

    if last_weekday < weekday:

        return last_day - timedelta(last_weekday + 7 - weekday)

    else:

        return last_day - timedelta(last_weekday - weekday)

# Classes for producing instances from recurrence structures.

class Selector:

    "A generic selector."

    def __init__(self, level, args, qualifier, selecting=None):

        """

        Initialise at the given 'level' a selector employing the given 'args'

        defined in the interpretation of recurrence rule qualifiers, with the

        'qualifier' being the name of the rule qualifier, and 'selecting' being

        an optional selector used to find more specific instances from those

        found by this selector.

        """

        self.level = level

        self.pos = positions[level]

        self.args = args

        self.qualifier = qualifier

        self.context = ()

        self.selecting = selecting

    def __repr__(self):

        return "%s(%r, %r, %r, %r)" % (self.__class__.__name__, self.level, self.args, self.qualifier, self.context)

    def materialise(self, start, end, count=None, setpos=None, inclusive=False):

        """

        Starting at 'start', materialise instances up to but not including any

        at 'end' or later, returning at most 'count' if specified, and returning

        only the occurrences indicated by 'setpos' if specified. A list of

        instances is returned.

        If 'inclusive' is specified, the selection of instances will include the

        end of the search period if present in the results.

        """

        start = to_tuple(start)

        end = to_tuple(end)

        counter = count and [0, count]

        results = self.materialise_items(self.context, start, end, counter, setpos, inclusive)

        results.sort()

        return results[:count]

    def materialise_item(self, current, earliest, next, counter, setpos=None, inclusive=False):

        """

        Given the 'current' instance, the 'earliest' acceptable instance, the

        'next' instance, an instance 'counter', and the optional 'setpos'

        criteria, return a list of result items. Where no selection within the

        current instance occurs, the current instance will be returned as a

        result if the same or later than the earliest acceptable instance.

        """

        if self.selecting:

            return self.selecting.materialise_items(current, earliest, next, counter, setpos, inclusive)

        elif earliest <= current:

            return [current]

        else:

            return []

    def convert_positions(self, setpos):

        "Convert 'setpos' to 0-based indexes."

        l = []

        for pos in setpos:

            lower = pos < 0 and pos or pos - 1

            upper = pos > 0 and pos or pos < -1 and pos + 1 or None

            l.append((lower, upper))

        return l

    def select_positions(self, results, setpos):

        "Select in 'results' the 1-based positions given by 'setpos'."

        results.sort()

        l = []

        for lower, upper in self.convert_positions(setpos):

            l += results[lower:upper]

        return l

    def filter_by_period(self, results, start, end, inclusive):

        """

        Filter 'results' so that only those at or after 'start' and before 'end'

        are returned.

        If 'inclusive' is specified, the selection of instances will include the

        end of the search period if present in the results.

        """

        l = []

        for result in results:

            if start <= result and (inclusive and result <= end or result < end):

                l.append(result)

        return l

class Pattern(Selector):

    "A selector of instances according to a repeating pattern."

    def materialise_items(self, context, start, end, counter, setpos=None, inclusive=False):

        first = scale(self.context[self.pos], self.pos)

        # Define the step between items.

        interval = self.args.get("interval", 1) * units.get(self.qualifier, 1)

        step = scale(interval, self.pos)

        # Define the scale of a single item.

        unit_interval = units.get(self.qualifier, 1)

        unit_step = scale(unit_interval, self.pos)

        current = combine(context, first)

        results = []

        while (inclusive and current <= end or current < end) and (counter is None or counter[0] < counter[1]):

            next = update(current, step)

            current_end = update(current, unit_step)

            interval_results = self.materialise_item(current, max(current, start), min(current_end, end), counter, setpos, inclusive)

            if counter is not None:

                counter[0] += len(interval_results)

            results += interval_results

            current = next

        return results

class WeekDayFilter(Selector):

    "A selector of instances specified in terms of day numbers."

    def materialise_items(self, context, start, end, counter, setpos=None, inclusive=False):

        step = scale(1, 2)

        results = []

        # Get weekdays in the year.

        if len(context) == 1:

            first_day = (context[0], 1, 1)

            last_day = (context[0], 12, 31)

        # Get weekdays in the month.

        elif len(context) == 2:

            first_day = (context[0], context[1], 1)

            last_day = update((context[0], context[1], 1), (0, 1, 0))

            last_day = update(last_day, (0, 0, -1))

        # Get weekdays in the week.

        else:

            current = context

            values = [value for (value, index) in self.args["values"]]

            while (inclusive and current <= end or current < end):

                next = update(current, step)

                if date(*current).isoweekday() in values:

                    results += self.materialise_item(current, max(current, start), min(next, end), counter, inclusive=inclusive)

                current = next

            if setpos:

                return self.select_positions(results, setpos)

            else:

                return results

        # Find each of the given days.

        for value, index in self.args["values"]:

            if index is not None:

                offset = timedelta(7 * (abs(index) - 1))

                if index < 0:

                    current = to_tuple(get_last_day(last_day, value) - offset, 3)

                else:

                    current = to_tuple(get_first_day(first_day, value) + offset, 3)

                next = update(current, step)

                # To support setpos, only current and next bound the search, not

                # the period in addition.

                results += self.materialise_item(current, current, next, counter, inclusive=inclusive)

            else:

                if index < 0:

                    current = to_tuple(get_last_day(last_day, value), 3)

                    direction = operator.sub

                else:

                    current = to_tuple(get_first_day(first_day, value), 3)

                    direction = operator.add

                while first_day <= current <= last_day:

                    next = update(current, step)

                    # To support setpos, only current and next bound the search, not

                    # the period in addition.

                    results += self.materialise_item(current, current, next, counter, inclusive=inclusive)

                    current = to_tuple(direction(date(*current), timedelta(7)), 3)

        # Extract selected positions and remove out-of-period instances.

        if setpos:

            results = self.select_positions(results, setpos)

        return self.filter_by_period(results, start, end, inclusive)

class Enum(Selector):

    def materialise_items(self, context, start, end, counter, setpos=None, inclusive=False):

        step = scale(1, self.pos)

        results = []

        for value in self.args["values"]:

            current = combine(context, scale(value, self.pos))

            next = update(current, step)

            # To support setpos, only current and next bound the search, not

            # the period in addition.

            results += self.materialise_item(current, current, next, counter, setpos, inclusive)

        # Extract selected positions and remove out-of-period instances.

        if setpos:

            results = self.select_positions(results, setpos)

        return self.filter_by_period(results, start, end, inclusive)

class MonthDayFilter(Enum):

    def materialise_items(self, context, start, end, counter, setpos=None, inclusive=False):

        last_day = monthrange(context[0], context[1])[1]

        step = scale(1, self.pos)

        results = []

        for value in self.args["values"]:

            if value < 0:

                value = last_day + 1 + value

            current = combine(context, scale(value, self.pos))

            next = update(current, step)

            # To support setpos, only current and next bound the search, not

            # the period in addition.

            results += self.materialise_item(current, current, next, counter, inclusive=inclusive)

        # Extract selected positions and remove out-of-period instances.

        if setpos:

            results = self.select_positions(results, setpos)

        return self.filter_by_period(results, start, end, inclusive)

class YearDayFilter(Enum):

    def materialise_items(self, context, start, end, counter, setpos=None, inclusive=False):

        first_day = date(context[0], 1, 1)

        next_first_day = date(context[0] + 1, 1, 1)

        year_length = (next_first_day - first_day).days

        step = scale(1, self.pos)

        results = []

        for value in self.args["values"]:

            if value < 0:

                value = year_length + 1 + value

            current = to_tuple(first_day + timedelta(value - 1), 3)

            next = update(current, step)

            # To support setpos, only current and next bound the search, not

            # the period in addition.

            results += self.materialise_item(current, current, next, counter, inclusive=inclusive)

        # Extract selected positions and remove out-of-period instances.

        if setpos:

            results = self.select_positions(results, setpos)

        return self.filter_by_period(results, start, end, inclusive)

special_enum_levels = {

    "BYDAY" : WeekDayFilter,

    "BYMONTHDAY" : MonthDayFilter,

    "BYYEARDAY" : YearDayFilter,

    }

# Public functions.

def connect_selectors(selectors):

    """

    Make the 'selectors' reference each other in a hierarchy so that

    materialising the principal selector causes the more specific ones to be

    employed in the operation.

    """

    current = selectors[0]

    for selector in selectors[1:]:

        current.selecting = selector

        current = selector

    return selectors[0]

def get_selector(dt, qualifiers):

    """

    Combine the initial datetime 'dt' with the given 'qualifiers', returning an

    object that can be used to select recurrences described by the 'qualifiers'.

    """

    dt = to_tuple(dt)

    return connect_selectors(combine_datetime_with_qualifiers(dt, qualifiers))

def get_rule(dt, rule):

    """

    Using the given initial datetime 'dt', interpret the 'rule' (a semicolon-

    separated collection of "key=value" strings), and return the resulting

    selector object.

    """

    if not isinstance(rule, tuple):

        rule = rule.split(";")

    qualifiers = get_qualifiers(rule)

    return get_selector(dt, qualifiers)

# vim: tabstop=4 expandtab shiftwidth=4