#region imports
from AlgorithmImports import *
#endregion

class LongShortAlphaModel(AlphaModel):
    def __init__(self, algorithm):
        self.algorithm = algorithm
        self.rebalance_time = datetime.min
    
    def update(self, algorithm, data):
        # Do nothing until next rebalance
        if algorithm.time < self.rebalance_time:
            return []

        insights = []

        # Enter long positions
        for symbol in self.algorithm.long_symbols:
            insights.append(Insight.price(symbol, timedelta(91), InsightDirection.UP))
        # Enter short positions
        for symbol in self.algorithm.short_symbols:
            insights.append(Insight.price(symbol, timedelta(91), InsightDirection.DOWN))

        # Set next rebalance time
        self.rebalance_time = Expiry.EndOfMonth(algorithm.time)

        return insights

#region imports
from AlgorithmImports import *
from selection import ReturnAndEarningsUniverseSelectionModel
from alpha import LongShortAlphaModel
#endregion

class PriceEarningsMomentumAlgorithm(QCAlgorithm):
    '''
    A stock momentum strategy based on quarterly returns and earning growth.
    Paper: http://papers.ssrn.com/sol3/papers.cfm?abstract_id=299107
    Online copy: https://www.trendrating.com/wp-content/uploads/dlm_uploads/2019/03/momentum.pdf
    '''
    def initialize(self):
        self.set_start_date(2023, 3, 1)
        self.set_end_date(2024, 3, 1)
        self.set_cash(100000)            # Set Strategy Cash

        self.set_brokerage_model(BrokerageName.INTERACTIVE_BROKERS_BROKERAGE, AccountType.MARGIN)
        self.set_security_initializer(BrokerageModelSecurityInitializer(self.brokerage_model, FuncSecuritySeeder(self.get_last_known_prices)))

        num_of_coarse = 100          # Number of coarse selected universe
        self.long_symbols = []
        self.short_symbols = []

        self.universe_settings.data_normalization_mode = DataNormalizationMode.RAW
        self.universe_settings.resolution = Resolution.MINUTE          # Resolution setting of universe selection
        self.universe_settings.schedule.on(self.date_rules.month_start())
        self.universe_model = ReturnAndEarningsUniverseSelectionModel(self, num_of_coarse)
        self.set_universe_selection(self.universe_model)   # Coarse and Fine Universe Selection
        
        self.add_alpha(LongShortAlphaModel(self))
        self.set_portfolio_construction(EqualWeightingPortfolioConstructionModel(Expiry.END_OF_WEEK))
        self.add_risk_management(MaximumDrawdownPercentPerSecurity(0.2))

#region imports
from AlgorithmImports import *
#endregion

class ReturnAndEarningsUniverseSelectionModel(FundamentalUniverseSelectionModel):
    """
    This universe selection model refreshes monthly to contain US securities in the asset management industry.
    """
    def __init__(self, algorithm, numOfCoarse):
        self.algorithm = algorithm
        self.num_of_coarse = numOfCoarse
        self.eps_by_symbol = {}           # Contains RollingWindow objects of EPS for every stock

        super().__init__(self.select)

    def select(self, fundamental):
        """
        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.
        """            
        # Sort the equities (prices > 5) by Dollar Volume descendingly
        # # Pick the top 100 liquid equities as the coarse-selected universe
        selectedByDollarVolume = sorted([x for x in fundamental if x.price > 5 and x.has_fundamental_data], 
                                        key = lambda x: x.dollar_volume, reverse = True)[:self.num_of_coarse]

        symbols = [x.symbol for x in selectedByDollarVolume]

        # Get the quarterly returns for each symbol
        history = self.algorithm.history(symbols, timedelta(91), Resolution.DAILY, data_normalization_mode=DataNormalizationMode.SCALED_RAW)
        history = history.drop_duplicates().close.unstack(level = 0)
        rankByQuarterReturn = self.get_quarterly_return(history)

        # Get the earning growth for each symbol
        rankByEarningGrowth = self.get_earning_growth(selectedByDollarVolume) 

        # Get the sum of rank for each symbol and pick the top ones to long and the bottom ones to short
        rankSumBySymbol = {key: rankByQuarterReturn.get(key, 0) + rankByEarningGrowth.get(key, 0) 
                                for key in set(rankByQuarterReturn) | set(rankByEarningGrowth)}

        # Get 10 symbols to long and short respectively
        sortedDict = sorted(rankSumBySymbol.items(), key = lambda x: x[1], reverse = True)
        self.algorithm.long_symbols = [Symbol(SecurityIdentifier.parse(x[0]), str(x[0]).split(' ')[0]) for x in sortedDict[:5]]
        self.algorithm.short_symbols = [Symbol(SecurityIdentifier.parse(x[0]), str(x[0]).split(' ')[0]) for x in sortedDict[-5:]]
        selected = self.algorithm.long_symbols + self.algorithm.short_symbols

        return selected

    def get_quarterly_return(self, history):
        '''
        Get the rank of securities based on their quarterly return from historical close prices
        Return: dictionary
        '''
        # Get quarterly returns for all symbols
        # (The first row divided by the last row)
        returns = history.iloc[0] / history.iloc[-1]

        # Transform them to dictionary structure
        returns = returns.to_dict()

        # Get the rank of the returns (key: symbol; value: rank)
        # (The symbol with the 1st quarterly return ranks the 1st, etc.)
        ranked = sorted(returns, key = returns.get, reverse = True)
        return {symbol: rank for rank, symbol in enumerate(ranked, 1)}

    def get_earning_growth(self, fine):
        '''
        Get the rank of securities based on their EPS growth
        Return: dictionary
        '''

        # Earning Growth by symbol
        egBySymbol = {}
        for stock in fine:

            # Select the securities with EPS (> 0)
            if stock.earning_reports.basic_eps.three_months == 0:
                continue

            # Add the symbol in the dict if not exist
            if not stock.symbol in self.eps_by_symbol:
                self.eps_by_symbol[stock.symbol] = RollingWindow[float](2)

            # Update the rolling window for each stock
            self.eps_by_symbol[stock.symbol].add(stock.earning_reports.basic_eps.three_months)

            # If the rolling window is ready
            if self.eps_by_symbol[stock.symbol].is_ready:
                rw = self.eps_by_symbol[stock.symbol]
                # Caculate the Earning Growth
                egBySymbol[stock.symbol] = (rw[0] - rw[1]) / rw[1]

        # Get the rank of the Earning Growth
        ranked = sorted(egBySymbol, key = egBySymbol.get, reverse = True)
        return {symbol: rank for rank, symbol in enumerate(ranked, 1)}