from AlgorithmImports import *
from datetime import datetime, timedelta
import talib as ta
import pytz

symbol = "TSLA"

class BullPivot(QCAlgorithm):
    
    

    def __init__(self):

        self.previous = None
        self.fast = None
        self.slow = None

    def Initialize(self):
        '''Initialise the data and resolution required, as well as the cash and start-end dates for your algorithm. All algorithms must initialized.'''

        self.SetStartDate(2022,10,10)    #Set Start Date
        #self.SetEndDate(2022,3,1)      #Set End Date
        self.SetEndDate(datetime.now())   #Set End Date
        self.SetCash(100000)             #Set Strategy Cash

        
        
        # AddEquity( ticker, resolution, market, fillDataForward, leverage, extendedMarketHours)
        # self.AddEquity("SPY", Resolution.Hour, Market.USA, True, 2, True)
        self.ticker = self.AddEquity(symbol, Resolution.Minute, extendedMarketHours = False)
        self.ticker.SetFeeModel(TDAmeritradeFeeModel())
        # self.SetBrokerageModel(BrokerageName.InteractiveBrokersBrokerage, AccountType.Margin)
        # self.DefaultOrderProperties = InteractiveBrokersOrderProperties()
        # self.DefaultOrderProperties.OutsideRegularTradingHours = True

        self.tickerSymbol = self.ticker.Symbol
        self.SetBenchmark(symbol)
        self.SetWarmUp(10, Resolution.Minute)



    def bullpivot(self, close, open, low):

        bullpivot = ((close-open) >= 0.38) & (close > close.shift(2)) \
        & (close < open.shift(2)) & (open <= close.shift(1))

        return bullpivot

    def bullpivot2(self, close, open, low):
        bullbar2 = ((close - open) >= .38) & (close > close.shift(2)) & (close < open.shift(4)) \
            & ((open.shift(2)-close.shift(2)) > .20) & (low.shift(4) > low.shift(1)) \
                & (low.shift(5) > low.shift(3)) & (low.shift(6) > low.shift(4))
        
        return bullbar2

    def closelong(self, bullpivot, bullpivot2):
        
        cl1 = bullpivot.shift(10)
        cl2 = bullpivot2.shift(10)

        return cl1, cl2

    def runT3(self, data):
    
        data.reset_index(inplace=True)
        open = data['open'] #series
        close = data['close'] #series
        high = data['high'] #series
        low = data['low'] #series

        # check for EOD
        global endofday
        t = data.iloc[-1]
        t = str(t['time'])
        t = t[-8:]
        if t == '15:59:00':
            endofday = True
            #self.Debug("endofday: true")
        else:
            endofday = False
            #self.Debug("endofday: false")

        # check for opening 10 bars
        global tenfromopen
        if t == '09:30:00' or t == '09:31:00' or t == '09:32:00' or t == '09:33:00' or t == '09:34:00' or t == '09:35:00' or t == '09:36:00' or t == '09:37:00' or t == '09:38:00' or t == '09:39:00':
            tenfromopen = False
        else:
            tenfromopen = True

        # Bull Pivot
        bb1 = self.bullpivot(close, open, low)
        bb2 = self.bullpivot2(close, open, low)
        data['bull_pivot'] = bb1
        data['bull_pivot2'] = bb2

        # close long
        cl1, cl2 = self.closelong(bb1, bb2)
        data['close_bullpivot'] = cl1
        data['close_bullpivot2'] = cl2

        return data
    
    def OnData(self, data):
        '''OnData event is the primary entry point for your algorithm. Each new data point will be pumped in here.
        Arguments:
            data: Slice object keyed by symbol containing the stock data
        '''


        if not data.ContainsKey(symbol): return

        if data[symbol] is None:
             self.Log("Price Missing Time: %s"%str(self.Time))
             return


        quantity = self.Portfolio[symbol].Quantity

        # Get OHLC data for last 52 bars
        hist = self.History(self.tickerSymbol, 15, Resolution.Minute)

        #Calculate signals on data
        signals = self.runT3(hist)

        bull_pivot = signals.iloc[-1]['bull_pivot']
        #self.Debug("BullPivot: ", bull_pivot)

        bull_pivot2 = signals.iloc[-1]['bull_pivot2']
        #self.Debug("BullPivot 2: ", bull_pivot2)

        close_long = signals.iloc[-1]['close_bullpivot']
        #self.Debug("Closelong: ", close_long)

        close2 = signals.iloc[-1]['close_bullpivot2']
        #self.Debug("Close2: ", close2)


        # Place orders      

        if quantity == 0 and bull_pivot:
            self.SetHoldings(symbol, 1.0)
            self.Debug("BullPivot 2")
            self.Log("BullPivot 2: %s"%str(self.Time))


        if quantity > 0 and close_long:
            self.Liquidate(symbol)
            self.Debug("Pivot sell ")
            self.Log("Pivot sell: %s"%str(self.Time))
        
        if quantity > 0 and close2:
            self.Liquidate(symbol)
            self.Debug("Pivot2 sell ")
            self.Log("Pivot2 sell: %s"%str(self.Time))