#!/usr/bin/python3

import pytz, re, sys
from datetime import datetime, timedelta

###############################################################################
# Steps through a time range with no regard for time zone.
# If either the start or end time is timezone aware, it is converted to UTC

def datetime_range( start, end, step ):
  if start.tzinfo is not None: start = start.astimezone( pytz.utc )
  if   end.tzinfo is not None:   end =   end.astimezone( pytz.utc )
  dt = start
  if step < timedelta( 0 ):
    while dt >= end:
      yield dt
      dt += step
  else:
    while dt <= end:
      yield dt
      dt += step

###############################################################################

def interval_to_timedelta( interval ):
  return _offset_to_timedelta( *_offset_parts( interval ) )

###############################################################################

def local_increment( dt, offset, tz ):
  return tz.localize( dt.replace( tzinfo=None ) + offset )

###############################################################################
# Steps through a time range within a given time zone, realigning each step
# with the current UTC offset for that time.  This is useful when DST changes
# the length of a local day, and you still want to step through the entire
# interval.

def local_range( start, end, step, tz ):
  dt = start
  if step < timedelta( 0 ):
    while dt >= end:
      yield dt
      dt = local_increment( dt, step, tz )
  else:
    while dt <= end:
      yield dt
      dt = local_increment( dt, step, tz )

###############################################################################
# Turns a time offset (relative to a given time, or the current time by
# default) into a datetime object.  Apply successive offsets in turn.

def offset_to_time( offset, time=datetime.now() ):
  for interval in re.findall( '([-+]?\d+[a-z]+)', offset, re.I ):
    ( count, units ) = _offset_parts( interval )
    if units.lower() == 'wy':
      yr = time.year
      if time.month < 10: yr -= 1
      yr += count
      time = datetime.strptime( "%d/10/01" % ( yr ), "%Y/%m/%d" )
    elif units.lower() == 'mo':
      mo = time.month - 1 + count
      yr = time.year + int( mo / 12 )
      if mo < 0 and mo % 12 > 0:      # Going back in time, we must decrement
        yr -= 1                       # year unless month falls on January
      mo %= 12
      mo += 1
      time = datetime.strptime( "%d/%d" % ( yr, mo ), "%Y/%m" )
    else:
      if units.lower() in 'wdhms':
        time += _offset_to_timedelta( count, units )
      elif units.lower().startswith( 'y' ):
        time += timedelta( years=count )
      if units.isupper():
        time = time.replace( hour=0, minute=0, second=0, microsecond=0 )
      elif units == 'h' or units == 'd':
        time = time.replace( minute=0, second=0, microsecond=0 )
      elif units == 'm':
        time = time.replace( second=0, microsecond=0 )
  return time

###############################################################################
# Breaks apart an offset value

def _offset_parts( offset ):
  count = 0
  units = None
  m = re.match( '([-+]?\d+)(\D+)?', offset )
  if m: ( count, units ) = m.group( 1, 2 )
  if units == None: units = 's'
  return( float( count ), units )

###############################################################################
# Turn an offset (weeks, days, hours, minutes, or seconds) into a number of
# seconds

def _offset_to_timedelta( count, units='s' ):
  count = float( count )
  u = units.lower()[0]
  if   u == 'w': count *= 604800   # weeks
  elif u == 'd': count *=  86400   # days
  elif u == 'h': count *=   3600   # hours
  elif u == 'm': count *=     60   # minutes
  elif u == 's': pass              # seconds
  elif u ==  '': pass              # also seconds
  else:
    print("Unknown unit type '%s' - assuming seconds" % ( units ), file=sys.stderr)
  return timedelta( seconds=count )