import pandas as pd
import numpy as np


class  MyFrameworkAlgorithm(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2020, 8, 10)  # Set Start Date
        self.SetEndDate(2020, 8, 12)
        self.SetCash(1000)  # Set Strategy Cash
        #self.SetEndDate(2018, 1, 1)
        
        # create a dictionary to store momentum indicators for all symbols 
        self.mom = {}
        self.roc = {}
        self.sto = {}
        self.stoExit = {}
        self.rsiEntry = {}
        self.bbEntry = {}
        self.smatwenty = {}
        self.adosc = {}
        self.psar = {}
        self.eveningstar = {}
        self.threeblackcrows = {}
        self.darkcloudcover = {}
        self.obv = {}
        self.ematwenty = {}
        self.emafifty = {}
        self.adx= {}
        
        self.Price_list = {} 
        self.ROCx_list = {}
        self.OBV_list = {}
        self.KminusD_list = {}
        
        #Warmup period 
        period = 200
        #MOM period
        period1 = 20
        
        self.SetRiskManagement(TrailingStopRiskManagementModel(0.22))
        
        # warm up the MOM indicator
        self.SetWarmUp(period)
        self.symbols = []
        self.tickers = ["TQQQ", #ProShares UltraPro QQQ
                        "SQQQ", #ProShares UltraPro Short QQQ
                        "UBT",  #ProShares Ultra 20+ Year Treasury 2x
                        "TECL", #Direxion Daily Technology Bull 3X
                        "TECS", #Direxion Daily Technology Bear 3X
                        "UGLDF",#VelocityShares 3x Long Gold ETN
                        "GLL",  #ProShares UltraShort Gold (GLL)
                        "TYD",  #Direxion Daily 7-10 Year Treasury Bull 3X
                        #"TYO", #Direxion Daily 7-10 Year Treasury Bear 3X
                        "DRIP", #Direxion Daily S&P Oil & Gas Exp. & Prod. Bear 2X
                        "TMF",  #Direxion Daily 20+ Year Treasury Bull 3X
                        #"TMV",  #Direxion Daily 20-Year Treasury Bear 3X
                        "SOXL", #Direxion Daily Semiconductor Bull 3X
                        #"SOXS", #Direxion Daily Semiconductor Bear 3X
                        "WEBL", #Direxion Daily Dow Jones Internet Bull 3X
                        "CWEB", #Direxion Daily CSI China Internet Bull 2X Shares (CWEB)
                        "CURE", #Direxion Daily Healthcare Bull 3X Shares
                        "UYG",  #ProShares Ultra Financials (UYG)
                        "XLF",  #Financial Select Sector SPDR® Fund (XLF)
                        "EWZ",  #iShares MSCI Brazil Capped ETF 
                        "EFA",  #iShares MSCI EAFE ETF 
                        "AMD", "SQ", "NVDA", "QCOM", "VZ", "JPM", "T", "AVGO"]
    
        
        for ticker in self.tickers:
            symbol = self.AddEquity(ticker, Resolution.Minute).Symbol
            self.symbols.append(symbol)
            self.mom[symbol] = self.MOM(symbol, period1, Resolution.Daily)
            self.roc[symbol] = self.ROC(symbol, 15, Resolution.Daily)
            self.sto[symbol] = self.STO(symbol, 21, 5, 5, Resolution.Daily)
            self.stoExit[symbol] = self.STO(symbol, 21, 5, 5, Resolution.Daily)
            self.rsiEntry[symbol] = self.RSI(symbol, 20, Resolution.Daily)
            self.bbEntry[symbol] = self.BB(symbol, 20, 2, Resolution.Daily)
            self.smatwenty[symbol] = self.SMA(symbol, 20, Resolution.Daily)
            self.ematwenty[symbol] = self.EMA(symbol, 14, Resolution.Daily)
            self.emafifty[symbol] = self.EMA(symbol, 50, Resolution.Daily)
            self.adosc[symbol] = self.ADOSC(symbol, 3, 10, Resolution.Daily)  #Chaikin Oscillator
            self.psar[symbol]  = self.PSAR(symbol, 0.02, 0.02, 0.2, Resolution.Daily)
            self.eveningstar[symbol] = self.CandlestickPatterns.EveningStar(symbol, 0.3, Resolution.Daily)
            self.threeblackcrows[symbol] = self.CandlestickPatterns.ThreeBlackCrows(symbol, Resolution.Daily)
            self.darkcloudcover[symbol] = self.CandlestickPatterns.DarkCloudCover(symbol, 0.80, Resolution.Daily)
            self.obv[symbol] = self.OBV(symbol, Resolution.Daily)
            self.adx[symbol] = self.ADX( symbol, 20, Resolution.Daily)
            
        self.Schedule.On(self.DateRules.Every(DayOfWeek.Wednesday), \
                 self.TimeRules.At(9, 31), \
                 self.Rebalance)
                
        self.Schedule.On(self.DateRules.Every(DayOfWeek.Thursday, DayOfWeek.Friday, DayOfWeek.Monday, DayOfWeek.Tuesday), \
                 self.TimeRules.At(9, 31), \
                 self.ExitPlan)
        #self.Schedule.On(self.DateRules.EveryDay(), \
        #self.TimeRules.At(10, 30), \
        #self.Rebalance)
                
    def Rebalance(self):
        if self.IsWarmingUp: return
    
        Enterdata = pd.DataFrame(list(zip(self.tickers)), columns =['tickers'])
        Enterdata.insert(1, "Price", [self.Portfolio[symbol].Price for symbol in self.symbols], True)
        Enterdata.insert(2, "mom", [self.mom[symbol].Current.Value for symbol in self.symbols], True)
        Enterdata.insert(3, "rsi", [self.rsiEntry[symbol].Current.Value for symbol in self.symbols], True)
        Enterdata.insert(4, "rsiAverageGain", [self.rsiEntry[symbol].AverageGain.Current.Value for symbol in self.symbols], True)
        Enterdata.insert(5, "rsiAverageLoss", [self.rsiEntry[symbol].AverageLoss.Current.Value for symbol in self.symbols], True)
        Enterdata.insert(6, "DStoch", [self.sto[symbol].StochD.Current.Value for symbol in self.symbols], True)
        Enterdata.insert(7, "KStoch", [self.sto[symbol].StochK.Current.Value for symbol in self.symbols], True)
        Enterdata.insert(8, "MiddleBand", [self.bbEntry[symbol].MiddleBand.Current.Value for symbol in self.symbols], True)
        Enterdata.insert(9, "UpperBand", [self.bbEntry[symbol].UpperBand.Current.Value for symbol in self.symbols], True)
        Enterdata.insert(10, "LowerBand", [self.bbEntry[symbol].LowerBand.Current.Value for symbol in self.symbols], True)
        Enterdata.insert(11, "sma20", [self.smatwenty[symbol].Current.Value for symbol in self.symbols], True)
        Enterdata.insert(12, "ema20", [self.ematwenty[symbol].Current.Value for symbol in self.symbols], True)
        Enterdata.insert(13, "ema50", [self.emafifty[symbol].Current.Value for symbol in self.symbols], True)
        Enterdata.insert(14, "adosc", [self.adosc[symbol].Current.Value for symbol in self.symbols], True)
        Enterdata.insert(15, "psar", [self.psar[symbol].Current.Value for symbol in self.symbols], True)
        Enterdata.insert(16, "OBV", [self.obv[symbol].Current.Value for symbol in self.symbols], True)
        Enterdata.insert(17, "ADX", [self.adx[symbol].Current.Value for symbol in self.symbols], True)
        Enterdata.insert(18, "eveningstar", [self.eveningstar[symbol].Current.Value for symbol in self.symbols], True)
        Enterdata.insert(19, "blackcrows", [self.threeblackcrows[symbol].Current.Value for symbol in self.symbols], True)
        Enterdata.insert(20, "darkcloud", [self.darkcloudcover[symbol].Current.Value for symbol in self.symbols], True)
        Enterdata.insert(21, "ROC", [self.roc[symbol].Current.Value for symbol in self.symbols], True)
        # using apply function to create a new column 
        Enterdata['ROCx'] = Enterdata.apply(lambda row: row.ROC * 10, axis = 1)
        Enterdata['KminusD'] = Enterdata.apply(lambda row: row.KStoch - row.DStoch, axis = 1)
     
        #alldata = pd.DataFrame({
                            # 'mom':[self.mom[symbol].Current.Value for symbol in self.symbols],
                            # 'rsi':[self.rsi[symbol].Current.Value for symbol in self.symbols]})
       
        star = Enterdata[Enterdata['eveningstar'] != -1]
        blackcrows = star[star['blackcrows'] != -1]
        darkcloud = blackcrows[blackcrows['darkcloud'] != -1]
        ados = darkcloud[darkcloud['adosc'] > 0]
        #sma20 = ados[ados['Price'] > ados['sma20']]
        ema20 = ados[ados['Price'] > ados['ema20']]
        #emacross = ados[ados['ema20'] > ados['ema50']]
        lowerbb =  ema20[ema20['Price'] > ema20['LowerBand']]
        psar = lowerbb[lowerbb['Price'] > lowerbb['psar']]
        adx = psar[psar['ADX'] > 5] #basically doing nothing
        stoch = adx[adx['DStoch'] < adx['KStoch']]
        stoch2 = stoch[stoch['KStoch'] < 95]
        lowest_rsi = stoch2[stoch2['rsi'] < 95]
        positiveROC = lowest_rsi[(lowest_rsi['ROCx']) >= 0]
        top3 = positiveROC.sort_values(by='mom', ascending = False)[:4]
        top3_symbols = [self.symbols[i] for i in top3.index]
        
        
        self.Price_list = dict([(i,a) for i, a in zip(Enterdata.tickers, Enterdata.Price)])
        self.ROCx_list = dict([(i,a) for i, a in zip(Enterdata.tickers, Enterdata.ROCx)])
        self.OBV_list = dict([(i,a) for i, a in zip(Enterdata.tickers, Enterdata.OBV)])
        self.KminusD_list = dict([(i,a) for i, a in zip(Enterdata.tickers, Enterdata.KminusD)])
        
        #self.Log(top3.index)
        #self.Log(self.Price_list)
    
        #TECLprice = self.Portfolio["TECL"].Price 
        #self.Plot("stock", "TECL", TECLprice)
        #UGLDprice = self.Portfolio["UGLD"].Price 
        #self.Plot("stock", "UGLD", UGLDprice)
        
        #UBTmom = alldata.iloc[1]['mom']
        #self.Plot("mom", "UBTmom", UBTmom)
        #TECLmom = alldata.iloc[2]['mom']
        #self.Plot("mom", "TECLmom", TECLmom)
        #UGLDmom = alldata.iloc[3]['mom']
        #self.Plot("mom", "UGLDmom", UGLDmom)
        
        #UBTrsi = alldata.iloc[1]['rsi']
        #self.Plot("rsi", "UBTrsi", UBTrsi)
        #TECLrsi = alldata.iloc[2]['rsi']
        #self.Plot("rsi", "TECLrsi", TECLrsi)
        #UGLDrsi = alldata.iloc[3]['rsi']
        #self.Plot("rsi", "UGLDrsi", UGLDrsi)
        
        #TQQQprice = self.Portfolio["TQQQ"].Price 
        #self.Plot("stock", "TQQQ", TQQQprice)
        #TQQQKStoch = Enterdata.iloc[0]['KStoch']
        #self.Plot("TQQQStoch", "TQQQKStoch", TQQQKStoch)
        #TQQQDStoch = Enterdata.iloc[0]['DStoch']
        #self.Plot("TQQQStoch", "TQQQDStoch", TQQQDStoch)
        #TQQQmom = Enterdata.iloc[0]['mom']
        #self.Plot("TQQQmom", "TQQQmom", TQQQmom)
        #TQQQrocx = (Enterdata.iloc[0]['ROCx'])
        #self.Plot("TQQQrocx", "TQQQrocx", TQQQrocx)
        #TQQQrsi = Enterdata.iloc[0]['rsi']
        #self.Plot("TQQQrsi", "TQQQrsi", TQQQrsi)
        #TQQQadosc = Enterdata.iloc[0]['adosc']
        #self.Plot("TQQQadosc", "TQQQadosc", TQQQadosc)
        #TQQQobv = (Enterdata.iloc[0]['OBV'])
        #self.Plot("TQQQobv", "TQQQobv", TQQQobv)
        #TQQQadx = (Enterdata.iloc[0]['ADX'])
        #self.Plot("TQQQadx", "TQQQadx", TQQQadx)
        
        
        self.Log(Enterdata)
        #self.Log(alldata.dtypes)
        #self.Log(str(alldata.shape))
        #self.Log(str(alldata.index.name))
        #self.Log(self.tickers)
        #self.Log(tickr)
        #self.Log(tickr.dtypes)
        #self.Log(TQQQmom)
       
        
        for kvp in self.Portfolio:
            security_hold = kvp.Value
            if security_hold.Invested and (security_hold.Symbol not in top3_symbols):
                self.Liquidate(security_hold.Symbol)
        
        added_symbols = []        
        for symbol in top3_symbols:
            if not self.Portfolio[symbol].Invested:
                added_symbols.append(symbol)
        for added in added_symbols:
            self.SetHoldings(added, 1/len(added_symbols))
            self.Log(added_symbols)
            
        return
            
    def ExitPlan(self):
        if self.IsWarmingUp: return
    
        Exitdata = pd.DataFrame(list(zip(self.tickers)), columns =['tickers'])
        Exitdata.insert(1, "Price", [self.Portfolio[symbol].Price for symbol in self.symbols], True)
        Exitdata.insert(2, "mom", [self.mom[symbol].Current.Value for symbol in self.symbols], True)
        Exitdata.insert(3, "rsi", [self.rsiEntry[symbol].Current.Value for symbol in self.symbols], True)
        Exitdata.insert(4, "rsiAverageGain", [self.rsiEntry[symbol].AverageGain.Current.Value for symbol in self.symbols], True)
        Exitdata.insert(5, "rsiAverageLoss", [self.rsiEntry[symbol].AverageLoss.Current.Value for symbol in self.symbols], True)
        Exitdata.insert(6, "DStoch", [self.stoExit[symbol].StochD.Current.Value for symbol in self.symbols], True)
        Exitdata.insert(7, "KStoch", [self.stoExit[symbol].StochK.Current.Value for symbol in self.symbols], True)
        Exitdata.insert(8, "MiddleBand", [self.bbEntry[symbol].MiddleBand.Current.Value for symbol in self.symbols], True)
        Exitdata.insert(9, "UpperBand", [self.bbEntry[symbol].UpperBand.Current.Value for symbol in self.symbols], True)
        Exitdata.insert(10, "LowerBand", [self.bbEntry[symbol].LowerBand.Current.Value for symbol in self.symbols], True)
        Exitdata.insert(11, "sma20", [self.smatwenty[symbol].Current.Value for symbol in self.symbols], True)
        Exitdata.insert(12, "ema20", [self.ematwenty[symbol].Current.Value for symbol in self.symbols], True)
        Exitdata.insert(13, "ema50", [self.emafifty[symbol].Current.Value for symbol in self.symbols], True)
        Exitdata.insert(14, "adosc", [self.adosc[symbol].Current.Value for symbol in self.symbols], True)
        Exitdata.insert(15, "psar", [self.psar[symbol].Current.Value for symbol in self.symbols], True)
        Exitdata.insert(16, "OBV", [self.obv[symbol].Current.Value for symbol in self.symbols], True)
        Exitdata.insert(17, "eveningstar", [self.eveningstar[symbol].Current.Value for symbol in self.symbols], True)
        Exitdata.insert(18, "blackcrows", [self.threeblackcrows[symbol].Current.Value for symbol in self.symbols], True)
        Exitdata.insert(19, "darkcloud", [self.darkcloudcover[symbol].Current.Value for symbol in self.symbols], True)
        Exitdata.insert(20, "ROC", [self.roc[symbol].Current.Value for symbol in self.symbols], True)
        Exitdata['ROCx'] = Exitdata.apply(lambda row: row.ROC * 10, axis = 1)
        Exitdata['ROC2Prev']= Exitdata['tickers'].map(self.ROCx_list)
        Exitdata['KminusD'] = Exitdata.apply(lambda row: row.KStoch - row.DStoch, axis = 1)
        Exitdata['KminusDPrev']= Exitdata['tickers'].map(self.KminusD_list)
        Exitdata['PricePrev']= Exitdata['tickers'].map(self.Price_list)
        # using apply function to create a new column 
        Exitdata['OBV2']= Exitdata['tickers'].map(self.OBV_list)
        
        self.Log(Exitdata)
        TQQQprice = self.Portfolio["TQQQ"].Price 
        self.Plot("stock", "TQQQ", TQQQprice)
        TQQQKminusD2 = (Exitdata.iloc[0]['KminusD'])
        self.Plot("TQQQKminusD", "TQQQKminusD2", TQQQKminusD2)
        TQQQKminusD = (Exitdata.iloc[0]['KminusD'])
        self.Plot("TQQQKminusD", "TQQQKminusD", TQQQKminusD)
        
        #TotExit = Exitdata[(Exitdata['eveningstar'] == -1) & (Exitdata['Price'] < Exitdata['psar'])]
        #rsiExit = Exitdata[(Exitdata['rsi'] >= 80)] #13.0
        #BBExit = Exitdata[(Exitdata['Price'] < Exitdata['LowerBand'])] #12.5
        #psarExit = Exitdata[Exitdata['Price'] < Exitdata['psar']] #10.2
        #adosExit = Exitdata[Exitdata['adosc'] < 0] #13.0
        #StarExit = Exitdata[Exitdata['eveningstar'] == -1] #9.8
        #PriceExit = Exitdata[Exitdata['Price'] > Exitdata['Price2']]
        #PriceExit = Exitdata[Exitdata['mom'] < -20]
        #StochExit = Exitdata[Exitdata['DStoch'] > Exitdata['KStoch']]
        #SMAExit =  Exitdata[Exitdata['sma20'] > Exitdata['Price']]
        EMAExit =  Exitdata[Exitdata['ema20'] > 1.03*Exitdata['Price']]
        #ROCExit =  Exitdata[Exitdata['ROCx'] < Exitdata['ROC2']]
        #OBVExit =  Exitdata[Exitdata['OBV2']*1.5 > Exitdata['OBV']]
        
        #self.Log(psarExit)
        #self.Log(BBExit)
        #self.Log(psarExit) 
        
        Exit_symbols = [self.symbols[i] for i in EMAExit.index]
        
        for kvp in self.Portfolio:
            security_hold = kvp.Value
            if security_hold.Invested and (security_hold.Symbol in Exit_symbols):
                self.Liquidate(security_hold.Symbol)