#!/usr/bin/python3

from datetime import datetime, timedelta
import re

###############################################################################
# Increments a value in a dictionary, creating a new value if necessary

def _increment( dict, key, value ):
  if key in dict and dict[key] is not None and value is not None:
    dict[key] += value
  else:
    dict[key] = value

def _seconds( delta ):
  return delta.seconds + delta.days * 24 * 3600

###############################################################################
# Determines the offset from <time> to the <interval> step just before
# <snaptime>, adding <additional> <interval>s if requested.
#
# Useful for jumping forward to the first datum to consider for snapping or
# computing statistics on a time series.

def _jump_to( time, snaptime, interval, additional = 0 ):
  whole_intervals = int( _seconds( time - snaptime ) //
                         _seconds( interval ) + additional )
  return( interval * whole_intervals )

def op( operator, a, b ):
  if not re.match( r'^[-+/*]$', operator ):
    raise TypeError( "Unknown operator: " + operator )
  if a is None or b is None: return None
  if   operator == '+': return a + b
  elif operator == '-': return a - b
  elif operator == '*': return a * b
  elif operator == '/':
    if b == 0:
      return None
    else:
      return a / b

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

class TimeSeries:
  
  def __init__( self ):
    self.ts = {}
    
  def math( self, operator, operand, style = 'intersect' ):
    result = TimeSeries()
    if isinstance( operand, TimeSeries ):
      if style == 'union':
        result.set_data_from_dict( self.ts )

      for date in operand.ts:
        if date in self.ts:
          result.set_datum( date,
                            op( operator, self.ts[date], operand.ts[date] ) )
        elif style == 'union':
          result.set_datum( date, op( operator, 0.0, operand.ts[date] ) )
          
    else:				# Scalar, apply to every value
      operand = float( operand )
      for date in self.ts:
        result.set_datum( date, op( operator, self.ts[date], operand ) )
    return result

  def add( self, operand, style='intersect' ):
    return self.math( '+', operand, style )

  def subtract( self, operand, style='intersect' ):
    return self.math( '-', operand, style )

  def multiply( self, operand, style='intersect' ):
    return self.math( '*', operand, style )

  def divide( self, operand, style='intersect' ):
    return self.math( '/', operand, style )

  def dump_its( self ):
    ret = ''
    for date in sorted ( self.ts ):
      ret += str( date )
      if self.ts[date] is None:
        ret += " ---\n"
      else:
        ret += " %9.2f\n" % ( self.ts[date] )
    return ret

  def duplicate( self ):
    dup = TimeSeries()
    dup.set_data_from_dict( self.ts )
    return dup

  def domain( self ):
    keys = sorted( self.ts.keys() )
    return( keys[0], keys[-1] )

  def first( self ): return self.domain()[0]
  def  last( self ): return self.domain()[1]

  def first_val( self ):
    first = self.first()
    return [ first, self.ts[first] ]

  def last_val( self ):
    last = self.last()
    return [ last, self.ts[last] ]

  def range( self ):
    vals = sorted( self.ts.values() )
    return( vals[0], vals[-1] )

  def min( self ): return self.range()[0]
  def max( self ): return self.range()[1]

  def dict( self ):
    return self.ts.copy()

  def pairs( self ):
    return list(self.ts.items())

  def __len__( self ):
    return len( self.ts )

  def set_data_from_dict( self, dict ):
    self.ts = dict

  def set_data_from_time_series( self, ts ):
    self.ts = ts.ts

  def set_datum( self, date, value ):
    self.ts[date] = value

  def shift( self, count, units ):
    shift_ts = TimeSeries()
    for date in list(self.ts.keys()):
      shift_ts.set_datum( shift_date( date, count, units ), self.ts[date] )
    return shift_ts

  #############################################################################
  # Create a new timeseries in which all values are shifted forward if during
  # daylight savings.  This is useful if a timeseries is regular to standard
  # time, but must be reported as regular to the local time.

  def offset_to_dst( self ):
    dst_ts = TimeSeries()
    for date in sorted( self.ts.keys ):
      dst_date -= dst_offset( date )
      dst_ts.set_datum( dst_date, self.ts[date] )
    return dst_ts
      
  #############################################################################
  # Snap the timeseries to a given interval, allowing variance within a given
  # before-after window.  Apply a within-interval shift if requested.

  def snap( self, interval, before=None, after=None,
            shift=timedelta( seconds=0 ), start=None ):

    if _seconds( interval ) == 0: return self.duplicate()
   
    if before is None: before = interval / 2
    if after  is None: after  = before

    # <before> and <after> values greater than half the snap interval could
    # cause values to be snapped to a prior value when they should more
    # appropriately be snapped to a later value.  So have these max out at
    # <interval> / 2

    if before > interval / 2: before = interval / 2
    if after  > interval / 2: after  = interval / 2

    if start is None:
      int_secs = _seconds( interval )
      if int_secs < 86400:
       per_day = 86400 / int_secs
       if per_day != int( per_day ):
          raise UserWarning( "snap: %d-second interval does not fit evenly " +
	    "into a day.  Aligning to first day in range.  Hope that's OK." %
	    ( interval.seconds ) )
      else:
        days_per = int_secs / 86400
        if days_per != int( days_per ):
           raise UserWarning( "snap:  %-day interval is not a whole number " +
	    "of days.  Aligning to first day in range.  Hope that's OK." % 
	    ( days_per.seconds ) )

    snap_ts = TimeSeries()
    snap_tm = None
    keys = sorted( [x for x in list(self.ts.keys()) if self.ts[x] is not None] )
    if len( keys ) == 0: return snap_ts

    if start is not None:
      snap_tm = start + shift
    else:
      tm = keys[0]
      snap_tm = tm.replace( hour=0, minute=0, second=0, microsecond=0 )
      snap_tm += _jump_to( tm, snap_tm, interval )
      snap_tm += shift

    final_tm = ( keys[-1] + before )

    # For each snap time:
    #   if src[snap] exists and is not None, store it
    #   otherwise:
    #     step through src times until time > snap
    #     compute offset to src[time] and src[prior_time]
    #     store whichever src has the smallest offset, if within snap window

    i = 0
    while snap_tm < final_tm:
      if snap_tm in self.ts and self.ts[snap_tm] is not None:
        snap_ts.set_datum( snap_tm - shift, self.ts[snap_tm] )
        snap_tm += interval
        continue
      moved = 0
      while( i < len( keys ) and keys[i] < snap_tm ):
        i += 1
        moved += 1
      if moved and i >= len( keys ): i -= 1
      if keys[i]:
        offset_before = abs( snap_tm - keys[i-1] )
        offset_after  = abs( snap_tm - keys[i]   )
        if offset_before > offset_after:
          if offset_after <= after:
            snap_ts.set_datum( snap_tm - shift, self.ts[keys[i]] )
            i -= 1
          elif offset_before <= before:
            snap_ts.set_datum( snap_tm - shift, self.ts[keys[i-1]] )
            i -= 2
      if i < 0: i = 0
      snap_tm += interval
    return snap_ts

  #############################################################################
  # Return interval statistics for the time series for a given interval.

  def stat( self, interval, snap=None ):
    ( avg, ct, delta, first,
      last, max, min, range, tot ) = ( {}, {}, {}, {}, {}, {}, {}, {}, {} )
    zero_delta = timedelta( seconds = 0 )
    
    for date in sorted( self.ts.keys() ):
      
      # ONCE:  Align all our series to the first day in the data range.  Set
      # the first timestamp to the next interval at or after the first datum.
      # We will store our first set of series values there.

      if not snap:
        snap = date.replace( hour=0, minute=0, second=0, microsecond=0 )
        if date - snap > zero_delta:
          snap += _jump_to( date, snap, interval, 1 )


      # Step forward by intervals until we are after the current datum.

      while date > snap:
        if( snap in max and max[snap] is not None and
            snap in min and min[snap] is not None ):
          range[snap] = max[snap] - min[snap]
        if( snap in last  and  last[snap] is not None and
            snap in first and first[snap] is not None ):
          delta[snap] = last[snap] - first[snap]
        if( snap in tot and tot[snap] is not None and
            snap in ct  and  ct[snap] is not None ):
          avg[snap] = tot[snap] / ct[snap]
        snap += interval

      # This datum is within the interval range of the current snapped
      # timestep.  If it has a value, update appropriate stats.

      if self.ts[date] is not None:
        if snap not in min or self.ts[date] < min[snap]:
          min[snap] = self.ts[date]
        if snap not in max or self.ts[date] > max[snap]:
          max[snap] = self.ts[date]
        _increment( tot, snap, self.ts[date] )
        _increment( ct, snap, 1 )
        if snap not in first:
          first[snap] = self.ts[date]
        last[snap] = self.ts[date]

    # Compute final range, delta, and average, if any.

    if snap is not None:
      if( snap in max and max[snap] is not None and
          snap in min and min[snap] is not None ):
        range[snap] = max[snap] - min[snap]
      if( snap in last  and  last[snap] is not None and
          snap in first and first[snap] is not None ):
        delta[snap] = last[snap] - first[snap]
      if( snap in tot and tot[snap] is not None and
          snap in ct  and  ct[snap] is not None ):
        avg[snap] = tot[snap] / ct[snap]

    return( min, max, tot, ct, avg, range, delta )

  #############################################################################
  # Construct a rolling average dict for the time series given a lookback
  # interval and an initial value

  def rollavg( self, interval, lookback, stime ):
    ( avg, ct, tot ) = ( {}, {}, {} )
    zero_delta     = timedelta( seconds = 0 )

    snap = stime.replace( hour=0, minute=0, second=0, microsecond=0 )
    snap += _jump_to( stime, snap, interval )

    for date in sorted( self.ts.keys() ):

      # Step forward by intervals until within <lookback> of the current datum

      while date - snap > zero_delta:
        if( snap in tot and tot[snap] is not None and
            snap in ct  and  ct[snap] is not None ):
          avg[snap] = tot[snap] / ct[snap]
        snap += interval

      # This datum is within the interval range of one or more snapped
      # timesteps.  Set appropriate values.

      s = snap
      while s - date < lookback:
        _increment( tot, s, self.ts[date] )
        _increment( ct, s, 1 )
        s += interval

    # Compute final average, if any.

    if( snap in tot and tot[snap] is not None and
        snap in ct  and  ct[snap] is not None ):
      avg[snap] = tot[snap] / ct[snap]

    return avg

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

def dump_dict( dict, title = '' ):
  ret = title + "\n"
  for date in sorted ( dict ):
    if dict[date] is None: value = "None"
    else:                  value = "%9.2f" % dict[date]
    ret += "%s %s\n" % ( date.isoformat(), value )
  return ret

if __name__ == '__main__':

  ts = TimeSeries()

  etime = datetime.today().replace( hour=0, minute=0, second=0, microsecond=0 )
  stime = etime - timedelta( days=3 )
  tm = stime
  value = 1
  one_hour = timedelta( hours=1 )
  while tm < etime:
    if value in [ 1, 24, 25, 26, 49, 61, 75, 89, 107, 108, 109 ]:
      ts.set_datum( tm, None )
    else:
      ts.set_datum( tm, value )
    value += 1
    tm += one_hour

  print("Whole time series:")
  print(ts.dump_its())

  print("Snapped to 1d:")
  snapped = ts.snap( 86400 )
  print(snapped.dump_its())

  rollavg = ts.rollavg( 86400, 86400, stime )
  print(dump_dict( rollavg, "Rolling average 1d:" ))

  ( min, max, tot, ct, avg, range, delta ) = ts.stat( 86400 )
  print(dump_dict( min,   "Minimum 1d:" ))
  print(dump_dict( max,   "Maximum 1d:" ))
  print(dump_dict( tot,   "Total 1d:" ))
  print(dump_dict( ct,    "Count 1d:" ))
  print(dump_dict( avg,   "Average 1d:" ))
  print(dump_dict( range, "Range 1d:" ))
  print(dump_dict( delta, "Delta 1d:" ))