# region imports
from AlgorithmImports import *
# endregion

class VirtualYellowGreenLlama(QCAlgorithm):

    def initialize(self):
        self.set_start_date(2010, 1, 1)
        self.set_cash(10_000_000)
        self.settings.automatic_indicator_warm_up = True
        self._beta_period = self.get_parameter('beta_period', 252)
        self._spy = self.add_equity('SPY', Resolution.DAILY)
        self.universe_settings.resolution = Resolution.DAILY
        self.universe_settings.schedule.on(self.date_rules.month_start(self._spy.symbol))
        self._universe = self.add_universe(self._select_assets)
        self.schedule.on(self.date_rules.month_start(self._spy.symbol), self.time_rules.at(0, 1), self._rebalance)

    def _select_assets(self, fundamentals):
        fundamentals = sorted(fundamentals, key=lambda f: f.dollar_volume)[-2_500:]
        return [f.symbol for f in fundamentals if f.asset_classification.morningstar_industry_code == MorningstarIndustryCode.SOFTWARE_APPLICATION]
        
    def _beta(self, securities): # Source: https://stackoverflow.com/questions/39501277/efficient-python-pandas-stock-beta-calculation-on-many-dataframes
        symbols = [s.symbol for s in securities]
        returns = self.history([self._spy.symbol] + symbols, self._beta_period, Resolution.DAILY).close.unstack(0).dropna(axis=1).pct_change().dropna()
        symbols = [s for s in symbols if s in returns.columns]
        df = returns[[self._spy.symbol] + symbols]
        # first column is the market
        X = df.values[:, [0]]
        # prepend a column of ones for the intercept
        X = np.concatenate([np.ones_like(X), X], axis=1)
        # matrix algebra
        b = np.linalg.pinv(X.T.dot(X)).dot(X.T).dot(df.values[:, 1:])
        return pd.Series(b[1], df.columns[1:], name='Beta')

    def _rebalance(self):
        securities = [self.securities[symbol] for symbol in self._universe.selected]
        self.plot('Securities', 'Original Size', len(securities))
        securities = [s for s in securities if s.price]
        self.plot('Securities', 'Tradable Size', len(securities))
        if not securities:
            return
        # Get the beta of each asset.
        beta_by_symbol = self._beta(securities)
        # Calculate the weights of the portfolio.
        # Step 1: Weigh the assets by their beta ranks, relative to the universe median.
        #   - Lower-beta security have larger weight in the low-beta portfolio and higher-beta securities have larger weights in the high beta portfolio.
        beta_rank_by_symbol = beta_by_symbol.sort_values().rank(method='first', ascending=False) # Higher rank number = Lower Beta (so more positive portfolio weight).
        rank_delta_from_mean = beta_rank_by_symbol - beta_rank_by_symbol.mean()
        normalizing_constant = rank_delta_from_mean.abs().sum() / 2 # Make the weights in the long and short sides of the portfolio sum to 1/-1.
        weight_by_symbol = rank_delta_from_mean / normalizing_constant
        
        # Step 2: Scale the weights of each side of the portfolio to a -1 beta, making the portfolio market-netural.
        try:
            self._scale_weights_to_negative_1_beta(beta_by_symbol, weight_by_symbol, lambda x: x > 0)
            self._scale_weights_to_negative_1_beta(beta_by_symbol, weight_by_symbol, lambda x: x < 0)
        except: # If the portfolio Beta for either side of the portfolio is >= 0, don't rebalance.
            return
        long_portfolio_weights = weight_by_symbol[weight_by_symbol > 0]
        long_portfolio_beta = beta_by_symbol[long_portfolio_weights.index].dot(long_portfolio_weights)
        short_portfolio_weights = weight_by_symbol[weight_by_symbol < 0]
        short_portfolio_beta = beta_by_symbol[short_portfolio_weights.index].dot(short_portfolio_weights)
        self.plot('Portfolio Beta', 'Long', long_portfolio_beta)
        self.plot('Portfolio Beta', 'Short', short_portfolio_beta)
        
        # Step 3: Scale all the weights to keep gross leverage at 1. 
        weight_sum_by_bias = {
            'long' : weight_by_symbol[weight_by_symbol > 0].sum(),
            'short': weight_by_symbol[weight_by_symbol < 0].sum()
        } 
        self.plot('Weight Sum', 'Long (Before)', weight_sum_by_bias['long'])
        self.plot('Weight Sum', 'Short (Before)', weight_sum_by_bias['short'])
        scaling_factor = 1 / sum([abs(weight) for weight in weight_sum_by_bias.values()])
        weight_by_symbol *= scaling_factor
        self.plot('Weight Sum', 'Long (After)', weight_by_symbol[weight_by_symbol > 0].sum())
        self.plot('Weight Sum', 'Short (After)', weight_by_symbol[weight_by_symbol < 0].sum())

        # Rebalance the portfolio.
        self.set_holdings([PortfolioTarget(symbol, weight) for symbol, weight in weight_by_symbol.items()], True)

    def _scale_weights_to_negative_1_beta(self, beta_by_symbol, weight_by_symbol, selection_function):
        portfolio_weights = weight_by_symbol[selection_function(weight_by_symbol)]
        portfolio_beta = beta_by_symbol[portfolio_weights.index].dot(portfolio_weights)
        if portfolio_beta >= 0: # Beta should be negative since we're betting against beta
            self.log(f'{self.time} Portfolio beta >= 0: {portfolio_beta}')
            raise Exception(f"Portfolio beta >= 0: {portfolio_beta}")
        weight_by_symbol[selection_function(weight_by_symbol)] /= -portfolio_beta # We don't want to flip the sign of the weights, so divide by the negative value.