from AlgorithmImports import *
from Selection.FundamentalUniverseSelectionModel import FundamentalUniverseSelectionModel
from itertools import groupby
from math import ceil

class Top100UniverseSelectionModel(FundamentalUniverseSelectionModel):
    '''Defines the top 100 universe as a universe selection model for framework algorithm
    For details: https://github.com/QuantConnect/Lean/pull/1663'''

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

    def SelectCoarse(self, algorithm, coarse):
        '''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'''
        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.
        # A new selection will be attempted on the next trading day as self.lastMonth is not updated
        if len(self.dollarVolumeBySymbol) == 0:
            return Universe.Unchanged

        # return the symbol objects our sorted collection
        return list(self.dollarVolumeBySymbol.keys())

    def SelectFine(self, algorithm, fine):
        '''Performs fine selection for the QC500 constituents
        The company's headquarter must in the U.S.
        The 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'''

        sortedBySector = sorted([x for x in fine if x.CompanyReference.CountryId == "USA"
                                        and x.CompanyReference.PrimaryExchangeID in ["NYS","NAS"]
                                        and (algorithm.Time - x.SecurityReference.IPODate).days > 180
                                        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 Universe.Unchanged

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

        percent = self.numberOfSymbolsFine / count
        sortedByDollarVolume = []

        # select stocks with top dollar volume in every single sector
        for code, g in groupby(sortedBySector, lambda x: x.CompanyReference.IndustryTemplateCode):
            y = sorted(g, key = lambda x: self.dollarVolumeBySymbol[x.Symbol], reverse = True)
            c = ceil(len(y) * percent)
            sortedByDollarVolume.extend(y[:c])

        sortedByDollarVolume = sorted(sortedByDollarVolume, key = lambda x: self.dollarVolumeBySymbol[x.Symbol], reverse=True)
        return [x.Symbol for x in sortedByDollarVolume[:self.numberOfSymbolsFine]]

from MeanMomentumAlphaModel import MeanMomentumAlphaModel
from Top100UniverseSelectionModel import Top100UniverseSelectionModel

class MeanMomentumAlgorithm(QCAlgorithmFramework):
    
    def Initialize(self):
        
        self.UniverseSettings.Resolution = Resolution.Daily
        self.UniverseSettings.Leverage = 1.0
        self.SetBrokerageModel(BrokerageName.AlphaStreams)
        
        self.SetBenchmark(Symbol.Create("SPY", SecurityType.Equity, Market.USA))
        
        self.SetStartDate(2020,8,20)
        self.SetEndDate(2021,9,1)
        self.SetCash(25000)
        
        # S&P100
        tickers = [
            'AAPL', 'ADBE', 'ADI', 'ADP', 'ADSK', 'AEP', 'ALGN', 'AMAT', 'AMD', 
            'AMGN', 'AMZN', 'ANSS', 'ASML', 'ATVI', 'AVGO', 'BIDU', 'BIIB', 
            'BKNG', 'CDNS', 'CDW', 'CERN', 'CHKP', 'CHTR', 'CMCSA', 'COST', 
            'CPRT', 'CRWD', 'CSCO', 'CSX', 'CTAS', 'CTSH', 'DLTR', 'DOCU', 
            'DXCM', 'EA', 'EBAY', 'EXC', 'FAST', 'FB', 'FISV', 'FOX', 'FOXA', 
            'GILD', 'GOOG', 'GOOGL', 'HON', 'IDXX', 'ILMN', 'INCY', 'INTC', 
            'INTU', 'ISRG', 'JD', 'KDP', 'KHC', 'KLAC', 'LRCX', 'LULU', 'MAR', 
            'MCHP', 'MDLZ', 'MELI', 'MNST', 'MRNA', 'MRVL', 'MSFT', 'MTCH', 
            'MU', 'NFLX', 'NTES', 'NVDA', 'NXPI', 'OKTA', 'ORLY', 'PAYX', 
            'PCAR', 'PDD', 'PEP', 'PTON', 'PYPL', 'QCOM', 'REGN', 'ROST', 
            'SBUX', 'SGEN', 'SIRI', 'SNPS', 'SPLK', 'SWKS', 'TCOM', 'TEAM', 
            'TMUS', 'TSLA', 'TXN', 'VRSK', 'VRSN', 'VRTX', 'WBA', 'WDAY', 
            'XEL', 'XLNX', 'ZM'
            ]

        
        symbols = [
            Symbol.Create(ticker, SecurityType.Equity, Market.USA)
            for ticker in tickers
            ]
            
        self.AddUniverseSelection(ManualUniverseSelectionModel(symbols))
        
        self.SetAlpha(MeanMomentumAlphaModel())
        
        self.SetPortfolioConstruction(EqualWeightingPortfolioConstructionModel())
        
        # self.SetExecution(ImmediateExecutionModel())
        self.SetExecution(NullExecutionModel())
        
        self.SetRiskManagement(TrailingStopRiskManagementModel(0.8))
        
        self.SetWarmUp(timedelta(30))
        
        
    def OnOrderEvent(self, orderEvent):
        if orderEvent.Status == OrderStatus.Filled:
            self.Log(
                f'Processed Order: {orderEvent.Symbol}, Quantity: {orderEvent.FillQuantity}'
                )

class MeanMomentumAlphaModel(AlphaModel):
    
    def __init__(self):
        self.Name = "MeanMomentumAlphaModel"
        self.symbolDataBySymbol = {}
        self.predictionInterval = timedelta(1)
        self.rollingWindow = 14
        self.schedule_symbol = None
        self.scheduled_event_open = None
        self.scheduled_event_close = None
        self.algorithm = None
        self.buy_symbol = None
        
        
    def ones(self, n):
        return [1 for i in range(n)]
        
        
    # def ClosePositions(self, algorithm):
    #     openPositions = [ x.Symbol for x in algorithm.Portfolio.Values if x.Invested ]
    
    #     insights=[]
    #     for x in openPositions:
    #         insight = Insight.Price(x, Expiry.EndOfDay, InsightDirection.Flat)
    #         insights.append(insight)
    #     algorithm.EmitInsights(insights)
    
    
    def Update(self, algorithm, slice):
        '''Updates this alpha model with the latest data from the algorithm.
        This is called each time the algorithm receives data for subscribed securities
        Args:
            algorithm: The algorithm instance
            data: The new data available
        Returns:
            The new insights generated'''

        insights = []
        
        n = self.rollingWindow
        max = self.ones(n)
        min = self.ones(n)
        
        for i in range(n):
            for symbol, symbolData in self.symbolDataBySymbol.items():
                if symbolData.Ratios[-n:][i] > max[i]:
                    max[i] = symbolData.Ratios[-n:][i]
                if symbolData.Ratios[-n:][i] < min[i]:
                    min[i] = symbolData.Ratios[-n:][i]
                # algorithm.Log(symbolData.Ratios)
            
        # algorithm.Log(max)
        # algorithm.Log(min)
        
        max.append(1)
        min.append(1)
        best = ''
        worst = ''
        
        for symbol, symbolData in self.symbolDataBySymbol.items():
            
            # Issue...
            if not slice.Bars.ContainsKey(symbol.ToString()):
                return []
            
            ratio = slice[symbol].Close/slice[symbol].Open
            symbolData.Ratios.append(ratio)
            if ratio > max[-1]:
                best = symbol.Value
                max[-1] = ratio
            if ratio < min[-1]:
                worst = symbol.Value
                min[-1] = ratio
            # algorithm.Log(symbolData.Ratios)
        
        # algorithm.Log(max)
        # algorithm.Log(min)
        # algorithm.Log(best)
        # algorithm.Log(worst)
        
        cumulative_momentum_gains = 1.0
        cumulative_reversion_gains = 1.0
       
        for i in range(n):
            for symbol, symbolData in self.symbolDataBySymbol.items():
                # algorithm.Log('ratios: ' + str(symbolData.Ratios))
                # algorithm.Log('last n+1 ratios: ' + str(symbolData.Ratios[-n-1:]))
                if symbolData.Ratios[-n-1:][i] == max[i]:
                    cumulative_momentum_gains *= symbolData.Ratios[-n-1:][i+1]
                if symbolData.Ratios[-n-1:][i] == min[i]:
                    cumulative_reversion_gains *= symbolData.Ratios[-n-1:][i+1]
                # algorithm.Log('max[i]: ' + str(max[i]))
                # algorithm.Log('min[i]: ' + str(min[i]))
                # algorithm.Log('mom gains: ' + str(cumulative_momentum_gains))
                # algorithm.Log('rev gains: ' + str(cumulative_reversion_gains))
                # algorithm.Log('\n')
                
        # for symbol, symbolData in self.symbolDataBySymbol.items():
        #     algorithm.Log(symbolData.Ratios)
        algorithm.Log(cumulative_momentum_gains)
        algorithm.Log(cumulative_reversion_gains)
        
        if cumulative_momentum_gains > cumulative_reversion_gains:
            insights.append(Insight.Price(best, Expiry.EndOfDay, InsightDirection.Up))
            algorithm.Log("Momentum approach on " + best)
            self.buy_symbol = best
        else:
            insights.append(Insight.Price(worst, Expiry.EndOfDay, InsightDirection.Up))
            algorithm.Log("Reversion approach on " + worst)
            self.buy_symbol = worst
        
        return insights

    
    def OpenPositions(self):
        if self.algorithm.Portfolio.Invested:
            self.algorithm.Liquidate()
        self.algorithm.SetHoldings(self.buy_symbol, 1)
    
    
    def ClosePositions(self):        
        return
    

    def OnSecuritiesChanged(self, algorithm, changes):
        '''Event fired each time the we add/remove securities from the data feed
        Args:
            algorithm: The algorithm instance that experienced the change in securities
            changes: The security additions and removals from the algorithm'''
        
        for added in changes.AddedSecurities:
            symbol = added.Symbol
            symbolData = self.symbolDataBySymbol.get(added.Symbol)
            if symbolData is None:
                symbolData = SymbolData(added, algorithm, self.rollingWindow)
                self.symbolDataBySymbol[added.Symbol] = symbolData
            if self.schedule_symbol is None:
                self.algorithm = algorithm
                # Schedule event 1 minute before market close
                self.scheduled_event_open = algorithm.Schedule.On(
                    algorithm.DateRules.EveryDay(symbol), 
                    algorithm.TimeRules.AfterMarketOpen(symbol, 1), 
                    self.OpenPositions)
                self.scheduled_event_close = algorithm.Schedule.On(
                    algorithm.DateRules.EveryDay(symbol), 
                    algorithm.TimeRules.BeforeMarketClose(symbol, 1), 
                    self.ClosePositions)
                self.schedule_symbol = symbol
                        
                    

                
class SymbolData:
    '''Contains data specific to a symbol required by this model'''
    
    def __init__(self, security, algorithm, rollingWindow):
        self.Security = security
        self.Symbol = security.Symbol
        # preload ratios of last n days before yesterday given n = rollingWindow
        history = algorithm.History(
            [self.Symbol], rollingWindow+1, Resolution.Daily
            ).head(rollingWindow)
        # for open, close in zip(history['open'], history['close']):
        #     algorithm.Log(str(open) + ' ' + str(close) + ' ' + str(close/open))
        self.Ratios = [
            close/open for open, close in zip(history['open'], history['close'])
            ]