#!/usr/bin/env python

"""

Recurrence rule calculation.

Copyright (C) 2014, 2015, 2017 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"

    )

# Levels defining days.

daylevels = [2, 3, 4, 5]

# 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]

# Define the lowest values at each resolution (years, months, days... hours,

# minutes, seconds).

firstvalues = [0, 1, 1, 1, 1, 1, 0, 0, 0]

# 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

        # Accept frequency indicators as qualifiers.

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

            qualifier = frequency = (value, {})

        # Accept interval indicators for frequency qualifier parameterisation.

        elif key == "INTERVAL":

            interval = int(value)

            continue

        # Accept enumerators as qualifiers.

        elif enum.has_key(key):

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

        # Ignore other items.

        else:

            continue

        qualifiers.append(qualifier)

    # Parameterise any frequency qualifier with the interval.

    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:

        # Distinguish between enumerators, used to select particular periods,

        # and frequencies, used to select repeating periods.

        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):

        # At the day number, adjust the frequency level offset to reference

        # days (and then hours, minutes, seconds).

        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.

    Initial datetime values appearing at broader resolutions than any qualifiers

    are ignored, since their details will be used when materialising the

    results.

    Qualifiers are accumulated in order to define a selection. Datetime values

    occurring between qualifiers or at the same resolution as qualifiers are

    ignored.

    Any remaining datetime values are introduced as enumerators, provided that

    they do not conflict with qualifiers. For example, specific day values

    conflict with day selectors and weekly qualifiers.

    The purpose of the remaining datetime values is to define a result within

    a period selected by the most precise qualifier. For example, the selection

    of a day and month in a year recurrence.

    """

    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

    # 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)

        # Qualifier at wider or same resolution as datetime value.

        else:

            if not have_q:

                add_initial_qualifier(from_q, level, l)

                have_q = True

            # Add the qualifier to the combined list.

            l.append(from_q)

            # Datetime value at same resolution.

            if _level == level:

                from_dt = get_next(iter_dt)

            # Get the next qualifier.

            from_q = get_next(iter_q)

    # Complete the list by adding remaining datetime enumerators.

    while from_dt:

        # Ignore datetime values that conflict with day-level qualifiers.

        if not l or from_dt.level != freq["DAILY"] or l[-1].level not in daylevels:

            l.append(from_dt)

        from_dt = get_next(iter_dt)

    # Complete the list by adding remaining qualifiers.

    while from_q:

        if not have_q:

            add_initial_qualifier(from_q, level, l)

            have_q = True

        # Add the qualifier to the combined list.

        l.append(from_q)

        # Get the next qualifier.

        from_q = get_next(iter_q)

    return l

def add_initial_qualifier(from_q, level, l):

    """

    Take the first qualifier 'from_q' at the given resolution 'level', using it

    to create an initial qualifier, adding it to the combined list 'l' if

    required.

    """

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

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

        l.append(repeat)

# 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, first=False):

        """

        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.args = args

        self.qualifier = qualifier

        self.selecting = selecting

        self.first = first

        # Define the index of values from datetimes involved with this selector.

        self.pos = positions[level]

    def __repr__(self):

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

    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(start, 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 time periods according to a repeating pattern."

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

        """

        Bounded by the given 'context', return periods within 'start' to 'end',

        updating the 'counter', selecting only the indexes specified by 'setpos'

        (if given).

        If 'inclusive' is specified, the selection of periods will include those

        starting at the end of the search period, if present in the results.

        """

        # Define the step between result periods.

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

        step = scale(interval, self.pos)

        # Define the scale of a single period.

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

        unit_step = scale(unit_interval, self.pos)

        # Employ the context as the current period if this is the first

        # qualifier in the selection chain.

        if self.first:

            current = context[:self.pos+1]

        # Otherwise, obtain the first value at this resolution within the

        # context period.

        else:

            first = scale(firstvalues[self.level], self.pos)

            current = combine(context[:self.pos], first)

        results = []

        # Obtain periods before the end (and also at the end if inclusive),

        # provided that any limit imposed by the counter has not been exceeded.

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

              (counter is None or counter[0] < counter[1]):

            # Increment the current datetime by the step for the next period.

            next = update(current, step)

            # Determine the end point of the current period.

            current_end = update(current, unit_step)

            # Obtain any period or periods within the bounds defined by the

            # current period and any contraining start and end points, plus

            # counter, setpos and inclusive details.

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

            # Update the counter with the number of identified results.

            if counter is not None:

                counter[0] += len(interval_results)

            # Accumulate the results.

            results += interval_results

            # Visit the next instance.

            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[:self.pos], 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)