Energies

#!/usr/bin/env python3

"""
"""


import argparse

import numpy as np
import pandas as pd
from scipy import signal

import colors
import fetcher
import printing


UP_ARROW = '\u25b2'
DOWN_ARROW = '\u25bc'

# Constants specific to the data source used
CLOSE  = fetcher.DataSource.CLOSE
OPEN   = fetcher.DataSource.OPEN
HIGH   = fetcher.DataSource.HIGH
LOW    = fetcher.DataSource.LOW
VOLUME = fetcher.DataSource.VOLUME


class PivotPoints(object):
  """Container class to hold pivot points."""


def Extrema(df, period):
  maxima = signal.argrelextrema(df.values, np.greater)
  minima = signal.argrelextrema(df.values, np.less)


def SMA(df, period):
  return df[CLOSE].rolling(window=period).mean()


def EMA(df, period):
  return df[CLOSE].ewm(span=period).mean()


def Stoch(df, period_k, period_d, smoothing):
  high, low, close = df[HIGH], df[LOW], df[CLOSE]
  lowest_low = low.rolling(period_k).min()
  highest_high = high.rolling(period_k).max()
  fast_k = 100 * (close - lowest_low) / (highest_high - lowest_low)
  fast_d = fast_k.rolling(smoothing).mean()
  slow_k = fast_d
  slow_d = slow_k.rolling(period_d).mean()
  return slow_k, slow_d


def MACD(df, fast_length, slow_length, smoothing):
  fast = EMA(df, fast_length)
  slow = EMA(df, slow_length)
  macd = fast - slow
  signal = macd.ewm(span=smoothing).mean()
  histogram = macd - signal
  return macd, signal, histogram


def TrendAnalysis(series, smoothing):
  # Get differences, smoothing them out to avoid insignificant "wiggles"
  change = series.diff().rolling(window=smoothing).mean()

  # Find the direction changes and get the current direction count
  signs = list(reversed(np.sign(change).values))
  direction = signs[0]  # current direction
  periods = 1
  for sign in signs[1:]:
    if sign != direction: break
    periods += 1

  # Sum the total change seen for the number of periods
  total_change = abs(sum(list(series.values)[:periods]))

  return direction, periods, total_change


def EnergyOfTrend(df):
  df['SMA50'] = SMA(df, 50)
  return TrendAnalysis(df['SMA50'], 1)


def NameDirection(direction):
  if direction < 0:
    return 'down'
  return 'up'


def FilterOutTrend(direction, args):
  """Determines whether to filter out this Trend value based on args."""
  if not (args.trend_up or args.trend_down):
    return False
  return not getattr(args, 'trend_' + NameDirection(direction))


def ReportEnergyOfTrend(direction, count, total_change):
  arrow = {'up': colors.PaintGreen(UP_ARROW),
           'down': colors.PaintRed(DOWN_ARROW),
          }[NameDirection(direction)]
  return '{} [t={}d]'.format(arrow, count)


def EnergyOfMomentum(df):
  df['MACD'], _, _ = MACD(df, 12, 26, 9)
  return df


def ReportEnergyOfMomentum(df):
  macd_val = df['MACD'].iloc[-1]
  direction, count, _ = TrendAnalysis(df['MACD'], 2)
  arrow = UP_ARROW
  if direction < 0:
    arrow = DOWN_ARROW
  macd = colors.PaintGreen('{:.2f}'.format(macd_val))
  if macd_val < 0:
    macd = colors.PaintRed('{:.2f}'.format(macd_val))
  return '{} [macd={}]'.format(arrow, macd)


def EnergyOfCycle(df):
  df['StochK'], _ = Stoch(df, 5, 3, 2)
  direction, _, _ = TrendAnalysis(df['StochK'], 2)
  return direction, df['StochK'].iloc[-1]


def GradeEnergyOfCycle(k):
  if k >= 80:
    return 'high'
  elif k > 20:
    return 'medium'
  elif k <= 20:
    return 'low'


def ReportEnergyOfCycle(direction, k):
  arrow = UP_ARROW
  if direction < 0:
    arrow = DOWN_ARROW
  color = {'low': colors.GREEN,
           'medium': colors.YELLOW,
           'high': colors.RED}[GradeEnergyOfCycle(k)]
  return '{} [k={}{:.1f}{}]'.format(arrow, color, k, colors.RESET)


def FilterOutCycle(direction, k, args):
  """Determines whether to filter out this Cycle value based on args."""
  filter_level = filter_dir = False
  if any((args.cycle_low, args.cycle_medium, args.cycle_high)):
    filter_level = not getattr(args, 'cycle_' + GradeEnergyOfCycle(k))
  if any((args.cycle_up, args.cycle_down)):
    filter_dir = not getattr(args, 'cycle_' + NameDirection(direction))
  return filter_level or filter_dir


def EnergyOfScale(orig_df, timeframe='W'):
  # Resample data to get longer time frame
  orig_df.index = pd.to_datetime(orig_df.index)
  df = orig_df.resample(timeframe).agg(
                            {OPEN:   lambda x: x[0],  # take first
                             HIGH:   'max',
                             LOW:    'min',
                             CLOSE:  lambda x: x[-1], # take last
                             VOLUME: 'sum'},
                        loffset=pd.offsets.timedelta(days=-6))  # to put the labels to Monday

  # Get the MACD of the new timeframe
  df['MACD'], _, _ = MACD(df, 12, 26, 9)

  return TrendAnalysis(df['MACD'], 2)


def ReportEnergyOfScale(direction, count, total_change):
  arrow = {'up': colors.PaintGreen(UP_ARROW),
           'down': colors.PaintRed(DOWN_ARROW)}[NameDirection(direction)]
  return '{} [t={:d}w]'.format(arrow, count)


def FilterOutScale(direction, args):
  """Determines whether to filter out this Scale value based on args."""
  if not (args.scale_up or args.scale_down):
    return False
  return not getattr(args, 'scale_' + NameDirection(direction))


def Volatility(df):
  return (df[CLOSE].pct_change().tail(5).std() * 100,
          df[CLOSE].pct_change().tail(21).std() * 100)


def ReportVolatility(wVol, mVol):
  arrow = UP_ARROW
  if wVol < mVol:
    arrow = DOWN_ARROW
  color = colors.GREEN
  if wVol > 3 and wVol < 5:
    color = colors.YELLOW
  elif wVol >= 5:
    color = colors.RED
  return '{} [w={}{:.2f}%{}]'.format(arrow, color, wVol, colors.RESET)


def CalculatePivotPoints(incoming_df):
  this_month = int(incoming_df.tail(1).index.month[0])
  year = int(incoming_df.tail(1).index.year[0])
  last_month = ((this_month - 2) % 12) + 1
  if last_month > this_month:
    year = year - 1

  df = incoming_df.loc[(incoming_df.index.month == last_month) & (incoming_df.index.year == year)]
  low = float(df[LOW].min())
  high = float(df[HIGH].max())
  close = float(df.tail(1)[CLOSE][0])
  pp = (high + low + close) / 3
  r1 = 2*pp - low
  s1 = 2*pp - high
  r2 = pp + (high - low)
  s2 = pp - (high - low)
  r3 = high + 2*(pp - low)
  s3 = low - 2*(high - pp)

  last_close = float(incoming_df.tail(1)[CLOSE])
  p = PivotPoints()
  p.pp, p.r1, p.s1, p.r2, p.s2, p.r3, p.s3 = pp, r1, s1, r2, s2, r3, s3
  return p, last_close


def ReportPivotPoints(ppoints, close):
  # Pivot points are ordered and colored statically
  strs = [colors.PaintRed('{:.2f}'.format(ppoints.s3)),
          colors.PaintLightRed('{:.2f}'.format(ppoints.s2)),
          colors.PaintYellow('{:.2f}'.format(ppoints.s1)),
          '{:.2f}'.format(ppoints.pp),
          colors.PaintCyan('{:.2f}'.format(ppoints.r1)),
          colors.PaintBlue('{:.2f}'.format(ppoints.r2)),
          colors.PaintMagenta('{:.2f}'.format(ppoints.r3))]


  # colors the 'close' value according to its place in the index
  close_colorizer = {0: colors.PaintBackgroundRed,
                     1: colors.PaintBackgroundRed,
                     2: colors.PaintBackgroundYellow,
                     3: colors.PaintBlackOnWhite,
                     4: colors.PaintBlackOnWhite,
                     5: colors.PaintBackgroundCyan,
                     6: colors.PaintBackgroundBlue,
                     7: colors.PaintBackgroundMagenta}

  # Find the index where the 'close' value belongs so we can insert it with the other strings
  vals = [ppoints.r3, ppoints.r2, ppoints.r1, ppoints.pp, ppoints.s1, ppoints.s2, ppoints.s3, close]
  vals.sort()
  close_index = vals.index(close)

  # Finally, color and insert the 'close' value appropriately
  strs.insert(close_index, close_colorizer[close_index]('{:.2f}'.format(close)))

  return '[{} {} {} {} {} {} {} {}]'.format(*tuple(strs))


def main(args):
  headers = ['TICKER', 'TREND', 'MOMENTUM', 'CYCLE', 'SCALE', 'VOLATILITY', 'PIVOT POINTS']
  widths = [6, 10, 14, 10, 9, 11, 13]
  printer = printing.TabularPrinter(headers=headers, widths=widths)

  def rows():
    for ticker in args.tickers:
      df = fetcher.DataFetcher().FetchData(ticker)

      # Trend
      trend_args = EnergyOfTrend(df)
      if FilterOutTrend(trend_args[0], args):
        continue
      trend = ReportEnergyOfTrend(*trend_args)

      # Momentum
      momentum = ReportEnergyOfMomentum(EnergyOfMomentum(df))

      # Cycle
      cycle_args = EnergyOfCycle(df)
      if FilterOutCycle(*cycle_args, args):
        continue
      cycle_report = ReportEnergyOfCycle(*cycle_args)

      # Scale
      scale_args = EnergyOfScale(df)
      if FilterOutScale(scale_args[0], args):
        continue
      scale = ReportEnergyOfScale(*scale_args)

      volatility = ReportVolatility(*Volatility(df))
      pivotpoints = ReportPivotPoints(*CalculatePivotPoints(df))
      yield (colors.PaintCyan(ticker),
             trend,
             momentum,
             cycle_report,
             scale,
             volatility,
             pivotpoints)
  printer.print(rows())


if __name__ == '__main__':
  parser = argparse.ArgumentParser()
  parser.add_argument('tickers', nargs='+')
  parser.add_argument('--cycle-low', action='store_true', help='Filter for tickers in the cycle low range.')
  parser.add_argument('--cycle-medium', action='store_true', help='Filter for tickers in the cycle medium range.')
  parser.add_argument('--cycle-high', action='store_true', help='Filter for tickers in the cycle high range.')
  parser.add_argument('--cycle-up', action='store_true', help='Filter for tickers where Cycle is going up.')
  parser.add_argument('--cycle-down', action='store_true', help='Filter for tickers where Cycle is going down.')
  parser.add_argument('--scale-up', action='store_true', help='Filter for tickers who Scale energy is up.')
  parser.add_argument('--scale-down', action='store_true', help='Filter for tickers who Scale energy is down.')
  parser.add_argument('--trend-up', action='store_true', help='Filter for tickers who Trend energy is up.')
  parser.add_argument('--trend-down', action='store_true', help='Filter for tickers who Trend energy is down.')
  args = parser.parse_args()
  main(args)