from AlgorithmImports import *
# endregion

from datetime import timedelta
from QuantConnect import *


class MyAlgorithm(QCAlgorithm):
    
    def Initialize(self):


        self.SetStartDate(2020, 1, 1)  # Set the start date for the backtest
        self.SetEndDate(2020, 3, 1)  # Set the end date for the backtest
        self.SetCash(100000)  # Set an initial cash balance
        # self.mult = 0.04

        self.AddEquity("SPY")  # Add the SPY equity
        # self.option = self.AddOption("SPY", Resolution.Minute)

        # self.option.SetFilter(timedelta(1), timedelta(15))

        self.ohlc_window = RollingWindow[TradeBar](3)

        self.sto14_window = RollingWindow[IndicatorDataPoint](3)
        self.STO("SPY", 14, 3, 3).Updated += (lambda sender, updated: self.sto14_window.Add(updated))
        
        self.sto70_window = RollingWindow[IndicatorDataPoint](3)
        self.STO("SPY", 70, 15, 7).Updated += (lambda sender, updated: self.sto70_window.Add(updated))

        self.adx_window = RollingWindow[IndicatorDataPoint](3)
        self.ADX("SPY", 14, Resolution.Daily).Updated += (lambda sender, updated: self.adx_window.Add(updated))

        self.plusdi_window = RollingWindow[IndicatorDataPoint](3)
        self.ADX("SPY", 14).PositiveDirectionalIndex.Updated += (lambda sender, updated: self.plusdi_window.Add(updated))

        self.minusdi_window = RollingWindow[IndicatorDataPoint](3)
        self.ADX("SPY", 14).NegativeDirectionalIndex.Updated += (lambda sender, updated: self.minusdi_window.Add(updated))

        self.bb_u = RollingWindow[IndicatorDataPoint](3)
        self.BB("SPY", 20, 2.1).UpperBand.Updated += (lambda sender, updated: self.bb_u.Add(updated))

        self.bb_l = RollingWindow[IndicatorDataPoint](3)
        self.BB("SPY", 20, 2.1).LowerBand.Updated += (lambda sender, updated: self.bb_l.Add(updated))


        self.Schedule.On(self.DateRules.EveryDay("SPY"), self.TimeRules.BeforeMarketClose("SPY", 1), self.atm)
        # self.SetWarmUp(30, Resolution.Daily)

        self.SetWarmUp(3, Resolution.Daily)

    def OnData(self, data):
        if data.ContainsKey("SPY"):
            self.ohlc_window.Add(data["SPY"])

    def atm(self):

        if self.ohlc_window.Size < 3:
            return
        if(self.plusdi_window[0]):
            self.Debug(f"low {self.ohlc_window[0].Low}")
            self.Debug(f"close {self.ohlc_window[0].Close}")
            self.Debug(f"sto14 {self.sto14_window[0].Value}")
            self.Debug(f"DI+ {self.plusdi_window[0].Value}")
            self.Debug(f"DI- {self.minusdi_window[0].Value}")
            self.Debug(f"ADX {self.adx_window[0].Value}")