# region imports
from AlgorithmImports import *

import itertools
# endregion

class VirtualYellowGreenLlama(QCAlgorithm):

    _blocked_assets = [
        'AMMA WEHWVFZT6VZ9', # Halted for several months in 2020, causes trading issues. See https://finance.yahoo.com/news/nasdaq-halts-scworx-corp-135720862.html and  https://www.stocktitan.net/news/WORX/sc-worx-to-resume-trading-on-nasdaq-monday-august-10-doocrwwrw9f6.html
        'PAY T8834TOLIRFP',  # https://www.quantconnect.com/datasets/issue/16989
        'HCP R735QTJ8XC9X',  # https://www.quantconnect.com/datasets/issue/17042
        'TOPT T0KDYN9C3IHX', # https://www.quantconnect.com/datasets/issue/15180
        'VSCI R735QTJ8XC9X', # https://www.quantconnect.com/datasets/issue/18448
        'IPDN VEN1SVFIVSKL',
    ]

    def initialize(self):
        self.set_start_date(2018, 1, 1)
        self.set_cash(10_000_000)
        self.settings.automatic_indicator_warm_up = True
        self._universe_size = 1_000
        self._beta_period = self.get_parameter('beta_period', 6) * 21
        self._assets_per_industry = self.get_parameter('assets_per_industry', 3)
        self._spy = self.add_equity('SPY', Resolution.DAILY)
        self.universe_settings.resolution = Resolution.DAILY
        self.universe_settings.schedule.on(self.date_rules.week_start(self._spy.symbol))
        self._universe = self.add_universe(self._select_assets)
        self.schedule.on(self.date_rules.week_start(self._spy.symbol), self.time_rules.at(0, 1), self._rebalance)

    def _select_assets(self, fundamentals):
        return [f.symbol for f in fundamentals if f.asset_classification.morningstar_industry_code and str(f.symbol.id) not in self._blocked_assets]
        # Narrow the universe to the most liquid for now and remove blocked assets.
        #return [f.symbol for f in sorted([f for f in fundamentals if f.asset_classification.morningstar_industry_code and str(f.symbol.id) not in self._blocked_assets], key=lambda f: f.dollar_volume)[-self._universe_size:]]
        
    def on_securities_changed(self, changes):
        # Create a Beta indicator for each asset that enters the universe.
        for security in changes.added_securities:
            security.beta = self.b(security.symbol, self._spy.symbol, self._beta_period)
        
    def _rebalance(self):
        # Select securities with negative beta and sort them by their beta values.
        securities = [self.securities[symbol] for symbol in self._universe.selected]
        securities = sorted([s for s in securities if s.price and s.beta.is_ready], key=lambda s: (s.fundamentals.asset_classification.morningstar_industry_code, s.beta.current.value))
        # Group assets by their industry.
        trades_by_industry = {}
        for industry_code, industry_securities in itertools.groupby(securities, lambda s: s.fundamentals.asset_classification.morningstar_industry_code):
            industry_securities = list(industry_securities)
            positive_betas = [s for s in industry_securities if s.beta.current.value > 0]
            negative_betas = [s for s in industry_securities if s.beta.current.value < 0]
            assets_per_leg = min(len(positive_betas), len(negative_betas), self._assets_per_industry)
            if assets_per_leg:
                trades_by_industry[industry_code] = {'short': negative_betas[:assets_per_leg], 'long': positive_betas[-assets_per_leg:]}
        # Give an equal weight to each industry and give an equal weight to each asset in each industry.
        targets = []
        for securities_by_trade_bias in trades_by_industry.values():
            for trade_bias, industry_securities in securities_by_trade_bias.items():
                weight = (1 if trade_bias == 'long' else -1) / len(trades_by_industry) / len(industry_securities) / 1.5
                for security in industry_securities:
                    targets.append(PortfolioTarget(security.symbol, weight))
        # Rebalance the portfolio. 
        self.set_holdings(targets, True)