from QuantConnect import *
from QuantConnect.Algorithm import *
from QuantConnect.Algorithm.Framework.Alphas import *
from QuantConnect.Algorithm.Framework.Execution import *
from QuantConnect.Algorithm.Framework.Portfolio import *
from QuantConnect.Algorithm.Framework.Selection import *

from datetime import timedelta
import pandas as pd

BARS_PER_YEAR = 252


class PriceHistoryManager:
    def __init__(self, num_bars):
        self._num_bars = num_bars
        self._close_histories = {}
        self._last_time = None

    @property
    def close_histories(self):
        return self._close_histories

    def apply_fetched_close_histories(self, algorithm, history_df, updated_closes):
        history_by_syms = history_df["close"].unstack(level=0)
        for s in history_by_syms.columns:
            s_new_hist = history_by_syms[s]
            s_sym = algorithm.Symbol(s)
            if s_sym in updated_closes:
                s_old_hist = updated_closes[s_sym]
                updated_closes[s_sym] = pd.concat([s_old_hist, s_new_hist]).sort_index()[-self._num_bars:]
            else:
                updated_closes[s_sym] = s_new_hist.sort_index()[-self._num_bars:]

    def update_close_histories(self, algorithm, symbols):
        updated_closes = {}
        symbols_to_update = []
        new_symbols = []
        # For each symbol that already has history, copy that history to the next_indicators.
        for s in symbols:
            if s in self._close_histories:
                symbols_to_update.append(s)
                updated_closes[s] = self._close_histories[s]
            else:
                new_symbols.append(s)
        # If there are symbols to update, get the incremental history
        if symbols_to_update:
            # Because we drop all history that isn't right now relevant, we are assured that all the
            # history we require will be from the last run to now.
            inc_history = algorithm.History(symbols_to_update, self._last_time, algorithm.Time, Resolution.Daily)
            if not inc_history.empty:
                self.apply_fetched_close_histories(algorithm, inc_history, updated_closes)
        # If there are new symbols, get the full history
        if new_symbols:
            full_history = algorithm.History(new_symbols, self._num_bars, Resolution.Daily)
            if not full_history.empty:
                self.apply_fetched_close_histories(algorithm, full_history, updated_closes)
        self._close_histories = updated_closes
        self._last_time = algorithm.Time
        assert len(self._close_histories) <= len(symbols), f'close histories does not match symbols count: {len(self._close_histories)} > {len(symbols)}'
        for s in symbols:
            assert len(self._close_histories.get(s, [])) <= self._num_bars, f'close histories for {s} is too large: {len(self._close_histories[s])}'

    def split_at_ratio(self, ratio, symbols=None):
        assert 1.0 >= ratio >= 0.0
        return self.split_at(int(self._num_bars * ratio), symbols)

    def split_at(self, bars, symbols=None):
        return self.tail(bars, symbols), self.head(self._num_bars - bars, symbols)

    def tail(self, bars, symbols=None):
        """Grabs the oldest bars."""
        assert bars >= 0
        symbols = symbols and set(symbols)
        return {sym: history[:bars] for sym, history in self._close_histories.items()
                if symbols is None or sym in symbols}

    def head(self, bars, symbols=None):
        """Grabs the newest bars"""
        assert bars >= 0
        return {sym: history[-bars:] for sym, history in self._close_histories.items()
                if symbols is None or sym in symbols}


class EtfMomentum(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2020, 1, 1)  # Set Start Date
        self.SetCash(100000)  # Set Strategy Cash
        # self.AddEquity("SPY", Resolution.Minute)
        symbols = [Symbol.Create(ticker, SecurityType.Equity, Market.USA) for ticker in
            ["SPY", "QQQ", "VBR", "IPAC", "IEUR", "ILTB", "IUSG", "IUSB", "VEA", "VWO", "XCEM"]]
        self.SetUniverseSelection(ManualUniverseSelectionModel(symbols))
        self.AddAlpha(SimpleMomentumAlphaModel(4))
        self.SetPortfolioConstruction(EqualWeightingPortfolioConstructionModel(rebalance=None))
        self.SetExecution(ImmediateExecutionModel())


def compute_returns(history):
    return history[-1]/history[0] - 1


class SimpleMomentumAlphaModel(AlphaModel):
    def __init__(self, _alpha_max_count):
        self._last_month = None
        self._max_long_count = _alpha_max_count
        # For the purposes of the algorithm, the number of days we want to assume
        # for each month. This is not exact, but it should be a good enough measure
        self._month_days = timedelta(int(BARS_PER_YEAR/12))
        # Longs and shorts will be set by OnSecuritiesChanged, and will be used in
        # Update. The latter will not do any data retrieval.
        self._longs = []
        self._price_history_manager = PriceHistoryManager(int(BARS_PER_YEAR/2))
        self._universe = set([])

    def order_candidates(self, candidates):
        up_phase, ignored = self._price_history_manager.split_at_ratio(11.0/12.0, candidates)
        up_returns = pd.Series({s: compute_returns(h) for s, h in up_phase.items()})
        up_ranks = up_returns.rank(ascending=False)
        sorted_up_ranks = up_ranks.sort_index()
        sorted_candidates = sorted_up_ranks.sort_values().index
        return sorted_candidates

    def OnSecuritiesChanged(self, algorithm, changes):
        added = set(changes.AddedSecurities)
        removed = set(changes.RemovedSecurities)
        self._universe = self._universe.union(added).difference(removed)
        algorithm.Log(f'Updating securities in alpha model.'
                      f' Added: {len(added)}. Removed: {len(removed)}. Universe: {len(self._universe)}')
        algorithm.Log(f'Active securities: {algorithm.ActiveSecurities.Count}')

    def Update(self, algorithm, data):
        # Emit signals once a month
        if self._last_month == algorithm.Time.month:
            return []
        self._last_month = algorithm.Time.month
        long_candidates = [s.Symbol for s in self._universe]
        self._price_history_manager.update_close_histories(algorithm, long_candidates)
        self._longs = self.order_candidates(long_candidates)[:self._max_long_count]
        insights = []
        for kvp in algorithm.Portfolio:
            holding = kvp.Value
            symbol = holding.Symbol
            if holding.Invested and symbol not in self._longs:
                insights.append(Insight.Price(symbol, self._month_days, InsightDirection.Flat))
        for symbol in self._longs:
            insights.append(Insight.Price(symbol, self._month_days, InsightDirection.Up))
        return insights