from AlgorithmImports import *

class GrowthStocksAlgorithm(QCAlgorithm):
    def Initialize(self):
        self.SetStartDate(2010, 1, 1)  # Start Date
        self.SetEndDate(2024, 1, 1)    # End Date
        self.SetCash(10000)            # Strategy Cash
        
        # Universe and Resolution Settings
        self.UniverseSettings.Resolution = Resolution.Daily
        self.AddUniverse(self.CoarseSelectionFunction, self.FineSelectionFunction)
        
        # Schedule rebalance
        self.Schedule.On(self.DateRules.MonthStart(), self.TimeRules.At(10, 0), self.RebalancePortfolio)
        
        # Trailing stop loss percentage and dictionary for high water marks
        self.trailingStopLossPercent = 0.07
        self.highWaterMarks = dict()
        
        # Changes tracker
        self.changes = None
        
        # Add VIX for market volatility
        self.vix = self.AddData(CBOE, "VIX").Symbol
        self.vixHighVolatilityThreshold = 20  # Define high volatility threshold
        self.canBuyEquities = True  # Flag to control equity buying
        
    def CoarseSelectionFunction(self, coarse):
        if not self.canBuyEquities:
            return []
        filtered_coarse = [x for x in coarse if x.DollarVolume > 1e6 and x.Price > 5]
        sorted_by_liquidity = sorted(filtered_coarse, key=lambda x: x.DollarVolume, reverse=True)
        selected_symbols = [x.Symbol for x in sorted_by_liquidity if x.Price * x.DollarVolume / x.Price > 2e9]
        return selected_symbols[:100]
        
    def FineSelectionFunction(self, fine):
        filtered_fine = [x for x in fine if x.OperationRatios.RevenueGrowth.OneYear > 0.0
                        and x.OperationRatios.NetIncomeGrowth.OneYear > 0.0
                        and x.EarningReports.BasicEPS.TwelveMonths > 0
                        and (x.ValuationRatios.PEGRatio > 0 and x.ValuationRatios.PEGRatio < 1.5)
                        and x.FinancialStatements.BalanceSheet.TotalEquity.Value > 0]
        
        sorted_by_growth = sorted(filtered_fine, key=lambda x: x.OperationRatios.RevenueGrowth.OneYear, reverse=True)
        return [x.Symbol for x in sorted_by_growth[:10]]

    def RebalancePortfolio(self):
        if self.changes is None or len(self.changes) == 0:
            return
        
        weight_per_security = 1.0 / len(self.changes)
        
        for symbol in self.changes:
            if not self.Securities[symbol].Invested:
                self.SetHoldings(symbol, weight_per_security)
                self.highWaterMarks[symbol] = self.Securities[symbol].Price
                
        self.changes = []  # Reset the changes after rebalancing

    def OnSecuritiesChanged(self, changes):
        self.changes = [x.Symbol for x in changes.AddedSecurities]
    
    def OnData(self, data):
        symbolsToLiquidate = []

        for symbol, highWaterMark in self.highWaterMarks.items():
            # Check if the symbol has data before accessing its Close price
            if symbol in data and data[symbol] is not None:
                currentPrice = data[symbol].Close
                if currentPrice > highWaterMark:
                    self.highWaterMarks[symbol] = currentPrice

                trailingStopPrice = highWaterMark * (1 - self.trailingStopLossPercent)
                if currentPrice < trailingStopPrice:
                    symbolsToLiquidate.append(symbol)

        for symbol in symbolsToLiquidate:
            self.Liquidate(symbol)
            del self.highWaterMarks[symbol]        
        # VIX-based market volatility handling
        if data.ContainsKey(self.vix):
            currentVixValue = data[self.vix].Close
            isHighVolatility = currentVixValue > self.vixHighVolatilityThreshold
            
            if isHighVolatility:
                self.AdjustEquityExposure(False)  # Adjust for high volatility
            else:
                self.AdjustEquityExposure(True)  # Reset for normal conditions

    def AdjustEquityExposure(self, canBuy):
        self.canBuyEquities = canBuy
        if not canBuy:
            # Exit profitable positions during high volatility
            for symbol, security in self.Securities.items():
                if security.Invested:
                    currentPrice = security.Price
                    buyPrice = security.Holdings.AveragePrice
                    if currentPrice > buyPrice:
                        self.Liquidate(symbol)