import pandas as pd
from us_equity import USEquity
from global_equity import GlobalEquity

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
        
        # Initialize universes
        GlobalEquity.Initialize(self)
        USEquity.Initialize(self)
        
        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
        
        # Rank Securities and trade: 
        allConstituents = [USEquity.Constituents, GlobalEquity.Constituents] 
        
        # Calculate portfolio weights of the securities
        target = pd.Series()
        
        for constituents in allConstituents:
            
            # Gather value and momentum factor values
            df = pd.DataFrame()
            for symbol, symbol_data in constituents.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), 'Value'] = symbol_data.value
                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 constituents.keys():
                    self.Liquidate(symbol)
                continue
                
            # Rank the securities on their factor values
            rank = df.rank(axis=0) 
            
            # Weight securities based on their factor rankings; Make dollar-neutral
            weight_by_symbol = (rank - rank.mean()).mean(axis=1)
            if all(weight_by_symbol.values == 0):
                for symbol in constituents.keys():
                    self.Liquidate(symbol)
                continue
            
            # Scale weights down so the portfolio stays within leverage constraints
            weight_by_symbol /= abs(weight_by_symbol).sum()
    
            # Scale weights further down to have the asset classes equally allocated
            weight_by_symbol /= len(allConstituents)
            
            target = target.append(weight_by_symbol)
        
        if not target.empty:
            self.rebalance_flag = False
            
        for symbol, weight in target.iteritems():
            self.SetHoldings(symbol, weight)
                
    def OnSecuritiesChanged(self, changes):
        USEquity.OnSecuritiesChanged(changes)
        GlobalEquity.OnSecuritiesChanged(changes)

from dateutil.relativedelta import relativedelta
import pandas as pd

class SymbolData:
    # Value Indicator:    Negative of past 5 year return
    # Momentum Indicator: 12month returns, excluding the most recent month
        
    def setup(self, algorithm, symbol, momentum_lookback=12*22, momentum_delay=1*22, value_length=5*12*22):
        self.symbol = symbol
        self.value_length = value_length
        
        # Setup consolidator and indicators
        self.consolidator = algorithm.ResolveConsolidator(self.symbol, Resolution.Daily)
        self.value_indicator = RateOfChange(value_length)
        self.mom = RateOfChange(momentum_lookback-momentum_delay)
        self.momentum_indicator = IndicatorExtensions.Of(Delay(momentum_delay), self.mom)
        algorithm.RegisterIndicator(self.symbol, self.value_indicator, self.consolidator)
        algorithm.RegisterIndicator(self.symbol, self.mom, self.consolidator)
        
        # Warm up indicators
        history = algorithm.History(self.symbol, value_length, Resolution.Daily)
        if history.empty or 'close' not in history.columns:
            return
        for time, close in history.loc[self.symbol].close.iteritems():
            self.value_indicator.Update(time, close)
            self.mom.Update(time, close)


    @property
    def value(self):
        return -self.value_indicator.Current.Value
        
    @property
    def momentum(self):
        return self.momentum_indicator.Current.Value
    
    @property
    def IsReady(self):
        return self.momentum_indicator.IsReady and self.value_indicator.IsReady
        
    def dispose(self):
        self.algorithm.SubscriptionManager.RemoveConsolidator(self.symbol, self.consolidator)

from symbol_data import SymbolData

class GlobalEquity(SymbolData):
    # Notes:
    #   The paper aggregates the individual stocks’ book-to-market ratios to compute the value indicator of
    #   each country. To trade each index, the authors used futures contracts. 
    #   The paper discusses how these aggregated ratios correlate closely with the negative of 
    #   the past 5 year return of each index. We use this metric as our value indicator and trade
    #   each country index via MSCI ETFs.
    
    algorithm = None
    Symbols = []
    Constituents = {}
    
    @staticmethod
    def Initialize(algo):
        
        GlobalEquity.algorithm = algo
        # Global Equity Indices
        tickers = ['EWA',  # Australia
                   'EWO',  # Austria 
                   'EWK',  # Belgium
                   'EWC',  # Canada
                   'EDEN', # Denmark
                   'EWQ',  # France
                   'EWG',  # Germany
                   'EWH',  # Hong Kong
                   'EWI',  # Italy
                   'EWJ',  # Japan
                   'EWN',  # Netherlands
                   'ENOR', # Norway
                   'PGAL', # Portugal
                   'EWP',  # Spain
                   'EWD',  # Sweden
                   'EWL',  # Switzerland
                   'EWU',  # U.K.
                   'EUSA'  # U.S.
        ]
        GlobalEquity.Symbols = [Symbol.Create(t, SecurityType.Equity, "usa") for t in tickers]
        algo.AddUniverseSelection( ManualUniverseSelectionModel(GlobalEquity.Symbols) )
    
    @staticmethod
    def OnSecuritiesChanged(changes):
        for security in changes.AddedSecurities:
            symbol = security.Symbol
            if symbol not in GlobalEquity.Symbols:
                continue
            GlobalEquity.Constituents[symbol] = GlobalEquity(symbol)
            
    def __init__(self, symbol):
        super().setup(GlobalEquity.algorithm, symbol)

from Selection.FundamentalUniverseSelectionModel import FundamentalUniverseSelectionModel
from symbol_data import SymbolData

class USEquity(SymbolData):
    # Value Indicator:    Book-to-market value
    # Momentum Indicator: 12month returns, excluding the most recent month
    
    algorithm = None
    Symbols = []
    Constituents = {}
    
    @staticmethod
    def Initialize(algo):
        # algorithm.Add Global tickers 
        USEquity.algorithm = algo 
        algo.AddUniverseSelection(LargeLiquidUSEquities())
        
    @staticmethod
    def OnSecuritiesChanged(changes):
        for security in changes.AddedSecurities:
            if security.Symbol not in USEquity.Symbols:
                continue
            USEquity.Constituents[security.Symbol] = USEquity(security)
        for security in changes.RemovedSecurities:
            USEquity.algorithm.Liquidate(security.Symbol)
            us_equity = USEquity.Constituents.pop(security.Symbol, None)
            if us_equity:
                us_equity.dispose()
                
    def __init__(self, security):
        self.security = security
        super().setup(USEquity.algorithm, security.Symbol)
    
    @property
    def value(self):
        return 1 / self.security.Fundamentals.ValuationRatios.PBRatio
    
    @property
    def IsReady(self):
        return self.momentum_indicator.IsReady and self.security.Fundamentals is not None and self.security.Fundamentals.ValuationRatios.PBRatio != 0
    
    def dispose(self):
        USEquity.algorithm.SubscriptionManager.RemoveConsolidator(self.symbol, self.consolidator)


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 in decreasing order by market cap until we have n% of the total market cap
        top_market_cap = sorted(selected, key=lambda x: x.MarketCap, reverse=True)
        USEquity.Symbols = [f.Symbol for f in top_market_cap[:100]]
        return USEquity.Symbols