# region imports
from AlgorithmImports import *
# endregion


class SwimmingLightBrownGalago(QCAlgorithm):

    def initialize(self):

        self.set_start_date(2010, 1, 1)
        self.set_cash(100_000_000)

        self.settings.daily_precise_end_time = False
        self.settings.minimum_order_margin_portfolio_percentage = 0.01

        lookback_years = 12 # 12 years includes 3 non-leap years & 3 election cycles
        self._lookback = lookback_years * 252 
        self._lookback_calendar_time = timedelta(lookback_years*365)
        self._period = 21 # Trading days

        self.universe_settings.resolution = Resolution.DAILY
        self._universe = self.add_universe(self._select_assets)

        spy = Symbol.create('SPY', SecurityType.EQUITY, Market.USA)
        self.schedule.on(
            self.date_rules.every_day(spy),
            self.time_rules.midnight,
            self._rebalance
        )

        self.set_portfolio_construction(InsightWeightingPortfolioConstructionModel(portfolio_bias=PortfolioBias.LONG))

        self._all_history = pd.DataFrame()

    def _select_assets(self, fundamentals):
        return [
            Symbol.create(ticker, SecurityType.EQUITY, Market.USA)
            for ticker in ['SPY', 'QQQ']
        ]

    def _select_uncorrelated_assets(self, fundamentals):
        # Select large caps.
        large_caps = sorted(
            [f for f in fundamentals if f.symbol.id.date <= self.time - self._lookback_calendar_time],
            key=lambda f: f.market_cap
        )[:100]
        # Calculate the correlation matrix.
        corr = self.history(
            [f.symbol for f in large_caps], self._lookback, Resolution.DAILY
        ).open.unstack(0).pct_change(self._period).shift(-self._period).dropna().corr()
        # Select the least correlated assets.
        selected = []
        for index, row in corr.iterrows():
            corr_rank = row.abs().sum()
            selected.append((index, corr_rank))
        return sorted([x[0] for x in sort_[:5]], key=lambda x: x[1])

    def on_securities_changed(self, changes):
        for security in changes.removed_securities:
            self._all_history.drop(security.symbol, inplace=True)

        history = self.history([security.symbol for security in changes.added_securities], self._lookback, Resolution.DAILY)
        if not history.empty:
            self._all_history = self._all_history.join(history.open.unstack(0), how='outer')

    def _rebalance(self):
        # Append new rows.
        self._all_history = pd.concat([
            self._all_history, 
            self.history(list(self._universe.selected), 1, Resolution.DAILY).open.unstack(0)
        ])
        # 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:]

        # Calculate signal for this trade.
        period_returns = self._all_history.dropna(axis=1).pct_change(self._period).shift(-self._period).dropna()
        expected_return_by_day = period_returns.groupby(period_returns.index.strftime('%m-%d')).mean()
        signals = expected_return_by_day.iloc[
            self._wrapped_indices(
                len(expected_return_by_day.index), 
                expected_return_by_day.index.get_loc((self.time + timedelta(1)).strftime('%m-%d')), 
                self._period
            )
        ].sum()

        # Rebalance.
        self.emit_insights([Insight.price(symbol, timedelta(30), InsightDirection.UP, weight=1000*signal) for symbol, signal in signals.items()])

    def _wrapped_indices(self, index_length, start_index, count):
        start_index = start_index % index_length
        indices = [(start_index + i) % index_length for i in range(count)]
        return indices