from Selection.FundamentalUniverseSelectionModel import FundamentalUniverseSelectionModel

CoarseUniverseSize = 3000
FineUniverseSize = 10
UniverseMinMarketCap= 20e9
UniverseMaxMarketCap = 1500e9
PrimaryExchangeID = ["NAS"]

UniverseSectors = [MorningstarSectorCode.Technology,
                MorningstarSectorCode.BasicMaterials,
                MorningstarSectorCode.ConsumerCyclical,
                MorningstarSectorCode.FinancialServices,
                MorningstarSectorCode.ConsumerDefensive,
                MorningstarSectorCode.Healthcare,
                MorningstarSectorCode.Utilities,
                MorningstarSectorCode.CommunicationServices,
                MorningstarSectorCode.Energy,
                MorningstarSectorCode.Industrials,
                MorningstarSectorCode.RealEstate]

# ------------------------------ Universe Selection Model ---------------------------------

class UniverseSelectionModel(FundamentalUniverseSelectionModel):

    def __init__(self, filterFineData = True, universeSettings = None, securityInitializer = None):
        super().__init__(filterFineData, universeSettings, securityInitializer)
        self.numberOfSymbolsCoarse = CoarseUniverseSize
        self.numberOfSymbolsFine = FineUniverseSize
        self.sorted1BySymbol = {} 
        self.lastMonth = -1

# --------------------------------- Coarse --------------------------------------
    def SelectCoarse(self, algorithm, coarse):
        
        if algorithm.Time.month == self.lastMonth:
            return Universe.Unchanged
        self.lastMonth = algorithm.Time.month
        
        sorted1 = sorted([x for x in coarse if x.HasFundamentalData and x.Volume > 0 and x.Price >0],
                        key = lambda x: x.DollarVolume, reverse=True)[:self.numberOfSymbolsCoarse]
        
        self.sorted1BySymbol = {x.Symbol:x.DollarVolume for x in sorted1}
        
        count = len(self.sorted1BySymbol)
        if count == 0:
            return Universe.Unchanged

        # return the symbol objects our sorted collection
        return list(self.sorted1BySymbol.keys())

# ---------------------------------- Fine ---------------------------------------
    def SelectFine(self, algorithm, fine):
       
        sorted2 = sorted([x for x in fine if x.CompanyReference.PrimaryExchangeID in PrimaryExchangeID
                        and (algorithm.Time - x.SecurityReference.IPODate).days > 30
                        and x.AssetClassification.MorningstarSectorCode in UniverseSectors
                        #and x.ValuationRatios.PERatio < 20
                        and x.MarketCap >= UniverseMinMarketCap
                        and x.MarketCap <= UniverseMaxMarketCap],
                        
                        #and x.ValuationRatios.ForwardPERatio > 0
                        #and x.CompanyReference.CountryId == "USA"
                       key = lambda x: x.MarketCap, reverse = True)
         
        rsiValues = {}               
        for symbol in sorted2:
            stock = symbol.Symbol
            rsi = RelativeStrengthIndex(14)
            history = algorithm.History(stock, 14, Resolution.Daily)
            for tuple in history.loc[stock].itertuples():
                rsi.Update(tuple.Index, tuple.close)
            rsiValue = rsi.Current.Value
            rsiValues[stock] = rsiValue
            
        sorted3 = sorted([x for x in sorted2], key = lambda x: rsiValues[x.Symbol], reverse = False)
        
        # fundametnal
        #  market cap
        #  sales growth .OperationRatios.RevenueGrowth.ThreeMonths
        #  PE compared to sector
        #  profitable
        #  technical 
        #  RSI
        #  PPO
        #  SLOPE
        
        count = len(sorted3)
        if count == 0:
            return Universe.Unchanged
        
        algorithm.Debug(f"Fine:{count}")
        
        #for i in range(10):
        #    algorithm.Debug(f"{i} {sorted3[i].Symbol}  {rsiValues[sorted3[i].Symbol]}")
        
        
        # converts our list of fine to a list of Symbols 
        sorted3BySymbol = [f.Symbol for f in sorted3]
       
        return sorted3BySymbol[:self.numberOfSymbolsFine]

# (c) QIO18
import Universe

class QInOut(QCAlgorithm):

    def Initialize(self):
        
        # ------------------------------------------- Parameters ---------------------------------------
         
        self.InOut = True
        self.AlgoName = "QIO18"
        
        # --- List of assets to buy when in the market (Buy equal weights) FXO QQQJ PSCT
        self.InAssets = {0:["Universe"],
                         1:["QQQ"],
                         2:["SPY"],
                         3:["ARKK"],
                         4:["MSFT","AAPL","FB","GOOG","AMZN"],
                         5:["CRSP","NTLA","EDIT"],
                         6:["FDN"],
                         7:["WDAY","AMAT","TSLA"],
                         8:["AAPL","ADBE","ADI","ADP","ADSK","AEP","ALGN","ALXN","AMAT","AMD","AMGN","AMZN","ANSS","ASML","ATVI","AVGO","BIDU","BIIB","BKNG","CDNS","CDW","CERN","CHKP","CHTR","CMCSA","COST","CPRT","CSCO","CSX","CTAS","CTSH"],
                         9:["LMND","CRWD","ZM","PINS","FVRR","NFLX","MTCH","MASI","DIS","NVDA","ABNB"]  # Fool
        }
        
        # --- List of assets to buy when OUT the market (Buy equal weights)
        self.OutAssets = {1:["TLT"],
                         2:["TMF"],
                         3:["VXX"],
                         4:["TLT", "TMF", "VXX"]
        }
        
        # --- List of IndexEquity
        self.IndexEquities = {1:"QQQ",
                         2:"SPY",
                         3:"ARKK"
        }
        
        # ---------------------------------- Algorithm External Parameters -----------------------------
  
        # Entry minute for the Trade algo after Market Open
        self.MinsAfter = int(self.GetParameter("minsAfter"))
        
        # Select IN and OUT Equity list of assests
        self.InAssetSelector = int(self.GetParameter("inAssetSelector"))        # 0:Monthly selected Universe, 1:QQQ ...
        self.OutAssetSelector = int(self.GetParameter("OutAssetSelector"))      # 1:TLT, 2:TMF ...
        self.IndexSelector = int(self.GetParameter("indexSelector"))
          
        # Indicators' Days input
        self.MacdFastDays = int(self.GetParameter("macdFastDays"))
        self.MacdSlowDays = int(self.GetParameter("macdSlowDays"))
        self.MacdSignalDays = int(self.GetParameter("macdSignalDays"))
        self.RsiDays = int(self.GetParameter("rsiDays"))
        self.RcDays = int(self.GetParameter("rcDays"))
        self.PpoSlowDays = int(self.GetParameter("ppoSlowDays"))
        self.PpoFastDays = int(self.GetParameter("ppoFastDays"))
       
        # Indicators gaps and factors
        self.RsiFactorGap = float(self.GetParameter("rsiFactorGap"))
        self.PpoGap = float(self.GetParameter("ppoGap"))
        self.PpoFactorGap = float(self.GetParameter("ppoFactorGap"))
        self.SlopeGap = float(self.GetParameter("slopeGap"))
        
        self.AtrDays = 18 # int(self.GetParameter("atrDays"))
        self.AtrFastDays = 2 # int(self.GetParameter("atrFastDays"))
        self.AtrSlowDays = 22 # int(self.GetParameter("atrSlowDays"))
        self.AtrGap = 1.3 # float(self.GetParameter("atrGap"))
        
        # MACD factors 
        self.MacdFactorOutGap = float(self.GetParameter("macdFactorOutGap"))
        self.MacdHistOutGap = float(self.GetParameter("macdHistOutGap"))
      
        # Trade time in minutes after market open
        self.MinDaysOut = int(self.GetParameter("daysOut"))
        
        
        # ------------------------------------ For Back Testing ---------------------------------
  
        self.StartYear = int(self.GetParameter("startYear"))
        self.EndYear = int(self.GetParameter("endYear"))
        self.SetStartDate(self.StartYear,1,1)
        self.SetEndDate(self.EndYear,12,31)
        
        self.StartCash = 100000
        self.SetCash(self.StartCash)
        #self.Settings.FreePortfolioValuePercentage = 0.05 # For Live IB

        # -------------------------------------- Internal variables ----------------------------------
        
        #For Plots on Equity Graph
        self.IndexEquityHistory = []                    # now used for QQQ
        self.OutEquityHistory = []                      # now used to TLT
        
        # For statistics
        self.AverageRSI = 0
        self.TotalRSI = 0

        self.TickerTable = {}
        
        self.MyInOutIndicator = 1 # Default Enter the Market
        self.IAMOut = False
        self.DaysOut = 0        # Cooldown days after exiting the Market
        self.TotalDays =0       # Counter for total days
        self.TotalDaysIn = 0    # Counter for days in the MArket
        
        self.ReasonOut= ""      # String to report reason to go out of the Market/InEquity

        
        # ------------------------------- Sellecting Assets  ----------------------------
        
        self.IndexEquityTicker = self.IndexEquities[self.IndexSelector]
        
        self.InAssetList = self.InAssets[self.InAssetSelector]          # IN Asset List
        
        self.SetAlpha(ConstantAlphaModel(InsightType.Price, InsightDirection.Up, timedelta(30)))
        
        if self.InAssetSelector == 0:
            self.UniverseSettings.Resolution = Resolution.Minute
            self.AddUniverseSelection(Universe.UniverseSelectionModel()) 
            self.InAssetList = []    # Will be built in OnSecurities Change
        else: # Manually selected Asset lists
            for asset in self.InAssetList:
                self.AddEquity(asset, Resolution. Minute)
       
        self.SPY = self.AddEquity("SPY", Resolution.Daily).Symbol
        self.Fear = self.AddEquity("VXX", Resolution.Daily).Symbol
        self.IndexEquity = self.AddEquity(self.IndexEquityTicker, Resolution.Minute).Symbol                                      # Key index for In Out decisions
        
        self.doNotEnter = [self.SPY, self.Fear, self.IndexEquity] 
       
        self.OutAssetList = self.OutAssets[self.OutAssetSelector]       # OUT Asset List Consider  "SQQQ","GLD","PSQ","SH","EMTY", "XLU","XLP","TBT","IEF","TLH" "TMF"]
        for asset in self.OutAssetList:
            symbol = self.AddEquity(asset, Resolution.Minute).Symbol
            self.doNotEnter.append(symbol)

            
        # ---------------------------- Add IndexEquity Indicators -------------------------

        rsi = self.RSI(self.IndexEquity, self.RsiDays, MovingAverageType.DoubleExponential, Resolution.Daily)
        rsiSMA = IndicatorExtensions.SMA(rsi,self.RsiDays)
        macd = self.MACD(self.IndexEquity,  self.MacdFastDays,self.MacdSlowDays,self.MacdSignalDays,MovingAverageType.Exponential, Resolution.Daily)
        macdSMA = IndicatorExtensions.SMA(macd,9)
        rc = self.RC(self.IndexEquity,self.RcDays, 2, Resolution.Daily)
        ppo = self.PPO(self.IndexEquity, self.PpoFastDays, self.PpoSlowDays, MovingAverageType.Exponential, Resolution.Daily)
        ppoSMA = IndicatorExtensions.SMA(ppo,self.PpoFastDays)
        
        atr = self.ATR(self.IndexEquity, self.AtrDays, MovingAverageType.Exponential, Resolution.Daily)
        atrSMA_fast = IndicatorExtensions.SMA(atr,self.AtrFastDays)
        atrSMA_slow = IndicatorExtensions.SMA(atr,self.AtrSlowDays)
        symbolData = SymbolData(self.IndexEquity, rsi, rsiSMA, macd, macdSMA, rc, ppo, ppoSMA, atr, atrSMA_fast, atrSMA_slow)
        
        self.TickerTable[self.IndexEquity] = symbolData


        # ------------------------------------ Final setup ----------------------------------

        self.SetBenchmark(self.SPY)
        
        # header data for graph
        str1 = f"{self.AlgoName} {str(self.StartYear)}-{str(self.EndYear)} {self.InAssetSelector}-{self.IndexEquity}-{self.OutAssetSelector}  {self.InOut}"
        str2 = f"{self.RsiFactorGap}/{self.PpoGap}/{self.SlopeGap}/{self.MinDaysOut}/{self.MacdHistOutGap}/{self.MacdFactorOutGap}"
        self.SetRuntimeStatistic(str1, str2)

        self.SetWarmUp(timedelta(days=120)) # Minimum should be RSI period X 2 + 1
        
        # --------------------------------------------------- Schedules --------------------------------------------
        
        self.Schedule.On(self.DateRules.EveryDay(self.IndexEquityTicker),self.TimeRules.AfterMarketOpen(self.IndexEquityTicker,self.MinsAfter), self.Trade)


    # ------------------------------------------------- On Data -----------------------------------------------------
    
    def OnData(self, data):
        
        self.allData = data  # For checking assets existance in Time
        
        if self.IsWarmingUp:
            return
        
        if not self.TickerTable[self.IndexEquity].Rsi.IsReady:
            self.Debug("RSI not ready")
            return
        if not self.TickerTable[self.IndexEquity].Ppo.IsReady:
            self.Debug("PPO not ready")
            return
        if not self.TickerTable[self.IndexEquity].Macd.IsReady:
            self.Debug("MACD not ready")
            return
        if not self.TickerTable[self.IndexEquity].Rc.IsReady:
            self.Debug("RC not ready")
            return
        
        
    # ---------------------------------------------------------------------------------------------------------------
    # ---------------------------------------------- Trade Function -------------------------------------------------
    # ---------------------------------------------------------------------------------------------------------------
    
    def Trade(self):
        
        self.TotalDays +=1
        self.CollectIndicatorsAndTrends()
        self.DecideInOrOut()
        
        NumberOfInAssets = len(self.InAssetList)
        NumberOfOutAssets = len(self.OutAssetList)
        
        if self.MyInOutIndicator == 1:
            self.TotalDaysIn +=1
            
            if self.IAMOut:
                self.Liquidate()
                self.IAMOut = False
            
            for asset in self.InAssetList:
                
                if not self.allData.ContainsKey(asset): # Asset does not exist in Data atm (ticker not traded ATM) ---- Cash will be left out !!!
                    continue
                
                if not self.Securities[asset].Invested:
                    # orderQuantity = self.CalculateOrderQuantity(asset, 1.00/NumberOfInAssets)
                    #margin = self.Portfolio.GetMarginRemaining(asset, OrderDirection.Buy)
                    #self.Debug("Security Margin " + str(asset) + " Margin: " + str(margin))
                    
                    self.SetHoldings(asset, 1.00/NumberOfInAssets, False,f"PPO:{round(self.IndexEquity_PPO,2)} SLOPE:{round(self.IndexEquity_SLOPE,2)} RSI(Double Expo):{round(self.IndexEquity_RSI,2)}")
                    self.Notify.Sms("+972542224488", self.AlgoName+" "+str(asset) +" In:" + str(self.Securities[asset].Price))
        
        elif self.InOut:
            
             if not self.IAMOut:
                 self.Liquidate()
                 self.IAMOut = True
                 
                 for asset in self.OutAssetList:
                     
                     if not self.allData.ContainsKey(asset): # Asset does not exist in Data atm (ticker not traded ATM) ---- Cash will be left out!
                         continue
                     
                     self.SetHoldings(asset, 1.00/NumberOfOutAssets, False,"Out: "+self.ReasonOut)
                     self.Notify.Sms("+972542224488", self.AlgoName+" "+str(self.IndexEquity) + " Out:" + str(self.Securities[self.IndexEquity].Price))
                 
        self.PlotIt()
    
    # --------------------------------  Decide IN or OUT ----------------------------------------
    
    def GoldenCross(self):
        if abs (self.IndexEquity_MACD - self.IndexEquity_MACD_Signal) <0.1 and True:
            return True
        else:
            return False
    
    def DeathCross(self):
        if (abs(self.IndexEquity_MACD_Histogram) < self.MacdHistOutGap
        and self.IndexEquity_MACD < self.IndexEquity_MACD_SMA * self.MacdFactorOutGap):
            return True
        else:
            return False
    
    
    def DecideInOrOut(self):
    
        # Should we go OUT ?
        self.ReasonOut=" "
        
        if self.DeathCross():
            self.ReasonOut = self.ReasonOut + f" MACD:{round(self.IndexEquity_MACD,2)}"
            
        if self.IndexEquity_PPO < self.IndexEquity_PPO_SMA * self.PpoFactorGap and self.IndexEquity_PPO < self.PpoGap :
           self.ReasonOut = self.ReasonOut + f" PPO:{round(self.IndexEquity_PPO,2)}"
            
        if self.IndexEquity_RSI < self.IndexEquity_RSI_SMA * self.RsiFactorGap:
             self.ReasonOut = self.ReasonOut + f" RSI:{round(self.IndexEquity_RSI,2)}"
            
        if self.IndexEquity_SLOPE < self.SlopeGap:
             self.ReasonOut = self.ReasonOut + f" SLOPE:{round(self.IndexEquity_SLOPE,2)}"
            
        #if self.IndexEquity_ATR > self.IndexEquity_ATR_SMA_SLOW and self.IndexEquity_ATR > self.IndexEquity_ATR_SMA_FAST:
        #   self.ReasonOut = self.ReasonOut + f" ATR:{round(self.IndexEquity_ATR,2)}"
        
        if self.ReasonOut != " ":
            self.MyInOutIndicator = 0
            self.DaysOut = 0    # Zero the DaysOut counter
       
        # Should we get IN?
        elif True:
        
            self.ReasonOut = " "
            if self.DaysOut >= self.MinDaysOut:
                self.MyInOutIndicator = 1
                self.DaysOut = 0
            else:
                self.DaysOut +=1                                        # ---- Unless in Market out Cooldown
        else:
            if self.MyInOutIndicator == 0:
                self.DaysOut +=1
    
    
    # -------------------------  Collect Indicators and trends before Trade ---------------------------------
    
    def CollectIndicatorsAndTrends(self):
        
        self.IndexEquity_RSI = self.TickerTable[self.IndexEquity].Rsi.Current.Value
        self.IndexEquity_RSI_SMA = self.TickerTable[self.IndexEquity].RsiSMA.Current.Value
        self.IndexEquity_PPO = self.TickerTable[self.IndexEquity].Ppo.Current.Value
        self.IndexEquity_PPO_SMA = self.TickerTable[self.IndexEquity].PpoSMA.Current.Value
        self.IndexEquity_MACD = self.TickerTable[self.IndexEquity].Macd.Current.Value
        self.IndexEquity_MACD_SMA = self.TickerTable[self.IndexEquity].MacdSMA.Current.Value
        self.IndexEquity_MACD_Signal = self.TickerTable[self.IndexEquity].Macd.Signal.Current.Value
        self.IndexEquity_MACD_Histogram = self.TickerTable[self.IndexEquity].Macd.Histogram.Current.Value
        self.IndexEquity_SLOPE = self.TickerTable[self.IndexEquity].Rc.Slope.Current.Value
        self.IndexEquity_ATR = self.TickerTable[self.IndexEquity].Atr.Current.Value
        self.IndexEquity_ATR_SMA_FAST = self.TickerTable[self.IndexEquity].AtrSMA_fast.Current.Value
        self.IndexEquity_ATR_SMA_SLOW = self.TickerTable[self.IndexEquity].AtrSMA_slow.Current.Value
        
        # for statistics
        self.TotalRSI+=self.IndexEquity_RSI
        self.AverageRSI = self.TotalRSI / self.TotalDays
        
        #self.Log(f"{self.Time}  RSI:{round(self.IndexEquity_RSI,2)} RSISMA:{round(self.IndexEquity_RSI_SMA,2)}  PPO:{round(self.IndexEquity_PPO,2)}  Slope:{round(self.IndexEquity_SLOPE,2)}  MACD:{round(self.IndexEquity_MACD,2)}  MACDSMA:{round(self.IndexEquity_MACD_SMA,2)}  MACDHIST:{round(self.IndexEquity_MACD_Histogram,2)}")
    
    # -------------------------------------------- Plot Function ------------------------------------------------
    
    def PlotIt(self):    
        
        self.Plot("In/Out Indicator","InOut",self.MyInOutIndicator*100)
        self.Plot("In/Out Indicator","RSI",self.IndexEquity_RSI)
        self.Plot("In/Out Indicator","RSI_SMA",self.IndexEquity_RSI_SMA)
        #self.Plot("PPO","PPO",self.IndexEquity_PPO)
        #self.Plot("RC Slope","SLOPE",self.IndexEquity_SLOPE)
        self.Plot("MACD","MACD",self.IndexEquity_MACD)        
        self.Plot("MACD","MACD Signal",self.IndexEquity_MACD_Signal) 
        self.Plot("MACD","MACD Hist",self.IndexEquity_MACD_Histogram)
        self.Plot("In Assets","In Assets",len(self.InAssetList))
        #self.Plot("ATR","ATR",self.IndexEquity_ATR)
        
        try:
            self.Plot("Fear",self.Fear,self.Securities[self.Fear].Price)
        except:
            pass
        
        # Benchmark Ploting
        try:
            hist = self.History([self.IndexEquity], 2, Resolution.Daily)['close'].unstack(level= 0).dropna() 
            self.IndexEquityHistory.append(hist[self.IndexEquity].iloc[-1])
            perf = self.IndexEquityHistory[-1] / self.IndexEquityHistory[0] * self.StartCash
            self.Plot("Strategy Equity", self.IndexEquity, perf)
        except: # No data for this symbol
            pass
        
        try:
            hist = self.History([self.OutAssetList[0]], 2, Resolution.Daily)['close'].unstack(level= 0).dropna() 
            self.OutEquityHistory.append(hist[self.OutAssetList[0]].iloc[-1])
            perf = self.OutEquityHistory[-1] / self.OutEquityHistory[0] * self.StartCash
            self.Plot("Strategy Equity", self.OutAssetList[0], perf)
        except: # No data for this symbol
            pass


# --------------------------------------------- On Securities Changed --------------------------------------------
    def OnSecuritiesChanged(self, changes):
        
        self.changes = changes
        #self.Debug(f"OnSecuritiesChanged({self.Time}):: {changes}")
        
        
        if self.InAssetSelector == 0:
            
            for security in changes.RemovedSecurities:
                symbol = security.Symbol
                self.Liquidate(symbol) # ------ Need a smarter way !
                if symbol in self.doNotEnter:
                    continue
                if symbol in self.InAssetList:
                    self.InAssetList.remove(symbol)
        
            for security in changes.AddedSecurities:
                symbol = security.Symbol
                if symbol in self.doNotEnter:
                    continue
                if not symbol in self.InAssetList:
                    self.InAssetList.append(symbol)
            
            #self.Debug(len(self.InAssetList))
            
            

# ---------------------------------------------- End of Algo --------------------------------------------------

    def OnEndOfAlgorithm(self):
        self.Liquidate()
        if self.TotalDays>0:
            self.Debug(f"TPV:{round(self.Portfolio.TotalPortfolioValue,2)} Total Days:{self.TotalDays}  Total Days In:{self.TotalDaysIn} {round(self.TotalDaysIn/self.TotalDays*100,2)}%")
            self.Debug(f"Average RSI: {self.AverageRSI}")
    
# ---------------------------------------------- SymbolData --------------------------------------------------
         
class SymbolData:
    def __init__(self, symbol, rsi, rsiSMA, macd, macdSMA, rc, ppo, ppoSMA, atr, atrSMA_fast, atrSMA_slow):
        self.Symbol = symbol
        self.Rsi = rsi
        self.RsiSMA =rsiSMA
        self.Macd = macd
        self.MacdSMA = macdSMA
        self.Rc = rc
        self.Ppo = ppo
        self.PpoSMA = ppoSMA
        self.Atr = atr
        self.AtrSMA_fast = atrSMA_fast
        self.AtrSMA_slow = atrSMA_slow