#region imports
from AlgorithmImports import *
#endregion


class TwelveMonthCycle(QCAlgorithm):

    def initialize(self):
        self.set_start_date(self.end_date - timedelta(5*365))
        self.set_cash(100000) 
        self.settings.seed_initial_prices = True
        # Add a universe of US Equities.
        self._symbol_data_by_symbol = {}
        self.universe_settings.resolution = Resolution.DAILY
        self._date_rule = self.date_rules.month_start("SPY")
        self.universe_settings.schedule.on(self._date_rule)
        self.add_universe(self._select_assets)   
        # Warm up the indicators we need.
        self.set_warm_up(timedelta(450))
        
    def _select_assets(self, fundamentals):
        # Update the monthly returns of each asset.
        selected = []
        for f in fundamentals:
            # Only consider stocks with fundamental data listed on NYS or ASE.
            if (not f.has_fundamental_data or
                f.security_reference.exchange_id not in ['NYS', 'ASE']):
                continue
            if f.symbol not in self._symbol_data_by_symbol:
                self._symbol_data_by_symbol[f.symbol] = SymbolData()
            symbol_data = self._symbol_data_by_symbol[f.symbol]
            symbol_data.roc.update(f.end_time, f.price)
            if symbol_data.delayed_roc.is_ready:
                selected.append(f)
        # Drop 90% of stocks that are the smallest.
        size = int(0.1*len(selected))
        selected = sorted(selected, key=lambda f: f.market_cap)[-size:]
        # Split into long and short sides based on the delayed_roc factor.
        assets_per_side = int(len(selected)/10)
        sorted_by_roc = sorted(
            selected, 
            key=lambda f: self._symbol_data_by_symbol[f.symbol].delayed_roc
        )
        self._longs = [f.symbol for f in sorted_by_roc[-assets_per_side:]]
        self._shorts = [f.symbol for f in sorted_by_roc[:assets_per_side]]
        return self._longs + self._shorts

    def on_warmup_finished(self):
        # Add a Scheduled event to rebalance the portfolio monthly.
        time_rule = self.time_rules.at(8, 0)
        self.schedule.on(self._date_rule, time_rule, self._rebalance)
        # Rebalance the portfolio today too.
        if self.live_mode:
            self._rebalance()
        else:
            self.schedule.on(self.date_rules.today, time_rule, self._rebalance)
    
    def _rebalance(self):
        # Drop assets that have no price yet to avoid trading errors.
        longs = [symbol for symbol in self._longs if self.securities[symbol].price]
        shorts = [symbol for symbol in self._shorts if self.securities[symbol].price]
        # Form a long-short portfolio.
        targets = [PortfolioTarget(symbol, 0.5/len(longs)) for symbol in longs]
        targets += [PortfolioTarget(symbol, -0.5/len(shorts)) for symbol in shorts]
        self.set_holdings(targets, True)
        

class SymbolData:

    def __init__(self):
        self.roc = RateOfChange(1)
        self.delayed_roc = IndicatorExtensions.of(Delay(11), self.roc)