from CboeVixAlphaModel import CboeVixAlphaModel
from Execution.VolumeWeightedAveragePriceExecutionModel import VolumeWeightedAveragePriceExecutionModel
from Portfolio.EqualWeightingPortfolioConstructionModel import EqualWeightingPortfolioConstructionModel
from Risk.MaximumDrawdownPercentPerSecurity import MaximumDrawdownPercentPerSecurity
from SmallCapGrowthStocks import SmallCapGrowthStocks

class CalibratedVerticalSplitter(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2019, 1, 10)  # Set Start Date
        self.SetEndDate(2019, 12, 31)
        self.SetCash(100000)  # Set Strategy Cash
        # self.AddEquity("SPY", Resolution.Minute)
        self.AddAlpha(CboeVixAlphaModel(self))

        self.SetExecution(VolumeWeightedAveragePriceExecutionModel())

        self.SetPortfolioConstruction(EqualWeightingPortfolioConstructionModel())

        self.SetRiskManagement(MaximumDrawdownPercentPerSecurity(0.01))

        self.SetUniverseSelection(SmallCapGrowthStocks())


    def OnData(self, data):
        '''OnData event is the primary entry point for your algorithm. Each new data point will be pumped in here.
            Arguments:
                data: Slice object keyed by symbol containing the stock data
        '''

        # if not self.Portfolio.Invested:
        #    self.SetHoldings("SPY", 1)

from QuantConnect.Data.Custom.CBOE import *
class CboeVixAlphaModel:
    
    def __init__(self, algorithm):
        self.vix = algorithm.AddData(CBOE, "VIX").Symbol
    
    def Update(self, algorithm, data):
        insights = []
        
        if not data.ContainsKey(self.vix):
            return insights
        
        vix_data = data.Get(CBOE, self.vix)
        
        ## The Cboe Volatility Index® (VIX® Index) is the most popular benchmark index to measure
        ## the market’s expectation of future volatility. The VIX Index is based on
        ## options of the S&P 500® Index, considered the leading indicator of the broad
        ## U.S. stock market. The VIX Index is recognized as the world’s premier gauge
        ## of U.S. equity market volatility.
        
        ## Generate Insights here!
        
        return insights
    
    def OnSecuritiesChanged(self, algorithm, changes):
        # For instruction on how to use this method, please visit
        # https://www.quantconnect.com/docs/algorithm-framework/alpha-creation#Alpha-Creation-Good-Design-Patterns
        pass

from Selection.FundamentalUniverseSelectionModel import FundamentalUniverseSelectionModel

class SmallCapGrowthStocks(FundamentalUniverseSelectionModel):
    '''
    This module selects the most liquid stocks listed on the Nasdaq Stock Exchange.
    '''

    def __init__(self, filterFineData = True, universeSettings = None, securityInitializer = None):
        '''Initializes a new default instance of the TechnologyUniverseModule'''
        super().__init__(filterFineData, universeSettings, securityInitializer)
        self.numberOfSymbolsCoarse = 1000
        self.numberOfSymbolsFine = 100
        self.dollarVolumeBySymbol = {}
        self.lastMonth = -1

    def SelectCoarse(self, algorithm, coarse):
        '''
        Performs a coarse selection:
        
        -The stock must have fundamental data
        -The stock must have positive previous-day close price
        -The stock must have positive volume on the previous trading day
        '''
        if algorithm.Time.month == self.lastMonth: 
            return Universe.Unchanged

        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)[:self.numberOfSymbolsCoarse]

        self.dollarVolumeBySymbol = {x.Symbol:x.DollarVolume for x in sortedByDollarVolume}
        
        # If no security has met the QC500 criteria, the universe is unchanged.
        if len(self.dollarVolumeBySymbol) == 0:
            return Universe.Unchanged

        return list(self.dollarVolumeBySymbol.keys())

    def SelectFine(self, algorithm, fine):
        '''
        Performs a fine selection for companies in the Morningstar Banking Sector
        '''
        # Filter stocks and sort on dollar volume
        sortedByDollarVolume = sorted([x for x in fine if x.AssetClassification.StyleBox == StyleBox.SmallGrowth],
            key = lambda x: self.dollarVolumeBySymbol[x.Symbol], reverse=True)

        if len(sortedByDollarVolume) == 0:
            return Universe.Unchanged
            
        self.lastMonth = algorithm.Time.month

        return [x.Symbol for x in sortedByDollarVolume[:self.numberOfSymbolsFine]]