#region imports
from AlgorithmImports import *
#endregion
class AlertTanMosquito(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(1997, 1, 1)
        self.SetEndDate(2022, 1, 1)
        self.SetCash(100000)
        
        self.spy = self.AddEquity("SPY", Resolution.Daily).Symbol
        self.SetBenchmark("SPY")
        self.SetBrokerageModel(BrokerageName.InteractiveBrokersBrokerage, AccountType.Margin)
        self.SetSecurityInitializer(self.CustomSecurityInitializer)
        self.Settings.FreePortfolioValuePercentage = 0.05
        
        # Set universe of stocks
        self.maximums = {}
        self.activeStocks = set()
        self.AddUniverse(self.MyCoarseFilter, self.MyFineFilter)
        self.UniverseSettings.Resolution = Resolution.Daily
        
        # Variables
        self.stoploss = float(self.GetParameter('Stop Loss'))
        self.filtered_stop = float(self.GetParameter('Filtered Stop Loss'))
        self.num_coarse = int(self.GetParameter('Number of Course'))
        self.num_fine = int(self.GetParameter('Number of Fine'))
        self.lookback = int(self.GetParameter('Lookback Period'))
        self.sma_period = int(self.GetParameter('Simple Moving Average'))
        self.sma = self.SMA(self.spy, self.sma_period, Resolution.Daily)
        self.SetWarmUp(self.sma_period)
        #self.closeWindow = RollingWindow[float](self.lookback)

    def CustomSecurityInitializer(self, security):
        security.SetLeverage(1)
        security.SetDataNormalizationMode(DataNormalizationMode.Raw)

    def MyCoarseFilter(self, coarse):

        if len(self.activeStocks) >= 20:
            return self.Universe.Unchanged
        
        # Rebalance universe if a portfolio is not full
        # if len(self.activeStocks) < 20:
        #   self.Log("Rebalancing as active stock count: " + str(len(self.activeStocks)))

        # Performs coarse selection for the QC500 constituents.
        # The stocks must have fundamental data
        # The stock must have positive previous-day close price
        # The stock must have positive volume on the previous trading day
        
        sortedByDollarVolume = sorted([x for x in coarse if x.HasFundamentalData 
                                        and x.Volume > 0 
                                        and x.Price > 0],
                                        key = lambda x: x.DollarVolume, reverse=True)

        count = len(sortedByDollarVolume)

        # If no security has met the QC500 criteria, the universe is unchanged.
        if count == 0:
            #self.Log("No securities have met the coarse conditions.")
            return Universe.Unchanged

        # Filter for stocks above yearly high
        universe = sortedByDollarVolume[:self.num_coarse]
        for security in universe:
            symbol = security.Symbol
            
            if symbol not in self.maximums:
                self.maximums[symbol] = SelectionData(symbol, self.lookback)
            
            # Update SelectionData Object with current EOD price
            maxi = self.maximums[symbol]
            maxi.update(security.EndTime, security.AdjustedPrice)
            
        # Filter out values of the dictionary where price is below yearly high
        values = list(filter(lambda x: x.is_above_max, self.maximums.values()))
            
        count = len(values)

        if count == 0:
            #self.Log("No securities have met the yearly high conditions.")
            return Universe.Unchanged
        
        # Return the symbol objects of our sorted collection
        return [x.symbol for x in values[:self.num_coarse]]

    def MyFineFilter(self, fine):

        # Performs fine selection for the QC500 constituents
        # Company's headquarter must in the U.S.
        # Stock must be traded on either the NYSE or NASDAQ
        # At least half a year since its initial public offering
        # The stock's market cap must be greater than 500 million

        sortedByMarketCap = 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.MarketCap > 5e8],
                               key = lambda x: x.MarketCap, reverse=True)

        count = len(sortedByMarketCap)

        # If no security has met the QC500 criteria, the universe is unchanged.
        if count == 0:
            #self.Log("No securities have met the fine conditions.")
            return Universe.Unchanged
               
        # Return the symbol objects of our sorted collection
        return [x.Symbol for x in sortedByMarketCap[:self.num_fine]]
        
    def OnData(self, data):
        
        for security in self.ActiveSecurities.Values:
            if data.ContainsKey(security.Symbol):
                pass
                #self.closeWindow.Add(data[security.Symbol].Close)
            else:
                return

        for security in self.ActiveSecurities.Values:
            symbol = security.Symbol

            # Set up yearly Low indicator
            history = self.History(symbol, self.lookback, Resolution.Daily)["close"]
            yearly_low = min(history)
            yearly_high = max(history)
            price = security.Price
            
            if self.sma.IsReady:
                # Get the current sma value
                sma_value = self.sma.Current.Value

                if self.ActiveSecurities["SPY"].Price >= sma_value:

                    # Exit trade using market on open order when price falls below yearly low
                    if price <= yearly_low and self.Portfolio[symbol].Invested:
                        self.Liquidate(symbol, "Exit Conditions Met")
                        self.activeStocks.remove(symbol)
                        self.Log(str(symbol) + ' fell below yearly low outside of filter.')
                    
                    # Place market on open order when price is above yearly high
                    if not self.Portfolio[symbol].Invested and len(self.activeStocks) < 20 and symbol is not "SPY":
                        self.SetHoldings(symbol, 1/self.num_fine, False, "Entry Conditions Met")
                        self.activeStocks.add(symbol)
                        self.Log(str(symbol) + ' added to portfolio')
                
                elif self.ActiveSecurities["SPY"].Price < sma_value:

                    # Exit trade using market on open order when price falls below yearly low or 30% from high
                    if price <= yearly_low and self.Portfolio[symbol].Invested:
                        self.Liquidate(symbol, "Exit Conditions Met")
                        self.activeStocks.remove(symbol)
                        self.Log(str(symbol) + ' fell below yearly low during filter.')
                    
                    elif price <= self.filtered_stop*yearly_high and self.Portfolio[symbol].Invested:
                        self.Liquidate(symbol, "Exit Conditions Met")
                        self.activeStocks.remove(symbol)
                        self.Log(str(symbol) + ' fell 30 percent from high during filter.')
                    
                    # Do not enter new trades
                    pass
            
    def OnOrderEvent(self, orderEvent):
               
        order = self.Transactions.GetOrderById(orderEvent.OrderId)

        # Remove invalid orders from active stocks set
        if orderEvent.Status == OrderStatus.Invalid and order.Type == OrderType.Market:
            self.activeStocks.remove(order.Symbol)

        # Set Stop Loss order on fill
        if orderEvent.Status == OrderStatus.Filled and self.Portfolio[order.Symbol].Invested:
            self.StopMarketOrder(order.Symbol, -order.Quantity, round(self.stoploss*self.Portfolio[order.Symbol].AveragePrice, 2))
            
        # Remove sold securities from active stocks set
        if orderEvent.Status == OrderStatus.Filled and order.Type == OrderType.StopMarket:
            self.activeStocks.remove(order.Symbol)
            self.Log(str(order.Symbol) + ' removed from portfolio as Stop was initiated.')
    
    def OnSecuritiesChanged(self, changes):
        
       pass

class SelectionData(object):

    def __init__(self, symbol, period):
        self.symbol = symbol
        self.high = Maximum(period)
        self.is_above_max = False
    
    def is_ready(self):
        return self.high.IsReady

    def update(self, time, price):
        if self.high.Update(time, price):
            high = self.high.Current.Value
            self.is_above_max = price >= high