#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 TSZscore_VWAPReversion(AlphaModel):

    def __init__(self):
        self.period = 20
        self.securities_list = []
        self.day = -1
        self.count = 0
        self.historical_VwapReversion_by_symbol = {}


    def on_securities_changed(self, algorithm: QCAlgorithm, changes: SecurityChanges) -> None:
        # Register each security in the universe
        for security in changes.added_securities:
            if security not in self.securities_list:
                self.historical_VwapReversion_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

        # Neutralize Vwap/Close of securities so it's mean 0, then append them to the list 
        temp_list = {}
        for security in self.securities_list:
            if security.Close != 0:
                temp_list[security.symbol] = algorithm.vwap(security.symbol).Current.Value/security.Close
            else: 
                temp_list[security.symbol] = 1
        temp_mean = sum(temp_list.values())/len(temp_list.values())

        for security in self.securities_list:
            self.historical_VwapReversion_by_symbol[security.symbol].appendleft(temp_list[security.symbol]-temp_mean)
        
        # Emit insight once per 20 days
        if self.count != 20:
            self.count += 1
            return []
        else: 
            self.count = 0

        # Compute ts_zscore of current Vwap/Close
        zscore_by_symbol = {}
        for security in self.securities_list:
            zscore_by_symbol[security.symbol] = sp.stats.zscore(self.historical_VwapReversion_by_symbol[security.symbol])[0]
        
        # create insights to long / short the asset
        insights = []
        weights = {}
        for symbol, zscore in zscore_by_symbol.items():
            if not np.isnan(zscore):
                weight = zscore
            else:
                weight = 0
            weights[symbol] = weight

        for symbol, weight in weights.items():
            if weight > 0:
                insights.append(Insight.price(symbol, timedelta(20), InsightDirection.UP, weight=weight))
            elif weight < 0:
                insights.append(Insight.price(symbol, timedelta(20), 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 = {}
        self.previous_zscore_by_symbol = {}
        

    def on_securities_changed(self, algorithm: QCAlgorithm, changes: SecurityChanges) -> None:
        # Register 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 = {}
        for symbol, zscore in zscore_by_symbol.items():
            if not np.isnan(zscore):
                weight = zscore
            else:
                weight = 0
            weights[symbol] = weight

        for symbol, weight in weights.items():
            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))

        self.previous_zscore_by_symbol = zscore_by_symbol

        return insights


class Conditional_Reversion(AlphaModel):

    def __init__(self):
        self.condition_period = 5
        self.period = 3
        self.securities_list = []
        self.day = -1
        self.historical_volume_by_symbol = {}
        self.historical_close_by_symbol = {}
        

    def on_securities_changed(self, algorithm: QCAlgorithm, changes: SecurityChanges) -> None:
        # Register each security in the universe
        for security in changes.added_securities:
            if security not in self.securities_list:
                self.historical_volume_by_symbol[security.symbol] = deque(maxlen=self.condition_period)
                self.historical_close_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 month
            return []
        self.day = algorithm.time.day
        

        # Append volume and close to the list
        zscore_by_symbol = {}
        return_by_symbol = {}
        for security in self.securities_list:
            if (security.Close != 0 and security.Volume != 0):
                self.historical_close_by_symbol[security.symbol].appendleft(security.Close)
                self.historical_volume_by_symbol[security.symbol].appendleft(security.Volume)
                return_by_symbol[security.symbol] = (self.historical_close_by_symbol[security.symbol][0] - self.historical_close_by_symbol[security.symbol][-1])

        if return_by_symbol == {}:  # Don't emit insight if there's no valid data
            return []

        # Rank the 3 days return among securities to return value from 0 to 1
        sorted_return_by_symbol = sorted(return_by_symbol.items(), key=lambda x: x[1])
        return_rank_by_symbol = {}
        for item in sorted_return_by_symbol:   # item is a key-value pair. [0] is the security symbol and [1] is the return
            return_rank_by_symbol[item[0]] = (sorted_return_by_symbol.index(item))/ len(sorted_return_by_symbol)
        
        # Calculating the final weight
        weights = {}
        
        for security in self.securities_list:
            # If condition is met, assign weight
            if len(self.historical_volume_by_symbol[security.symbol]) != 0 and max(self.historical_volume_by_symbol[security.symbol]) == security.Volume:
                weight = -return_rank_by_symbol[security.symbol] # Change this sign and complete different behaviour if purely long. Investigate
            else:
                weight = 0 
            weights[security.symbol] = weight

        p_count = 0
        n_count = 0
        # Make the weights mean 0
        weights_mean = sum(weights.values())/len(weights.values())
        for symbol, weight in weights.items():
            weights[symbol] = weight - weights_mean
            if weights[symbol] > 0:
                p_count += 1
            else:
                n_count += 1


        # Create insights to long / short the asset
        insights = []
        for symbol, weight in weights.items():
            if weight > 0:
                insights.append(Insight.price(symbol, Expiry.END_OF_DAY, InsightDirection.UP, weight=weight))
            elif weight < 0:
                insights.append(Insight.price(symbol, Expiry.END_OF_DAY, InsightDirection.DOWN, weight=weight))
        #Expiry.END_OF_DAY
        return insights


class Momomentum(AlphaModel):

    def __init__(self):
        self.condition_period = 5
        self.period = 3
        self.securities_list = []
        self.day = -1
        self.historical_volume_by_symbol = {}
        self.historical_close_by_symbol = {}
        

    def on_securities_changed(self, algorithm: QCAlgorithm, changes: SecurityChanges) -> None:
        # Register each security in the universe
        for security in changes.added_securities:
            if security not in self.securities_list:
                self.historical_volume_by_symbol[security.symbol] = deque(maxlen=self.condition_period)
                self.historical_close_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 volume and close to the list
        zscore_by_symbol = {}
        return_by_symbol = {}
        for security in self.securities_list:
            self.historical_close_by_symbol[security.symbol].appendleft(security.Close)
            self.historical_volume_by_symbol[security.symbol].appendleft(security.Volume)
            return_by_symbol[security.symbol] = (self.historical_close_by_symbol[security.symbol][0] - self.historical_close_by_symbol[security.symbol][-1])
            # abs return? 

        # Rank the 3 days return among securities to return value from 0 to 1
        sorted_return_by_symbol = sorted(return_by_symbol.items(), key=lambda x: x[1])
        ranked_return_by_symbol = {}
        for item in sorted_return_by_symbol:   # item is a key-value pair. [0] is the security symbol and [1] is the return
            ranked_return_by_symbol[item[0]] = (sorted_return_by_symbol.index(item))/ len(sorted_return_by_symbol)
        

        # Calculating the final weight
        weights = {}
        for symbol, rank in ranked_return_by_symbol.items():
            # If condition is met, assign weight
            condition = (max(self.historical_volume_by_symbol[security.symbol]) == security.Volume)
            if condition:
                weight = rank
            else:
                weight = 0 
            weights[symbol] = weight

        # Make the weights mean 0
        weights_mean = sum(weights.values())/len(weights.values())
        for symbol_weight in weights.items():
            weights[symbol] = weight - weights_mean
           
        # Make the weights absolute value sum to 1
        weights_sum = sum(np.abs(list(weights.values())))
        for symbol, weight in weights.items():
            weights[symbol] = weight/weights_sum

        # # Create insights to long / short the asset
        insights = []
        for symbol, weight in weights.items():
            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


class Condition_Reversion(AlphaModel):

    def __init__(self):
        self.condition_period = 5
        self.securities_list = []
        self.day = -1
        self.historical_volume_by_symbol = {}
        

    def on_securities_changed(self, algorithm: QCAlgorithm, changes: SecurityChanges) -> None:
        # Register each security in the universe
        for security in changes.added_securities:
            if security not in self.securities_list:
                self.historical_volume_by_symbol[security.symbol] = deque(maxlen=self.condition_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 volume to the list
        for security in self.securities_list:
            if security.Volume != 0:
                self.historical_volume_by_symbol[security.symbol].appendleft(security.Volume)
        
        weights = {}
        n_count = 0
        p_count = 0
        for security in self.securities_list:
            # If condition is met, assign weight
            if len(self.historical_volume_by_symbol[security.symbol]) != 0 and max(self.historical_volume_by_symbol[security.symbol]) == security.Volume:
                weights[security.symbol] = -1
                n_count += 1
            else:
                weights[security.symbol] = 1
                p_count += 1

        for symbol, weight in weights.items():
            if weight == -1:
                weights[symbol] = weight/n_count
            elif weight == -1:
                weights[symbol] = weight/p_count
            

        # Create insights to long / short the asset
        insights = []
        for symbol, weight in weights.items():
            if weight < 0:
                insights.append(Insight.price(symbol, Expiry.END_OF_DAY, InsightDirection.UP, weight=weight))
            '''elif weight < 0:
                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(2020, 1, 1)
        self.set_end_date(2024, 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(20))

        self.universe_settings.asynchronous = True
        self.add_universe_selection(FundamentalUniverseSelectionModel(self.fundamental_filter_function))
        #self.add_universe_selection(ETFConstituentsUniverseSelectionModel("SPY"))
        
        self.add_alpha(TSZscore_VWAPReversion())
        #self.add_alpha(ConstantAlphaModel(InsightType.PRICE, InsightDirection.UP, timedelta(1)))

        self.set_portfolio_construction(InsightWeightingPortfolioConstructionModel(Expiry.EndOfMonth))
        #self.set_portfolio_construction(EqualWeightingPortfolioConstructionModel())
        #rebalance=timedelta(5)
        #rebalance=Expiry.EndOfMonth
        
        self.add_risk_management(NullRiskManagementModel())

        self.set_execution(ImmediateExecutionModel())

        self.previousHoldingsValue = {}
        self.function_time = -1

    
    '''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[0:1000]]