#region imports
from AlgorithmImports import *
from indicators import *
#endregion

class RankQuantilesAlphaModel(AlphaModel):

    def __init__(self, quantiles, lookback_months):
        self.quantiles = quantiles
        self.lookback_months = lookback_months
        self.securities_list = []
        self.day = -1

    def on_securities_changed(self, algorithm: QCAlgorithm, changes: SecurityChanges) -> None:
        # Create and register indicator for each security in the universe
        security_by_symbol = {}
        for security in changes.added_securities:

            # Create an indicator
            security_by_symbol[security.symbol] = security
            security.indicator = VwapReversion("indicator", security.symbol)
            self._register_indicator(algorithm, security)
            self.securities_list.append(security)
        
        # Warm up the indicators of newly-added stocks
        if security_by_symbol:
            history = algorithm.history[TradeBar](list(security_by_symbol.keys()), (self.lookback_months+1) * 30, Resolution.DAILY, data_normalization_mode=DataNormalizationMode.SCALED_RAW)
            for trade_bars in history:
                for bar in trade_bars.values():
                    if type(bar) == TradeBar:
                        security_by_symbol[bar.symbol].consolidator.update(bar)

        # Stop updating consolidator when the security is removed from the universe
        for security in changes.removed_securities:
            if security in self.securities_list:
                algorithm.subscription_manager.remove_consolidator(security.symbol, security.consolidator)
                self.securities_list.remove(security)


    def update(self, algorithm: QCAlgorithm, data: Slice) -> List[Insight]:
        # Reset indicators when corporate actions occur
        for symbol in set(data.splits.keys() + data.dividends.keys()):
            security = algorithm.securities[symbol]
            if security in self.securities_list:
                security.indicator.reset()
                algorithm.subscription_manager.remove_consolidator(security.symbol, security.consolidator)
                self._register_indicator(algorithm, security)

                history = algorithm.history[TradeBar](security.symbol, (security.indicator.warm_up_period+1) * 30, Resolution.DAILY, data_normalization_mode=DataNormalizationMode.SCALED_RAW)
                for bar in history:
                    security.consolidator.update(bar)
        
        # Only emit insights when there is quote data, not when a corporate action occurs (at midnight)
        if data.quote_bars.count == 0:
            return []
        
        # Only emit insights once per day
        if self.day == algorithm.time.day:
            return []
        self.day = algorithm.time.day

        # Get the indicator value of each asset in the universe
        indicator_by_symbol = {security.symbol : security.indicator.current.value 
            for security in self.securities_list if security.symbol in data.quote_bars and security.indicator.is_ready}
                
        # Determine how many assets to hold in the portfolio
        quantile_size = int(len(indicator_by_symbol)/self.quantiles)
        if quantile_size == 0:
            return []

        # Create insights to long the assets in the universe with the greatest indicator value
        weight = 1 / (quantile_size+1)
        insights = []
        for symbol, _ in sorted(indicator_by_symbol.items(), key=lambda x: x[1], reverse=True)[:quantile_size]:
            insights.append(Insight.price(symbol, Expiry.END_OF_DAY, InsightDirection.UP, weight=weight))

        return insights

    
    def _register_indicator(self, algorithm, security):
        # Update the indicator with monthly bars
        security.consolidator = TradeBarConsolidator(Calendar.MONTHLY)
        algorithm.subscription_manager.add_consolidator(security.symbol, security.consolidator)
        algorithm.register_indicator(security.symbol, security.indicator, security.consolidator)


class MomentumRank(AlphaModel):

    def __init__(self, quantiles, lookback_months):
        self.quantiles = quantiles
        self.lookback_months = lookback_months
        self.securities_list = []
        self.day = -1

    def on_securities_changed(self, algorithm: QCAlgorithm, changes: SecurityChanges) -> None:
        # Create and register indicator for each security in the universe
        security_by_symbol = {}
        for security in changes.added_securities:

            # Create an indicator
            security_by_symbol[security.symbol] = security
            security.indicator = CustomMomentumPercent("signal", self.lookback_months)     #CHANGE INDICATOR HERE
            self._register_indicator(algorithm, security)
            self.securities_list.append(security)
        
        # Warm up the indicators of newly-added stocks
        if security_by_symbol:
            history = algorithm.history[TradeBar](list(security_by_symbol.keys()), (self.lookback_months+1) * 30, Resolution.DAILY, data_normalization_mode=DataNormalizationMode.SCALED_RAW)
            for trade_bars in history:
                for bar in trade_bars.values():
                    if type(bar) == TradeBar:
                        security_by_symbol[bar.symbol].consolidator.update(bar)

        # Stop updating consolidator when the security is removed from the universe
        for security in changes.removed_securities:
            if security in self.securities_list:
                algorithm.subscription_manager.remove_consolidator(security.symbol, security.consolidator)
                self.securities_list.remove(security)


    def update(self, algorithm: QCAlgorithm, data: Slice) -> List[Insight]:
        # Reset indicators when corporate actions occur
        for symbol in set(data.splits.keys() + data.dividends.keys()):
            security = algorithm.securities[symbol]
            if security in self.securities_list:
                security.indicator.reset()
                algorithm.subscription_manager.remove_consolidator(security.symbol, security.consolidator)
                self._register_indicator(algorithm, security)

                history = algorithm.history[TradeBar](security.symbol, (security.indicator.warm_up_period+1) * 30, Resolution.DAILY, data_normalization_mode=DataNormalizationMode.SCALED_RAW)
                for bar in history:
                    security.consolidator.update(bar)
        
        # Only emit insights when there is quote data, not when a corporate action occurs (at midnight)
        if data.quote_bars.count == 0:
            return []
        
        # Only emit insights once per day
        if self.day == algorithm.time.day:
            return []
        self.day = algorithm.time.day

        # Get the indicator value of each asset in the universe
        indicator_by_symbol = {security.symbol : security.indicator.current.value 
            for security in self.securities_list if security.symbol in data.quote_bars and security.indicator.is_ready}

        # Create insights to long the assets in the universe with the greatest indicator value
        insights = []
        sorted_security_list = sorted(indicator_by_symbol.items(), key=lambda x: x[1])
        size = len(sorted_security_list)
        for security in sorted_security_list:
            weight = (sorted_security_list.index(security))/size-0.5
            if weight >= 0:
                insights.append(Insight.price(security[0], Expiry.END_OF_DAY, InsightDirection.UP, weight=weight*3))
            else:
                insights.append(Insight.price(security[0], Expiry.END_OF_DAY, InsightDirection.DOWN, weight=weight))

        return insights

    
    def _register_indicator(self, algorithm, security):
        # Update the indicator with monthly bars
        security.consolidator = TradeBarConsolidator(Calendar.MONTHLY)
        algorithm.subscription_manager.add_consolidator(security.symbol, security.consolidator)
        algorithm.register_indicator(security.symbol, security.indicator, security.consolidator)


class VWAPReversionTrendRank(AlphaModel):

    def __init__(self):
        self.securities_list = []
        self.day = -1

    def on_securities_changed(self, algorithm: QCAlgorithm, changes: SecurityChanges) -> None:
        # Create and register indicator for each security in the universe
        security_by_symbol = {}
        for security in changes.added_securities:

            # Create an indicator
            security_by_symbol[security.symbol] = security
            security.indicator = VwapReversion("custom", security.symbol, algorithm)
            algorithm.register_indicator(security.symbol, security.indicator)
            self.securities_list.append(security)


    def update(self, algorithm: QCAlgorithm, data: Slice) -> List[Insight]: 
        for symbol in set(data.splits.keys() + data.dividends.keys()):      # Reset indicators when corporate actions occur
            security = algorithm.securities[symbol]
            if security in self.securities_list:
                security.indicator.reset()
                algorithm.register_indicator(security.symbol, security.indicator)
        if data.quote_bars.count == 0:      # Only emit insights when there is quote data, not when a corporate action occurs (at midnight)
            return []
        if self.day == algorithm.time.day:      # Only emit insights once per day
            return []
        self.day = algorithm.time.day

        # Get the indicator value of each asset in the universe
        indicator_by_symbol = {security.symbol : security.indicator.current.value 
            for security in self.securities_list if security.symbol in data.quote_bars and security.indicator.is_ready}

        # Create insights to long / short the asset
        insights = []
        sorted_security_list = sorted(indicator_by_symbol.items(), key=lambda x: x[1])
        size = len(sorted_security_list)
        for security in sorted_security_list:
            weight = (sorted_security_list.index(security))/size - 0.5
            if weight >= 0:
                insights.append(Insight.price(security[0], Expiry.END_OF_DAY, InsightDirection.UP, weight=weight*3))
            else:
                insights.append(Insight.price(security[0], Expiry.END_OF_DAY, InsightDirection.DOWN, weight=weight))

        return insights

    
class TSZscore_DividendGrowth(AlphaModel):

    def __init__(self):
        self.period = 252
        self.securities_list = []
        self.day = -1
        self.historical_dividend_by_symbol = {}
        

    def on_securities_changed(self, algorithm: QCAlgorithm, changes: SecurityChanges) -> None:
        # Create and register indicator for each security in the universe
        for security in changes.added_securities:
            if security not in self.securities_list:
                self.historical_dividend_by_symbol[security.symbol] = deque(maxlen=self.period)
                self.securities_list.append(security)

        for security in changes.removed_securities:
            if security in self.securities_list:
                self.securities_list.remove(security)

    def update(self, algorithm: QCAlgorithm, data: Slice) -> List[Insight]:
        if data.quote_bars.count == 0:   # Only emit insights when there is quote data, not when a corporate action occurs (at midnight)
            return []
        if self.day == algorithm.time.day:  # Only emit insights once per day
            return []
        self.day = algorithm.time.day

        # Append dividend to the list, compute ts_zscore of current dividend
        zscore_by_symbol = {}
        for security in self.securities_list:
            self.historical_dividend_by_symbol[security.symbol].appendleft(security.fundamentals.earning_reports.dividend_per_share.Value)
            zscore_by_symbol[security.symbol] = sp.stats.zscore(self.historical_dividend_by_symbol[security.symbol])[0]

        # create insights to long / short the asset
        insights = []
        weights = {}
        size = len(zscore_by_symbol)
        for symbol, zscore in zscore_by_symbol.items():
            if not np.isnan(zscore):
                weight = zscore
            else:
                weight = 0
            weights[symbol] = weight

        weight_mean = sum(weights.values())/len(weights)

        for symbol, weight in weights.items():
            weight = weight
            if weight >= 0:
                insights.append(Insight.price(symbol, Expiry.END_OF_DAY, InsightDirection.UP, weight=weight))
            else:
                insights.append(Insight.price(symbol, Expiry.END_OF_DAY, InsightDirection.DOWN, weight=weight))

        return insights

#region imports
from AlgorithmImports import *
from collections import deque
import scipy as sp
import numpy as np
#endregion

def EWMA(value_history):
    output = value_history[0]
    for i in range(1, len(value_history)):
        output = 0.7 * value_history[i] + 0.3 * output
    return output
    
        
class CustomMomentumPercent(PythonIndicator):
    def __init__(self, name, period):
        self.name = name
        self.time = datetime.min
        self.value = 0
        self.momentum = MomentumPercent(period)

    def Update(self, input):
        self.momentum.update(IndicatorDataPoint(input.Symbol, input.EndTime, input.Close))
        self.time = input.EndTime
        self.value = self.momentum.Current.Value * input.Volume
        return self.momentum.IsReady


class Skewness(PythonIndicator):    # Doesn't work on 3th August 2020
    def __init__(self, name, period):
        self.name = name
        self.count = 0
        self.time = datetime.min
        self.value = 0
        self.queue = deque(maxlen=period)
        self.change_in_close = deque(maxlen=period)

    def Update(self, input):
        self.queue.appendleft(input.Close)
        if len(self.queue) > 1:
            self.change_in_close.appendleft(self.queue[0]/self.queue[1]-1)
        self.time = input.EndTime

        self.count = len(self.change_in_close)
        if self.count == self.queue.maxlen:
            self.value = sp.stats.skew(self.change_in_close, nan_policy="omit")
        
        return count == self.change_in_close.maxlen  


'''class VwapReversion(PythonIndicator):
    def __init__(self, name, symbol, algorithm):
        self.name = name
        self.time = datetime.min
        self.value = 0
        self.previous_value = 0
        self._vwap = algorithm.vwap(symbol)
        self.queue = deque(maxlen=30)
        

    def update(self, input):
        self._vwap.update(input)
        self.time = input.EndTime
        self.queue.appendleft(self._vwap.Current.Value / input.Close)

        count = len(self.queue)
        if count == self.queue.maxlen:
            z_array = sp.stats.zscore(self.queue)
            if np.isfinite(z_array[0]):
                self.previous_value = self.value
                self.value = 0.7 * z_array[0] + 0.3 * self.previous_value
                
        return count == self.queue.maxlen'''


    

# region imports
from AlgorithmImports import *
from alpha import *
# endregion


class LiquidEquityAlgorithm(QCAlgorithm):

    undesired_symbols_from_previous_deployment = []
    checked_symbols_from_previous_deployment = False

    def initialize(self):
        self.set_start_date(2013, 1, 1)
        self.set_end_date(2017, 1, 1) 
        self.set_cash(1_000_000)
        self.SetBenchmark(self.AddEquity("SPY").Symbol)
        self.set_brokerage_model(BrokerageName.INTERACTIVE_BROKERS_BROKERAGE, AccountType.MARGIN)
        self.settings.minimum_order_margin_portfolio_percentage = 0
        self.settings.rebalance_portfolio_on_security_changes = False
        self.settings.rebalance_portfolio_on_insight_changes = False
        self.day = -1
        self.set_warm_up(timedelta(150))

        self.universe_settings.asynchronous = True
        self.add_universe_selection(FundamentalUniverseSelectionModel(self.fundamental_filter_function))
        
        self.add_alpha(TSZscore_DividendGrowth())

        self.set_portfolio_construction(InsightWeightingPortfolioConstructionModel(rebalance=Expiry.EndOfMonth))

        self.add_risk_management(NullRiskManagementModel())

        self.set_execution(ImmediateExecutionModel())

    def on_data(self, data):
        # Exit positions that aren't backed by existing insights
        # If you don't want this behavior, delete this method definition.
        if not self.is_warming_up and not self.checked_symbols_from_previous_deployment:
            for security_holding in self.portfolio.values():
                if not security_holding.invested:
                    continue
                symbol = security_holding.symbol
                if not self.insights.has_active_insights(symbol, self.utc_time):
                    self.undesired_symbols_from_previous_deployment.append(symbol)
            self.checked_symbols_from_previous_deployment = True
        
        for symbol in self.undesired_symbols_from_previous_deployment:
            if self.is_market_open(symbol):
                self.liquidate(symbol, tag="Holding from previous deployment that's no longer desired")
                self.undesired_symbols_from_previous_deployment.remove(symbol)

        
    def fundamental_filter_function(self, fundamental: List[Fundamental]):
        filtered = [f for f in fundamental if f.symbol.value != "AMC" and f.has_fundamental_data and not np.isnan(f.dollar_volume)]
        sorted_by_dollar_volume = sorted(filtered, key=lambda f: f.dollar_volume, reverse=True)
        return [f.symbol for f in sorted_by_dollar_volume[:1000]]