from clr import AddReference
AddReference("System")
AddReference("QuantConnect.Algorithm")
AddReference("QuantConnect.Indicators")
AddReference("QuantConnect.Common")

from System import *
from QuantConnect import *
from QuantConnect.Data import *
from QuantConnect.Algorithm import *
from QuantConnect.Indicators import *
from System.Collections.Generic import List
from QuantConnect.Data.UniverseSelection import *
from Selection.FundamentalUniverseSelectionModel import FundamentalUniverseSelectionModel

from QuantConnect.Data.UniverseSelection import *
from Selection.FundamentalUniverseSelectionModel import FundamentalUniverseSelectionModel
from itertools import groupby
from math import ceil


class MeanReversionAlgo(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2020, 6, 6)  # Set Start Date

        self.globalResolution = Resolution.Daily
        
        self.initialcash = 100000
        self.SetCash(self.initialcash)  # Set Strategy Cash
        # Add SPY to set scheduled events
        self.AddEquity("SPY", self.globalResolution)
        # Setting Universe
        self.UniverseSettings.Resolution = self.globalResolution
        self.SetUniverseSelection(QC500UniverseSelectionModel())
        
        # this add universe method accepts two parameters:
        # - coarse selection function: accepts an IEnumerable<CoarseFundamental> and returns an IEnumerable<Symbol>
        # - fine selection function: accepts an IEnumerable<FineFundamental> and returns an IEnumerable<Symbol>

        self.UniverseSettings.Resolution = self.globalResolution
        
        #self.AddUniverseSelection(FineFundamentalUniverseSelectionModel(self.CoarseSelectionFunction, self.FineSelectionFunction))
        self.AddUniverse(self.CoarseSelectionFunction)


        self.numberOfSymbolsFine = 500
        self._changes = None
        
        self.dollarVolumeBySymbol = {}
        
        #self.AutomaticIndicatorWarmUp=True
        
        self.SetWarmup(200, Resolution.Daily)
        
        self.dataDict = {}
        
        self.dataDictBoughtSymbols = {}
        
        self.verboseLogging = False
        
        
        self.selections = []
        
        self.lastMonth = -1
        
        self.AddEquity("SHY")
        
        # schedule an event to fire every trading day for a security the
        # time rule here tells it to fire 5 minutes after SPY's market open
        #self.Schedule.On(self.DateRules.EveryDay("SPY"), self.TimeRules.AfterMarketOpen("SPY", 5), self.EveryDayAfterMarketOpen)
        
        # schedule an event to fire every trading day for a security the
        # time rule here tells it to fire 10 minutes before SPY's market close
        self.Schedule.On(self.DateRules.EveryDay("SPY"),self.TimeRules.BeforeMarketClose("SPY", 10),self.EveryDayBeforeMarketClose)
        
        self.Schedule.On(self.DateRules.MonthStart("SPY",0),self.TimeRules.AfterMarketOpen("SPY"),self.buy)
        

  
    def EveryDayAfterMarketOpen(self):
        
        [self.SetHoldings(symbol, 0) for symbol in self.Portfolio.Keys if self.Portfolio[symbol].Invested and symbol not in self.selections]
        

        
        for symbol in self.selections:
            if self.Portfolio[symbol].Invested is False:
                self.SetHoldings(symbol, 0.1)

        
        if can_trade_SHY:
            self.SetHoldings("SHY",1-len(self.selections)/10)
            can_trade_SHY = False
        

    
    def buy(self):
        
        [self.SetHoldings(symbol, 0) for symbol in self.Portfolio.Keys if self.Portfolio[symbol].Invested and symbol not in self.selections]
        
        if len(self.selections)==0:
            return
        
        if len(self.selections) < 10:
            '''
            for i in range(10-len(self.selections),10):
                self.Debug("Buying SHY from")
                self.Debug(10-len(self.selections))
                self.selections.append("SHY")
            '''    
            self.selections.extend(["SHY" for i in range(10-len(self.selections),10)])
            
        self.Debug("symbols we are going to buy")
        for i in self.selections:
            self.Debug(i)
            
        
        for symbol in self.selections:
            if self.Portfolio[symbol].Invested is False:
                self.SetHoldings(symbol, 0.1)


        
   

        
                
      

            
    def EveryDayBeforeMarketClose(self):
        
        
        values = [x for x in self.dataDict.values()]
        
        values_rsi_higher_than_55 = [x for x in values if x.RSI_NOT_OK] 

        symbols = [x for x in values_rsi_higher_than_55] 
        
        for x in symbols:
            if x.symbol in self.selections:
                self.Debug("Liquidating due to 2 period RSI higher than 95 Symbol "+str(x.symbol)+ " Current RSI "+str(x.RSI)) 
                self.Liquidate(x.symbol) 
                self.selections.remove(x.symbol)


    
    
    # sort the data by daily dollar volume and take the top 'NumberOfSymbols'
    def CoarseSelectionFunction(self, coarse):
        
        
        
        CoarseWithFundamental = [x for x in coarse if x.HasFundamentalData and x.Price > 5]
        
        for i in CoarseWithFundamental:
            if i.Symbol not in self.dataDict:
                self.dataDict[i.Symbol] = SymbolData(i.Symbol)
            
            self.dataDict[i.Symbol].update(i.EndTime, i.AdjustedPrice)  
        

        if self.Time.month == self.lastMonth:
            return self.Universe.Unchanged
        self.lastMonth = self.Time.month



        values = [x for x in self.dataDict.values()]
        
        

        
        values.sort(key=lambda x: x.STD, reverse=True)
       
        finally_filtered = values[:10]
        
        if self.verboseLogging and len(finally_filtered) > 0:
            self.Debug("filter --------> Top 10 by Standard Deviation")
            for i in finally_filtered:
                self.Debug(str(i.symbol)+" Current Standard Deviation "+str(i.STD))
        
        if self.verboseLogging:
            self.Debug(' coarse # of stocks {}'.format(len(finally_filtered)))
        
        if len(finally_filtered) == 0:
            return self.Universe.Unchanged
        

        # some filtering
        self.selections = [x.symbol for x in finally_filtered]
        return self.selections
        
 
    
    def FineSelectionFunction(self, fine):
        

        sortedBySector = sorted([x for x in fine if x.CompanyReference.CountryId == "USA"
                                        and x.CompanyReference.PrimaryExchangeID in ["NYS","NAS"]
                                        and (self.Time - x.SecurityReference.IPODate).days > 180
                                        and x.SecurityReference.IsPrimaryShare
                                        and x.MarketCap > 5e8],
                               key = lambda x: x.CompanyReference.IndustryTemplateCode)

        count = len(sortedBySector)

        # If no security has met the QC500 criteria, the universe is unchanged.
        # A new selection will be attempted on the next trading day as self.lastMonth is not updated
        if count == 0:
            return self.Universe.Unchanged

        # Update self.lastMonth after all QC500 criteria checks passed
        self.lastMonth = self.Time.month

        
        return [x.Symbol for x in sortedBySector]
    

        
    
    def OnData(self, slice):
        
        '''
        #Take profit logic
        if self.Portfolio.Invested: 
            if self.Portfolio.TotalPortfolioValue >  self.initialcash* 1.05:# means a 5% take profit target, if the initial portfolio value is 100000 with 1.05 you will take profit when the value of portfolio is greater than 105 000 $.
                self.Liquidate()
        '''
        
        #Stop loss logic
        if self.Portfolio.Invested: 
            if self.Portfolio.TotalPortfolioValue < self.initialcash*0.70:  # means a 30% stop loss. In this case 0.9 means that the portfolio is valued a 90% of the original value, so if the initial value is 100 000 $ it means 90 000$, a 10% stop loss. if you set self.initialcash*0.5 means a 50% stop loss and so on.
                self.Liquidate()        
        
        
        # this event fires whenever we have changes to our universe
    def OnSecuritiesChanged(self, changes):
            self._changes = changes
            #self.Log(f"OnSecuritiesChanged({self.UtcTime}):: {changes}")
        


        
class SymbolData(object):
    
    def __init__(self, symbol):
        self.symbol = symbol
        self.History = RollingWindow[float](126)# you can't change this
        self.STD = StandardDeviation(126)



        self.RSI = RelativeStrengthIndex(2)
        self.RSI_OK = False
        self.RSI_NOT_OK = False

        self.currentPrice = 0
        

    def update(self, time, value):

        self.History.Add(float(value))
        self.STD.Update(time,value)
        self.RSI.Update(time, value)
        self.currentPrice=value
        

        
        self.RSI_OK = self.RSI.IsReady and self.RSI.Current.Value < 10

        self.RSI_NOT_OK = self.RSI.IsReady and self.RSI.Current.Value > 95