# region imports
from AlgorithmImports import *
# endregion

class MaintainHistoricalDailyUniversePriceDataAlgorithm(QCAlgorithm):

    def initialize(self):
        self.set_start_date(2010, 1, 1)
        #self.set_end_date(2020, 3, 1)
        self.set_cash(1_000_000)

        # Add a universe of daily data.
        self.universe_settings.resolution = Resolution.DAILY
        self._universe = self.add_universe(
            lambda fundamentals: [f.symbol for f in sorted(fundamentals, key=lambda f: f.market_cap)[-10:]]
        )

        # Create a DataFrame to store the historical data.
        self._all_history = pd.DataFrame()
        # Define the lookback period.
        self._lookback = 252  # Trading days.

        # Create a Scheduled Event to record new daily prices and
        # rebelance the portfolio.
        spy = Symbol.create('SPY', SecurityType.EQUITY, Market.USA)
        self.schedule.on(
            self.date_rules.every_day(spy),
            self.time_rules.at(0, 1), # One minute after `on_securities_changed` runs (in backtests)
            self._rebalance
        )

    def on_securities_changed(self, changes):
        # Remove the historical prices of assets that leave the universe.
        for security in changes.removed_securities:
            if security.symbol in self._all_history.columns:
                self._all_history.drop(security.symbol, axis=1, inplace=True)

        # Warm-up the historical data of assets that enter the universe.
        history = self.history(
            [security.symbol for security in changes.added_securities], self._lookback, Resolution.DAILY
        )
        if not history.empty:
            # If you trade at market open, it might make more sense to generate signals
            # based on daily opening prices.
            self._all_history = self._all_history.join(history.open.unstack(0), how='outer')

    def _rebalance(self):
        # Add yesterday's open price to the DataFrame of historical prices.
        self._all_history = pd.concat([
            self._all_history, 
            self.history(list(self._universe.selected), 1, Resolution.DAILY).open.unstack(0)
        ])
        # Yesterday's open price has been added twice for assets that just entered the universe.
        # The first time was in `on_securities_changed` and the second time was in the line above.
        # Let's drop rows with duplicate indices, then trim to lookback window size.
        self._all_history = self._all_history.loc[
            ~self._all_history.index.duplicated(keep='last')
        ].iloc[-self._lookback:]

        history = self._all_history[[symbol for symbol in self._all_history.columns if self.securities[symbol].price]]

        # Calculate asset signals for this rebalance.
        # For example, set the signal to give greater weight to uncorrelated assets.
        signal_by_symbol = 1/history.dropna(axis=1).corr().abs().sum()
        signal_by_symbol /= signal_by_symbol.sum()

        # Rebalance the portfolio based on the signals.
        self.set_holdings([PortfolioTarget(symbol, signal) for symbol, signal in signal_by_symbol.items()], True)