#region imports
from AlgorithmImports import *
#endregion

class AssetGrowth(QCAlgorithm):
    def initialize(self):
        self.set_start_date(self.end_date-timedelta(365))
        self.set_cash(1_000_000)
        self.settings.seed_initial_prices = True
        self.settings.minimum_order_margin_portfolio_percentage = 0

        date_rule = self.date_rules.month_end("SPY")
        self.universe_settings.schedule.on(date_rule)
        self.universe_settings.resolution = Resolution.DAILY
        self.add_universe(self._select)

        self._total_assets_by_symbol = {}
        self._targets = []
        self.set_warmup(252, Resolution.DAILY)

    def _select(self, fundamental):
        # Rebalancing should occur in July
        if self.time.month % 6 != 0:
            return Universe.UNCHANGED

        total_assets_by_symbol = {
                x.symbol: x.financial_statements.balance_sheet.total_assets.value
                for x in fundamental 
                if (x.financial_statements.balance_sheet.total_assets.value > 0 and
                    x.security_reference.exchange_id in ['NYS', 'NAS', 'ASE'] and 
                    x.company_reference.industry_template_code not in ["B", "I"])
        }

        if self.is_warming_up:
            self._total_assets_by_symbol = total_assets_by_symbol
            return Universe.UNCHANGED

        # Calculate and sort the stocks on asset growth
        def calc_growth(symbol, current):
            previous = self._total_assets_by_symbol.get(symbol, current)
            return current / previous - 1

        growth_by_symbol = {symbol: calc_growth(symbol, x) for symbol, x in total_assets_by_symbol.items()}
        sorted_by_growth = sorted(growth_by_symbol.items(), key=lambda x: x[1])

        portfolio_size = int(len(growth_by_symbol)/10)
        weight = .5 / portfolio_size

        # Pick the upper decile to short and the lower decile to long
        self._targets = [PortfolioTarget(s, -weight) for s,g in sorted_by_growth[-portfolio_size:]]
        self._targets += [PortfolioTarget(s, weight) for s,g in sorted_by_growth[:portfolio_size]]

        # Save the fundamental data for the next year's analysis
        self._total_assets_by_symbol = total_assets_by_symbol
        return [t.symbol for t in self._targets]

    def on_data(self, data):
        if not self._targets:
            return
        self.set_holdings(self._targets, True)
        self._targets = []