f = False
if f:
    from AlgorithmImports import *
from collections import deque
from typing import List
import configs as cfg
from indicators import GoldenCross, ATRBuySell

from datetime import timedelta

import pickle

class EnergeticBlueDonkey(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2017, 6, 1)
        self.SetEndDate(2017, 6, 10)
        self.SetCash(100000)
        self.symbol = self.AddEquity('SPY', Resolution.Minute).Symbol
        self.gc = GoldenCross(cfg.fast_sma_period, cfg.slow_sma_period)
        self.atrbs = ATRBuySell()

        self.Consolidate(self.symbol, Resolution.Daily, self.OnConsolidation)

        option = self.AddOption(self.symbol)
        option.SetFilter(-20, +20, timedelta(15), timedelta(45))
        option.PriceModel = OptionPriceModels.CrankNicolsonFD()  # necessary for greeks

        self.options : List[Symbol] = []

        # makes it so we emulate Daily Resolution on Minute Resolution
        # this is necessary since we are dealing with options, which only work on Minute or finer data
        self.curr_day = -1 

        self.SetWarmUp(max(cfg.slow_sma_period + 3, cfg.atr_period), Resolution.Daily)

        # stuff for data recording
        self.orders = []
        self.bb = BollingerBands(cfg.bb_period, cfg.bb_multiple)
        self.kelt = KeltnerChannels(cfg.kelt_period, cfg.kelt_multiple)
        
        indicators = ['time', 'price', 'sma_short', 'sma_long', 'atr', 'obv', 'bb', 'bb_lower', 'bb_upper', 'kelt', 'kelt_lower', 'kelt_upper']
        self.indicator_values = {indicator: list() for indicator in indicators}

        self.bars = []

    def OnEndOfAlgorithm(self):
        self.ObjectStore.SaveBytes('orders', pickle.dumps(self.orders))
        self.ObjectStore.SaveBytes('indicators', pickle.dumps(self.indicator_values))
        self.ObjectStore.SaveBytes('bars', pickle.dumps(self.bars))

    def OnConsolidation(self, bar: TradeBar):
        self.atrbs.Update(bar)
        
        self.bb.Update(bar.EndTime, bar.Close)
        self.kelt.Update(bar)

        self.RecordValues(bar)

    def RecordValues(self, bar: TradeBar):
        if self.IsWarmingUp:
            return
        
        self.bars.append({
            'time': bar.EndTime,
            'open': bar.Open, 
            'high': bar.High, 
            'low': bar.Low, 
            'close': bar.Close
        })
        self.indicator_values['time'].append(bar.EndTime)
        self.indicator_values['price'].append(bar.Close)
        self.indicator_values['obv'].append(bar.Volume)
        self.indicator_values['sma_short'].append(self.gc.fast_sma.Current.Value)
        self.indicator_values['sma_long'].append(self.gc.slow_sma.Current.Value)
        self.indicator_values['atr'].append(self.atrbs.atr.Current.Value)
        self.indicator_values['bb'].append(self.bb.Current.Value)
        self.indicator_values['bb_lower'].append(self.bb.LowerBand.Current.Value)
        self.indicator_values['bb_upper'].append(self.bb.UpperBand.Current.Value)
        self.indicator_values['kelt'].append(self.kelt.Current.Value)
        self.indicator_values['kelt_lower'].append(self.kelt.LowerBand.Current.Value)
        self.indicator_values['kelt_upper'].append(self.kelt.UpperBand.Current.Value)

    def RecordTrade(self, direction):
        self.orders.append({'time':self.Time, 'direction': direction})

    def Print(self, msg:str):
        '''
        Just Debug, but with a check that debugging is enabled
        '''
        if cfg.debug:
            self.Debug(msg)

    def OnData(self, data:Slice):
        if self.curr_day == self.Time.day:
            return
        self.curr_day = self.Time.day

        self.ProcessOptions(data)

        if not data.Bars.ContainsKey(self.symbol):
            return

        self.gc.Update(data[self.symbol])
        
        if self.IsWarmingUp or not self.gc.IsReady or not self.atrbs.IsReady:
            return

        self.Rebalance()

        self.Plots()

    def Rebalance(self):
        '''
        take profit, stop loss
        GoldenCross entry and DeathCross exit
        equity/option weightage rebalance
        '''
        self.Print('Rebalancing...')
        if self.atrbs.Value == 1:
            self.atrbs.Reset()
            self.Print('TakeProfit - Liquidating')
            self.Liquidate()
            self.RecordTrade(0)
        elif self.atrbs.Value == -1:
            self.atrbs.Reset()
            self.Print('StopLoss - Liquidating')
            self.Liquidate()
            self.RecordTrade(2)
        elif self.gc.Value == 0:
            self.Print('DeathCross - Liquidating')
            self.Liquidate()
            self.RecordTrade(3)
        elif self.gc.Value == 2 and not self.Portfolio.Invested:
            self.Print('GoldenCross - Going Long')
            self.atrbs.SetLevels(1)
            self.SetHoldings(self.symbol, .5)
            for option in self.options:
                self.SetHoldings(option, .5/len(self.options))
            self.RecordTrade('GC')
        elif self.Portfolio.Invested:
            equity_value = self.Portfolio[self.symbol].Price * self.Portfolio[self.symbol].Quantity
            equity_portfolio_pct = equity_value / self.Portfolio.TotalPortfolioValue

            if .45 < equity_portfolio_pct < .65:
                return

            self.Print('Equity/Option Imbalance, resetting to 50/50')

            invested_options = [
                kvp.Key for kvp in self.Portfolio 
                if kvp.Value.Invested and kvp.Key.SecurityType == SecurityType.Option
            ]


            self.SetHoldings(self.symbol, .5)
            for option in invested_options:
                self.SetHoldings(option, .5/len(invested_options))

    def ProcessOptions(self, data:Slice):
        self.Print('Processing options...')
        
        self.options = []
        for x in data.OptionChains:
            chain: List[OptionContract] = [x for x in x.Value]
            
            contracts = [optionContract for optionContract in chain if cfg.option_filter(optionContract)]
          
            
            self.Print(len(contracts))
            
            underlying_price = self.Securities[self.symbol].Price
            nearTheMoney = sorted(contracts, key=lambda contract: abs(contract.Strike - underlying_price))[:cfg.options_count]
            self.options.extend([contract.Symbol for contract in nearTheMoney])

        self.Print(f'Found {len(self.options)} options')

    def Plots(self):
        if not cfg.debug:
            return
        self.Plot('GoldenCross', 'Value', self.gc.Value)

f = False
if f:
    from AlgorithmImports import *

import configs as cfg
from collections import deque

class GoldenCross:
    def __init__(self, fast_period:int, slow_period:int):
        '''
        GoldenCross indicator
        .Value = 0 -> not golden cross or death cross
        .Value = 1 -> golden cross formed, entry not
        .Value = 2 -> entry formed after golden cross
        '''
        self.Value = 0

        self.fast_sma = SimpleMovingAverage(fast_period)
        self.slow_sma = SimpleMovingAverage(slow_period)

        # fast sma - slow sma
        self.sma_diffs = deque(maxlen=3)

    def dq_rdy(self, vals:deque):
        '''
        returns True iff the deque is has maxlen elements
        '''
        return len(vals) == vals.maxlen

    def Update(self, input:TradeBar):
        '''
        updates the Golden Cross indicator with a new bar of data
        returns self.IsReady
        '''
        self.Time = input.EndTime
        close = input.Close

        self.fast_sma.Update(self.Time, close)
        self.slow_sma.Update(self.Time, close)
        
        if not self.slow_sma.IsReady:
            # since the slow_sma takes more values, if its ready
            # the fast_sma must be ready
            return False

        self.sma_diffs.append(
            self.fast_sma.Current.Value - self.slow_sma.Current.Value  
        )
        
        if not self.dq_rdy(self.sma_diffs):
            return False 
        
        is_crossed = (
            self.sma_diffs[2] > 0 and self.sma_diffs[1] < 0 and self.sma_diffs[0] < 0
        ) # if the fast just recently rises above the slow

        is_death_crossed = (
            self.sma_diffs[2] < 0 and self.sma_diffs[1] > 0 and self.sma_diffs[0] > 0
        ) # if the fast just recently dips above the slow

        if is_death_crossed:
            self.Value = 0
        if self.Value <= 0 and is_crossed:
            self.Value = 1
        elif self.Value == 1 and cfg.entry_condition(close, self.fast_sma.Current.Value, self.slow_sma.Current.Value)   :
            self.Value = 2

        return True
    
    def Warmup(self):
        pass

    @property
    def IsReady(self):
        '''
        returns True iff the indicator is ready to use
        '''
        return self.dq_rdy(self.sma_diffs)

class ATRBuySell:
    def __init__(self):
        '''
        ATR Take Profit and Stop Loss
        .Value = -1 -> stop loss
        .Value = 0 -> neutral
        .Value = 1 -> take profit
        '''
        self.Value = 0

        self.atr = AverageTrueRange(cfg.atr_period)
        self.Close = 0
        self.StopLoss = None
        self.TakeProfit = None

    def Update(self, input:TradeBar):
        '''
        updates the ATRBuySell indicator with a new bar of data
        returns self.IsReady
        '''
        self.Time = input.EndTime
        self.Close = input.Close

        self.atr.Update(input)
        
        if not self.atr.IsReady:
            return False

        if self.TakeProfit and input.Close > self.TakeProfit:
            self.Value = 1
        elif self.StopLoss and input.Close < self.StopLoss:
            self.Value = -1

        return True
    
    def Reset(self):
        self.TakeProfit = None
        self.StopLoss = None
        self.Value = 0

    def SetLevels(self):
        '''
        Sets the TakeProfit and StopLoss levels, which are used to determine .Value
        '''
        self.Value = 0
        atr = self.atr.Current.Value
        self.TakeProfit = self.Close + cfg.atr_take_profit_factor * atr
        self.StopLoss = self.Close - cfg.atr_stop_loss_factor * atr

    def Warmup(self):
        pass

    @property
    def IsReady(self):
        '''
        returns True iff the indicator is ready to use
        '''
        return self.atr.IsReady

f = False
if f: from AlgorithmImports import *

# disable below if you want to reduce logging/plotting
debug = False

#BEGIN GoldenCross configurations
fast_sma_period = 5
slow_sma_period = 20
assert(fast_sma_period < slow_sma_period)

# entry condition after cross has formed
def entry_condition(curr_price:float, fast_sma:float, slow_sma:float)->bool:
    '''
    return True iff entry condition is met 
    '''
    sma_avg = (fast_sma + slow_sma) / 2

    # 4% within average of two SMAs
    return abs(1-(curr_price / sma_avg)) < .04
#END GoldenCross configurations

#BEGIN ATR configs
atr_period = 14
atr_take_profit_factor = 2
atr_stop_loss_factor = 1
#END ATR configs

#BEGIN Options configurations
options_count = 3 # how many of the top delta options

target_delta = .5

def option_filter(optionContract: OptionContract):
    '''
    return True iff contract is a call and Greek conditions are met
    '''
    if optionContract.Symbol.ID.OptionRight != OptionRight.Call:
        return False
    # NEVER TRADES
    # elif optionContract.Greeks.Theta < -.05:
    elif optionContract.Greeks.Theta < -10:
        return False
    # NEVER TRADES
    # elif (  # .25 < delta < 1
    #     optionContract.Greeks.Delta < target_delta * .5 
    #     or optionContract.Greeks.Delta > target_delta * 2
    # ):
    elif (  # .25 < delta < 1
        optionContract.Greeks.Delta < target_delta * .5 
        or optionContract.Greeks.Delta > target_delta * 2
    ):
        return False

    return True
#END Options configurations

#BEGIN plotting configs
bb_period = 14
bb_multiple = 2
kelt_period = 14 
kelt_multiple = 2