# region imports
from AlgorithmImports import *

import itertools
# endregion
class VirtualYellowGreenLlama(QCAlgorithm):

    def initialize(self):
        self.set_start_date(2010, 1, 2)
        self.set_cash(10_000_000)
        self.settings.automatic_indicator_warm_up = True
        self.settings.minimum_order_margin_portfolio_percentage = 0.0001
        self.set_security_initializer(BrokerageModelSecurityInitializer(self.brokerage_model, FuncSecuritySeeder(self.get_last_known_prices)))
        self._beta_period = self.get_parameter('beta_period', 252)
        self._assets_per_industry = self.get_parameter('assets_per_industry', 20)
        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.add_universe(self._select_assets)
        self.schedule.on(self.date_rules.month_start(self._spy.symbol, 2), self.time_rules.midnight, self._rebalance)

    def _select_assets(self, fundamentals):
        # Select the most liquid assets of each industry group.
        fundamentals = sorted([f for f in fundamentals if f.asset_classification.morningstar_industry_group_code], key=lambda f: (f.asset_classification.morningstar_industry_group_code, f.dollar_volume))
        selected = []
        for _, industry_group_fundamentals in itertools.groupby(fundamentals, lambda f: f.asset_classification.morningstar_industry_group_code):
            selected.extend(list(industry_group_fundamentals)[-self._assets_per_industry:])  # We already sorted by dollar volume above.
        # Get the absolute beta of each asset.
        self._beta_by_symbol = self._beta([f.symbol for f in selected]).abs()
        # Get the median beta.
        median_beta = np.median(self._beta_by_symbol.values)
        # Select the assets in each industry have that a beta below the median.
        weights_by_industry = {}
        symbols = []
        for industry_code, industry_assets in itertools.groupby(selected, lambda f: f.asset_classification.morningstar_industry_group_code):
            # Get the beta of each asset in the industry.
            industry_beta_by_symbol = self._beta_by_symbol[[f.symbol for f in industry_assets if f.symbol in self._beta_by_symbol]]
            # Select assets with a beta below the median.
            low_betas = industry_beta_by_symbol[industry_beta_by_symbol < median_beta]
            if low_betas.empty:
                continue
            symbols.extend(list(low_betas.index))
            # Create weights for the assets in this industry (Beta-weighted weights).
            beta_ranks = low_betas.sort_values().rank(method='first', ascending=False) # Higher rank = Lower beta => Larger positive position.
            weights_by_industry[industry_code] = beta_ranks / beta_ranks.sum()

        # Create the portfolio targets. Give equal weight to each industry. Liquidate assets we no longer want.
        self._targets = [PortfolioTarget(symbol, 0) for symbol, holding in self.portfolio.items() if holding.invested and symbol not in symbols]
        for industry_assets in weights_by_industry.values():
            self._targets.extend([PortfolioTarget(symbol, weight/len(weights_by_industry)) for symbol, weight in industry_assets.items()])
        return symbols
        
    def _beta(self, symbols): # Source: https://stackoverflow.com/questions/39501277/efficient-python-pandas-stock-beta-calculation-on-many-dataframes
        returns = self.history([self._spy.symbol] + symbols, self._beta_period, Resolution.DAILY, fill_forward=False).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):
        if self._targets:
            # Rebalance the portfolio.
            self.set_holdings(self._targets)
            self._targets = []