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

from System import *
from System.Collections.Generic import List
from QuantConnect import *
from QuantConnect.Algorithm import QCAlgorithm
from QuantConnect.Data.UniverseSelection import *
from datetime import datetime, timedelta

import datetime

class ConsolidatorHelp(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2020, 12, 1)  # Set Start Date
        
        self.SetEndDate(datetime.datetime.now())
        #self.SetEndDate(2020, 12, 14)  # Set End Date
        self.SetCash(600000)  # Set Strategy Cash
        self.SetWarmup(20)
        self.SetBenchmark("SPY")
        self.UniverseSettings.Resolution = Resolution.Minute
        self.AddUniverse(self.CoarseSelectionFunction, self.FineSelectionFunction)
        self.UniverseSettings.Leverage = 2
        self._changes = None
        self.symbols = {}
        self.daySymbols = {}
        self.profits = {}
        self.unrealized = {}
        self.openEq = {}
        self.yesterdaysOHLC = {}
        self.currentHigh = {}
        self.currentOpen = {}
        self.uniFilter = []
 
    def OnData(self, data):
        
        # Make sure indicators and rolling windows are ready
        if not all([symbol.IsReady for symbol in self.symbols.values()]):
            return
        
        if not all([symbol.IsReady for symbol in self.daySymbols.values()]):
            return
        

        ######### Get daily values from yesterday(OHLC)  and today(OH) #####################
        # Initialize dictionaries
        if len(self.currentHigh) == 0:
            for symbol, value in self.daySymbols.items():
                self.currentHigh.setdefault(symbol.Value, 0)
        if len(self.currentOpen) == 0:
            for symbol, value in self.daySymbols.items():
                self.currentOpen.setdefault(symbol.Value, 0)
        
        for symbol, value in self.daySymbols.items():
            
            if  data.ContainsKey(symbol) and \
                data[symbol] != None:
                self.Debug("AM I GETTING DATA?????????????????????????????????????")
                self.Debug(str(symbol))    
                
                
                self.yesterdaysOHLC[symbol.Value] = \
                    (value.dBars[0].Open, value.dBars[0].High, value.dBars[0].Low, value.dBars[0].Close)
                # Get todays High    
                if data[symbol].High > self.currentHigh[symbol.Value]:
                    self.currentHigh[symbol.Value] = data[symbol].High
                # Get todays open    
                if self.currentOpen[symbol.Value] > 0:
                    break
                else:
                    self.currentOpen[symbol.Value] = data[symbol].Open
            
        
        for symbol, value in self.symbols.items():
            
            if (value.Bars[0].EndTime.hour < 15 or (value.Bars[0].EndTime.hour == 15 and value.Bars[0].EndTime.minute < 30)) and \
                data.ContainsKey(symbol) and not self.Portfolio[symbol].Invested and \
                value.Bars[0].EndTime == self.Time and \
                data[symbol] != None:
                
                ########## GAP ##############################
                if symbol.Value in self.yesterdaysOHLC and not self.tradeToday():
                    # Get todays Open and High
                    todaysOpen = self.currentOpen[symbol.Value]
                    todaysHigh = self.currentHigh[symbol.Value]
                    yesterdaysOpen  = self.yesterdaysOHLC[symbol.Value][0]
                    yesterdaysHigh  = self.yesterdaysOHLC[symbol.Value][1]
                    yesterdaysLow   = self.yesterdaysOHLC[symbol.Value][2]
                    yesterdaysClose = self.yesterdaysOHLC[symbol.Value][3]
                    
                    
    
    # Add and remove stocks whenever we have changes to our universe
    def OnSecuritiesChanged(self, changes):
        self._changes = changes
        for security in changes.AddedSecurities:
            symbol = security.Symbol
            
            # For 5 minute quotes
            if symbol not in self.symbols:
                self.symbols[symbol] = SymbolData(self, symbol)
                
                
            
            # For daily quotes
            if symbol not in self.daySymbols:
                self.daySymbols[symbol] = DaySymbolData(self, symbol)
                
        
        # Dont remove securities
        '''
        for security in changes.RemovedSecurities:
            symbol = security.Symbol
            if symbol in self.symbols:
                symbolData = self.symbols.pop(symbol, None)
                self.SubscriptionManager.RemoveConsolidator(symbol, symbolData.consolidator)
        '''
    ############# FILTERS###################################
    def CoarseSelectionFunction(self, coarse ):
        
        
        top = [x for x in coarse if x.HasFundamentalData
                           and 500 > x.Price > 40
                           and  x.Volume > 500000]
        self.Debug("len of coarse = " + str(len(top)))
        return [x.Symbol for x in top]
        
    def FineSelectionFunction(self, fine ):
        
        top = [x for x in fine if  x.ValuationRatios.PERatio > 20 and
                           x.ValuationRatios.TrailingDividendYield < .02 and 
                           not x.AssetClassification.MorningstarSectorCode == MorningstarSectorCode.BasicMaterials and
                           not x.AssetClassification.MorningstarSectorCode == MorningstarSectorCode.Energy]
        top = sorted(top, key=lambda x: x.ValuationRatios.PERatio, reverse=True)
        self.Debug("len of fine = " + str(len(top)))
        
        return  [x.Symbol for x in top][:5]
    
        
    
class SymbolData:
    
    def __init__(self, algorithm, symbol):
        self.algorithm = algorithm
        self.symbol = symbol
        
        ##########  INDICATORS  #################################
        self.sma = SimpleMovingAverage(9)
        self.trix = algorithm.TRIX(symbol, 9, Resolution.Minute)
        self.Bars = RollingWindow[TradeBar](2)
        
        ########### CONSOLIDATORS and WINDOWS ###########################
        self.consolidator = TradeBarConsolidator(timedelta(minutes=5))
        self.consolidator.DataConsolidated += self.OnDataConsolidated
        self.trixWindow = RollingWindow[IndicatorDataPoint](2) # Define our rolling window to hold indicator points
        self.trix.Updated += self.OnTrixUpdated # Set our event handler
        
        
        algorithm.SubscriptionManager.AddConsolidator(symbol, self.consolidator)
        algorithm.RegisterIndicator(symbol, self.sma, self.consolidator)
        algorithm.RegisterIndicator(symbol, self.trix, self.consolidator)
        
        
        
    # Update TRIX Window
    def OnTrixUpdated(self, sender, updated):
        if self.trix.IsReady: 
            self.trixWindow.Add(updated) # Add updated indicator data to rolling window
        
   
    # Lock and load
    @property    
    def IsReady(self):
        return self.trix.IsReady and self.trixWindow.IsReady 
        
    # Make sure I can get stock values in Window.
    def OnDataConsolidated(self, sender, bar):
        self.Bars.Add(bar)
        '''
        self.algorithm.Debug(f"Data Consolidatoed for {self.symbol} \
            at {bar.EndTime} with bar: {bar} and sma {self.sma} and trix {self.trix} and \
            kelt {self.kelt.LowerBand}")
        '''
        
class DaySymbolData:
    
    def __init__(self, algorithm, symbol):
        self.algorithm = algorithm
        self.symbol = symbol
        
        ##########  INDICATORS  #################################
        self.dBars = RollingWindow[TradeBar](2)
        
        
        ########### CONSOLIDATORS and WINDOWS ###########################
        self.consolidator = TradeBarConsolidator(timedelta(days=1))
        self.consolidator.DataConsolidated += self.OnDataConsolidated
        
        algorithm.SubscriptionManager.AddConsolidator(symbol, self.consolidator)
        
    # Lock and load
    @property    
    def IsReady(self):
        return self.dBars.IsReady
    
    # Make sure I can get stock values in Window.
    def OnDataConsolidated(self, sender, bar):
        self.dBars.Add(bar)
        self.algorithm.Debug(f"Data Consolidated for {self.symbol} \
            at {bar.EndTime} with bar: {bar} ")