from AlgorithmImports import * 


class CrossSectionalMomentum(QCAlgorithm):

    def Initialize(self):

        self.SetStartDate(2019, 1, 1)    
        self.SetEndDate(2020, 1, 1)      
        self.SetCash(1000000)
        self.SetBrokerageModel(AlphaStreamsBrokerageModel())
        
        self.UniverseSettings.Resolution = Resolution.Daily
        
        self.AddUniverse(self.CoarseSelection)
        
        self.SetAlpha(ReversalFakeoutAlpha())
        
        self.SetPortfolioConstruction(EqualWeightingPortfolioConstructionModel())
        self.SetRiskManagement(TrailingStopRiskManagementModel(0.02))
        self.SetExecution(ImmediateExecutionModel())
        
    def CoarseSelection(self, coarse):
        sortedCoarse = sorted(coarse, key=lambda c:c.DollarVolume, reverse=True)
        return [c.Symbol for c in sortedCoarse][:1000]
        


class ReversalFakeoutAlpha(AlphaModel):
    
    def __init__(self):
        self.symbols = {}
        self.lastWeek = -1
        
    def Update(self, algorithm, data):
        
        insights = []
        
        # Update daily rolling windows with daily data
        for symbol, symbolData in self.symbols.items():
            if data.ContainsKey(symbol) and data[symbol] != None:
                symbolData.dailyWindow.Add(data[symbol])
        
        # Retrieve week number from datetime
        thisWeek = algorithm.Time.isocalendar()[1]
        
        # Only rebalance once a week
        if self.lastWeek == thisWeek:
            return insights
        
        self.lastWeek = thisWeek
        
        # Retrieve all symbols that are ready
        symbols = [symbol for symbol, symbolData in self.symbols.items() if symbolData.IsReady]
        
        # Sort by annualized returns in descending order
        sortedByReturns = sorted(symbols, key=lambda s: self.symbols[s].AnnualizedReturn, reverse=True)
        
        winningSymbols = sortedByReturns[:100]
        losingSymbols = sortedByReturns[-100:]
        
        # Sort symbols by continuity and filter for symbols with recent reversals
        longContinuity = [symbol for symbol in sorted(winningSymbols, key=lambda s: self.symbols[s].Continuity, reverse=False) if self.symbols[symbol].RecentDownTrend]
        shortContinuity = [symbol for symbol in sorted(losingSymbols, key=lambda s: self.symbols[s].Continuity, reverse=True) if self.symbols[symbol].RecentUpTrend]
        
        # Create insights for the top 5 most continous symbols which are in a recent reversal
        insights += [Insight.Price(symbol, timedelta(days = 5), InsightDirection.Up) for symbol in longContinuity[:5]]
        insights += [Insight.Price(symbol, timedelta(days = 5), InsightDirection.Down) for symbol in shortContinuity[:5]]
        
        return insights
        
        
    def OnSecuritiesChanged(self, algorithm, changes):
        
        for security in changes.AddedSecurities:
            symbol = security.Symbol
            if symbol not in self.symbols:
                self.symbols[symbol] = SymbolData(algorithm, symbol)
        
        for security in changes.RemovedSecurities:
            symbol = security.Symbol
            if symbol in self.symbols:
                # Remove consolidators from algorithm and remove symbol from dictionary
                algorithm.SubscriptionManager.RemoveConsolidator(symbol, self.symbols[symbol].monthlyConsolidator)
                algorithm.SubscriptionManager.RemoveConsolidator(symbol, self.symbols[symbol].dailyConsolidator)
                self.symbols.pop(symbol, None)


class SymbolData:
    
    def __init__(self, algorithm, symbol):
        self.algorithm = algorithm
        self.symbol = symbol
        
        # Define daily and monthly rolling windows
        self.monthlyWindow = RollingWindow[TradeBar](13)
        self.dailyWindow = RollingWindow[TradeBar](280)
        
        # Define daily and monthly consolidators
        self.monthlyConsolidator = algorithm.Consolidate(symbol, Calendar.Monthly, self.OnMonthlyData)
        self.dailyConsolidator = TradeBarConsolidator(timedelta(days = 1))
        
        # Register daily consolistor to algorithm
        algorithm.SubscriptionManager.AddConsolidator(symbol, self.dailyConsolidator)
        
        # Define and register ADX indicator
        self.adxThreshold = 25
        self.adx = AverageDirectionalIndex(20)
        algorithm.RegisterIndicator(symbol, self.adx, self.dailyConsolidator)
        
                # Use historical data to warmup rolling windows, consolidators, and indicators
        history = algorithm.history(symbol, 280, Resolution.DAILY)
        if not history.empty:
            for index, row in history.iterrows():
                tbar = TradeBar()
                tbar.Symbol = symbol
                tbar.Time = index[1] if isinstance(index, tuple) else index
                tbar.Open = row.open
                tbar.High = row.high
                tbar.Low = row.low
                tbar.Close = row.close
                tbar.Volume = row.volume if hasattr(row, 'volume') else 0
                self.dailyWindow.Add(tbar)
                self.monthlyConsolidator.Update(tbar)
                self.adx.Update(tbar)
        
        
    def OnMonthlyData(self, bar):
        """Adds monthly bars to monthly rolling window"""
        self.monthlyWindow.Add(bar)
    
    @property
    def Continuity(self):
        """Returns the difference between losing days and winning days as a percentage of trading days"""
        positives = 0
        negatives = 0
        for bar in self.dailyWindow:
            dreturn = (bar.Close-bar.Open/bar.Open)
            if dreturn > 0:
                positives += 1
            else:
                negatives += 1
        return (negatives - positives)/(negatives + positives)
        
    @property
    def AnnualizedReturn(self):
        """Returns the 12 month compounded monthly return over a 13 month lookback period
        skipping the latest month."""
        returns = []
        
        for bar in self.monthlyWindow:
            monthlyReturn = (bar.Close/bar.Open)
            returns.append(monthlyReturn)
        returns.pop(0)
        return np.prod(returns) - 1
    
    
    
    
    @property
    def RecentDownTrend(self):
        """Returns true if the ADX is above a given threshold and DX+ is lower than DX-"""
        return self.adx.Current.Value > self.adxThreshold and \
        self.adx.PositiveDirectionalIndex.Current.Value < self.adx.NegativeDirectionalIndex.Current.Value
        
    @property
    def RecentUpTrend(self):
        """Returns true if the ADX is above a given threshold and DX+ is higher than DX-"""
        return self.adx.Current.Value > self.adxThreshold and \
        self.adx.PositiveDirectionalIndex.Current.Value > self.adx.NegativeDirectionalIndex.Current.Value
     
    @property
    def IsReady(self):
        """Returns true if all the rolling windows and indicators are ready: have been updated
        with enough inputs to yield valid values"""
        return self.monthlyWindow.IsReady and self.dailyWindow.IsReady and self.adx.IsReady