from datetime import timedelta

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 QuantConnect.Data.UniverseSelection import *

from holdmonths import HoldMonthsTracker
from pricehist import PriceHistoryManager
from qcconsts import *
from qcselectors import *

"""
Algorithm selects for high ebit/ev and low variance in the coarse selection phase.
Then it selects for highest returns over the previous 11 months
"""


def has_all_fundamentals(fine_data):
    return (ev_ebit.has_ebit_ev_fundamentals(fine_data)
            and fine_data.MarketCap
            and fine_data.FinancialStatements.BalanceSheet.CurrentLiabilities
            and fine_data.OperationRatios.AssetsTurnover)


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


def log_symbols(algorithm, prefix, symbols):
    algorithm.Log(prefix + ' '.join(map(lambda s: s.Value, symbols)))


class ValueWithDownFlatMomentum(QCAlgorithm):

    def __init__(self):
        # Number of stocks to select from the alpha model
        super().__init__()
        self._alpha_target_count = 20
        self._hold_months = HoldMonthsTracker(3)
        self._p_coarse_vol = 0.5
        self._p_coarse_price = 0.7
        self._price_history_manager = PriceHistoryManager(int(BARS_PER_YEAR / 4))
        self._market_cap_percentile = 50

    def Initialize(self):
        self.SetCash(100000)
        self.SetStartDate(2020, 1, 1)
        # If not specified, the backtesting EndDate would be today
        self.UniverseSettings.Resolution = Resolution.Daily
        # self.SetEndDate(2015, 7, 1)
        self.AddUniverseSelection(FineFundamentalUniverseSelectionModel(self.select_coarse, self.select_fine))
        self.AddAlpha(LongShortValueDownFlatAlphaModel(self._alpha_target_count, self._hold_months.clone()))
        self.SetPortfolioConstruction(EqualWeightingPortfolioConstructionModel(rebalance=None))
        self.SetExecution(VolumeWeightedAveragePriceExecutionModel())
        # self.SetExecution(StandardDeviationExecutionModel(deviations=1, period=5, resolution=Resolution.Daily))

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

    def select_by_variance(self, algorithm):
        close_variances = {}
        for sym, closes in self.close_histories.items():
            variance = closes.diff().var()
            if m.isfinite(variance):
                close_variances[sym] = variance
            else:
                algorithm.Log(f'Dropping {sym} which has variance {variance}')
        sorted_by_var = sorted(close_variances.items(), key=lambda kv: kv[1])
        head_tail_size = int(len(sorted_by_var) / 4)
        medium_var = sorted_by_var[head_tail_size:-head_tail_size]
        algorithm.Log('Dropping highest and lowest {} by variance. Min/max: {}/{}'
                      .format(head_tail_size, medium_var[0][1], medium_var[-1][1]))
        return [s[0] for s in medium_var]

    def select_coarse(self, coarse):
        if not self._hold_months.should_trade(self.Time.month):
            return Universe.Unchanged
        self.Log('Running coarse selection')
        with_fundamental = [x for x in coarse if x.HasFundamentalData]
        by_vol = sort_and_select_subset(self, with_fundamental, lambda s: s.DollarVolume,
                                        reverse=True, proportion=self._p_coarse_vol, value_name='Volume')
        by_price = sort_and_select_subset(self, with_fundamental, lambda s: s.Price,
                                          reverse=True, proportion=self._p_coarse_price, value_name='Price')
        symbols_by_pv = [s.Symbol for s in (set(by_vol).intersection(by_price))]
        self._price_history_manager.update_close_histories(self, symbols_by_pv)
        return self.select_by_variance(self)

    def select_fine(self, fine):
        if not self._hold_months.should_trade_and_set(self.Time.month):
            return Universe.Unchanged
        self.Log('Running fine selection')
        with_fundamentals = [f for f in fine if has_all_fundamentals(f) and is_tradable(f)]
        with_cap = select_by_market_cap(self, with_fundamentals, self._market_cap_percentile)
        return [f.Symbol for f in with_cap]


class LongShortValueDownFlatAlphaModel(AlphaModel):

    def __init__(self, _alpha_max_count, _hold_months):
        super().__init__()
        self._hold_months = _hold_months
        self._max_long_count = _alpha_max_count
        self._max_short_count = _alpha_max_count
        self._candidate_factor = 4
        # 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._shorts = []
        self._price_history_manager = PriceHistoryManager(BARS_PER_YEAR)
        self._universe = set([])

    def split_long_short_candidates(self):
        ordered_candidates = candidates_ordered_by_combined_rank(
            rank_by_long_term_debt_to_equity(self._universe),
            rank_by_ebit_ev(self._universe))
        long_candidates_count = self._max_long_count * self._candidate_factor
        short_candidates_count = self._max_short_count * self._candidate_factor
        if long_candidates_count + short_candidates_count > len(ordered_candidates):
            long_candidates_count = int(len(ordered_candidates) / 2)
            short_candidates_count = len(ordered_candidates) - long_candidates_count
        return ordered_candidates[:long_candidates_count], list(reversed(ordered_candidates[-short_candidates_count:]))

    def order_candidates(self, candidates):
        down_phase, flat_phase = self._price_history_manager.split_at_ratio(0.75, candidates)
        down_ranks = (pd.Series({s: compute_returns(h) for s, h in down_phase.items()})
                      .rank(ascending=True)
                      .sort_index())
        flat_ranks = (pd.Series({s1: compute_returns(h1) for s1, h1 in flat_phase.items()})
                      .abs()
                      .rank(ascending=True)
                      .sort_index())
        aggregate_rank = pd.DataFrame({'down': down_ranks, 'flat': flat_ranks}).max(1)
        return aggregate_rank.sort_values().index

    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}')
        long_candidates, short_candidates = self.split_long_short_candidates()
        self._price_history_manager.update_close_histories(algorithm, long_candidates + short_candidates)
        self._longs = self.order_candidates(long_candidates)[:self._max_long_count]
        # Reversing the list that's sorted by max of down-rank and flat-rank is
        # not the same as taking the highest min values of the down-rank and the
        # flat rank. We will need to implement long and short computation better.
        self._shorts = list(reversed(self.order_candidates(short_candidates)))[:self._max_short_count]

    def Update(self, algorithm, data):
        if not self._hold_months.should_trade_and_set(algorithm.Time.month):
            return []
        insights = []
        # Months have different numbers of days, so this will be approximate
        hold_duration = timedelta(BARS_PER_MONTH * self._hold_months.get - 1)
        for kvp in algorithm.Portfolio:
            holding = kvp.Value
            symbol = holding.Symbol
            if holding.Invested and symbol not in self._longs and symbol not in self._shorts:
                insights.append(Insight.Price(symbol, hold_duration, InsightDirection.Flat))
        log_symbols(algorithm, "Long positions selected: ", self._longs)
        for symbol in self._longs:
            insights.append(Insight.Price(symbol, hold_duration, InsightDirection.Up))
        # For now, just focus on long positions
        # log_symbols(algorithm, "Short positions selected: ", self._longs)
        # for symbol in self._shorts:
        #     insights.append(Insight.Price(symbol, hold_duration, InsightDirection.Down))
        return insights