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 = data.ContainsKey(symbol) and data[symbol] is not None
            if not (symbol_data.IsReady and has_data):
                continue
            df.loc[str(symbol), 'Value'] = symbol_data.value
            
        # 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['Value'].rank(ascending=False)
        
        # 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)
        for security in changes.RemovedSecurities:
            self.Liquidate(security.Symbol)
            self.symbol_data_by_symbol.pop(security.Symbol, None)

class SymbolData:
    # Value Indicator: Book-to-market ratio
    
    def __init__(self, security):
        self.security = security

    @property
    def value(self):
        return 1 / self.security.Fundamentals.ValuationRatios.PBRatio
    
    @property
    def IsReady(self):
        return self.security.Fundamentals.ValuationRatios.PBRatio != 0

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]]