import pandas as pd
from universe import LargeLiquidUSEquities
from symbol_data import SymbolData

class ValueAndMomentumEverywhere(QCAlgorithm):
    
    def Initialize(self):
        
        self.SetStartDate(2016, 1, 1)
        self.SetEndDate(2021, 2, 1) 
        self.SetCash(1000000)
        self.UniverseSettings.Resolution = Resolution.Daily
        self.Settings.FreePortfolioValuePercentage = 0.2
        
        self.AddUniverseSelection(LargeLiquidUSEquities())
        self.symbol_data_by_symbol = {}
        
        self.rebalance_flag = True
        self.Schedule.On(self.DateRules.MonthStart(0), self.TimeRules.At(0, 0), self.rebalance)  

    
    def rebalance(self): 
        self.rebalance_flag = True
        
    def OnData(self, data):
        if not self.rebalance_flag:
            return
        
        # Gather value and momentum factor values
        df = pd.DataFrame()
        for symbol, symbol_data in self.symbol_data_by_symbol.items(): 
            has_data = self.CurrentSlice.ContainsKey(symbol) and self.CurrentSlice[symbol] is not None
            if not (symbol_data.IsReady and has_data):
                continue
            df.loc[str(symbol), 'Momentum'] = symbol_data.momentum
            
        # Ensure some constituents are ready
        if df.empty: 
            self.Debug('Consitituents not warm yet')
            for symbol in self.symbol_data_by_symbol.keys():
                self.Liquidate(symbol)
            return
        self.rebalance_flag = False
            
        # Rank the securities on their factor values
        rank = df['Momentum'].rank() 
        
        # Weight securities based on their factor rankings; Make dollar-neutral
        weight_by_symbol = rank - rank.mean() 
        
        # Scale weights down so the portfolio stays within leverage constraints
        weight_by_symbol /= (abs(weight_by_symbol).sum())

        # Place trades
        for symbol, weight in weight_by_symbol.iteritems():
            self.SetHoldings(symbol, weight)
                
    def OnSecuritiesChanged(self, changes):
        for security in changes.AddedSecurities:
            self.symbol_data_by_symbol[security.Symbol] = SymbolData(security, self)
        for security in changes.RemovedSecurities:
            self.Liquidate(security.Symbol)
            symbol_data = self.symbol_data_by_symbol.pop(security.Symbol, None)
            if symbol_data:
                symbol_data.dispose()

class SymbolData:
    # Momentum Indicator: 12month returns, excluding the most recent month
    
    def __init__(self, security, algorithm, momentum_lookback=12*22, momentum_delay=1*22):
        self.security = security
        self.symbol = security.Symbol
        self.algorithm = algorithm
        
        # Setup consolidator and indicators
        self.consolidator = algorithm.ResolveConsolidator(self.symbol, Resolution.Daily)
        self.mom = RateOfChange(momentum_lookback-momentum_delay)
        self.momentum_indicator = IndicatorExtensions.Of(Delay(momentum_delay), self.mom)
        algorithm.RegisterIndicator(self.symbol, self.mom, self.consolidator)
        
        # Warm up indicator
        history = algorithm.History(self.symbol, momentum_lookback + momentum_delay, Resolution.Daily)
        if history.empty or 'close' not in history.columns:
            return
        for time, close in history.loc[self.symbol].close.iteritems():
            self.mom.Update(time, close)
    
    @property
    def momentum(self):
        return self.momentum_indicator.Current.Value
    
    @property
    def IsReady(self):
        return self.momentum_indicator.IsReady

    def dispose(self):
        self.algorithm.SubscriptionManager.RemoveConsolidator(self.symbol, self.consolidator)

from Selection.FundamentalUniverseSelectionModel import FundamentalUniverseSelectionModel

class LargeLiquidUSEquities(FundamentalUniverseSelectionModel):
    def __init__(self, min_price = 1):
        self.month = -1
        self.min_price = min_price
        super().__init__(True)

    def SelectCoarse(self, algorithm, coarse):
        """
        Coarse universe selection is called each day at midnight.
        
        Input:
         - algorithm
            Algorithm instance running the backtest
         - coarse
            List of CoarseFundamental objects
            
        Returns the symbols that have fundamental data, excluding penny stocks.
        """
        # Refresh monthly
        if self.month == algorithm.Time.month:
            return Universe.Unchanged
        
        # Select securities with fundamental data (exclude penny stocks)
        selected = [c for c in coarse if c.HasFundamentalData and c.Price >= self.min_price]
        sorted_by_dollar_volume = sorted(selected, key=lambda x: x.DollarVolume, reverse=True)
        return [c.Symbol for c in sorted_by_dollar_volume[:500]]
    
        
    def SelectFine(self, algorithm, fine):
        """
        Fine universe selection is performed each day at midnight after `SelectCoarse`.
        
        Input:
         - algorithm
            Algorithm instance running the backtest
         - fine
            List of FineFundamental objects that result from `SelectCoarse` processing
        
        Returns a list of the largest symbols from SelectCoarse that make up 90% of the
        total market cap, exluding depositary receipts, REITs, and financial firms.
        """
        self.month = algorithm.Time.month
        
        # Restrict to common stock; Remove ADRs, REITs, and Financials
        selected = [f for f in fine if 
                        f.SecurityReference.SecurityType == 'ST00000001' and 
                        not f.SecurityReference.IsDepositaryReceipt and 
                        not f.CompanyReference.IsREIT and 
                        f.AssetClassification.MorningstarSectorCode != MorningstarSectorCode.FinancialServices]
        
        # Select the largest stocks
        top_market_cap = sorted(selected, key=lambda x: x.MarketCap, reverse=True)
        return [f.Symbol for f in top_market_cap[:50]]