import pandas as pd
import numpy as np
from collections import deque

class EmaCrossUniverseSelectionAlgorithm(QCAlgorithm):

    def Initialize(self): 
        '''
        Initialise the data and resolution required, as well as the cash and start-end dates for your algorithm.
        All algorithms must initialized.
        '''

        self.SetStartDate(2017,6,1)  #Set Start Date
        self.SetEndDate(2017,6,30)    #Set End Date
        self.SetCash(1000000)           #Set Strategy Cash

        # ADJUST RESOLUTION TO CHANGE ORDER TYPES
        self.UniverseSettings.Resolution = Resolution.Daily
        self.UniverseSettings.Leverage = 2

        self.coarse_count = 100
        self.averages = {}
        self.TOLERANCE = 0.025  # +/- target weight
        
        # this add universe method accepts two parameters:
        # - coarse selection function: accepts an IEnumerable<CoarseFundamental> and returns an IEnumerable<Symbol>
        self.AddUniverse(self.CoarseSelectionFunction)
        
        self.symbols = None
        self.spy = self.AddEquity("SPY", Resolution.Minute).Symbol 
        
        self.ema_fast_window = 2
        self.ema_slow_window = 10
        
        self.LOOKBACK = 11 # 60
        self.equity_arr = deque([0], maxlen=252)
     
    #======================================================================    
        # schedule rebalance
        
        # make buy list
        self.Schedule.On(
            #self.DateRules.EveryDay(self.spy),
            #self.DateRules.MonthStart(self.spy),
            self.DateRules.WeekStart(self.spy),
            self.TimeRules.AfterMarketOpen(self.spy, 30),
            Action(self.rebalance),
        )    
    #======================================================================
    # REF: https://www.quantconnect.com/docs/algorithm-reference/universes#Universes-Coarse-Universe-Selection
    
    def CoarseSelectionFunction(self, coarse):
        """
        This is the coarse selection function which filters the stocks. 
        
        This may be split between coarse and fine selection but does 
        not seem necessary at this point.
        See REF: https://www.quantconnect.com/docs/algorithm-reference/universes#Universes-Fundamentals-Selection
        """
        
        # -------------------------------------------------
        # only trade stocks that have fundamental data
        # price greater than $5 USD and
        # Dollar Volume greater than $10MM USD
        # Can be adjusted however you want, see the above reference.
        
        filterCoarse = [x for x in coarse if x.HasFundamentalData and x.Price > 5 and x.DollarVolume > 10000000]        
        
        # -------------------------------------------------
        # Add SymbolData to self.averages dict keyed by Symbol
        # Call history for new symbols to warm up the SymbolData internals 
        
        tmp_symbols = [cf.Symbol for cf in filterCoarse if cf.Symbol not in self.averages]
        # need more than longest indicator data to compute rolling 
        history = self.History(tmp_symbols, self.LOOKBACK, Resolution.Daily) 
        if not history.empty:
            history = history.close.unstack(level=0)

            for symbol in tmp_symbols:
                if str(symbol) in history.columns:
                    self.averages[symbol] = SymbolData(self, history[symbol])

        # Updates the SymbolData object with current EOD price
        for cf in filterCoarse:
            avg = self.averages.get(cf.Symbol, None)
            if avg is not None:
                avg.Update(cf) # add prices by passing the cf object 
        
        # --------------------------------------------------
        # sort dict by mmi metric and convert to dataframe
        
        mmiDict = {symbol:value.Metric for symbol, value in self.averages.items()}
        # want to get the lowest values so remove the reveres=True flag
        # sort and convert to df
        sorted_by_metric = pd.DataFrame(
            sorted(mmiDict.items(), key=lambda x: x[1]), columns=['symbol','metric']
        ).iloc[:self.coarse_count] 
        # we need the Symbol.Value to compare symbols
        sorted_by_metric_symbols = [x.Value for x in sorted_by_metric["symbol"]]
        #self.Log(f"[{self.Time}] sorted by metric | shape {sorted_by_metric.shape}:\n{sorted_by_metric}")
        
        # --------------------------------------------------
        # extract symbols that are in uptrend
        
        #ema_uptrend = list(filter(lambda x: x.is_uptrend, self.averages.values()))
        # self.Debug(f"[{self.Time}] uptrend syms: {len(ema_uptrend_sym)} |\n {ema_uptrend_sym}")
                
        # --------------------------------------------------
        # extract symbols from uptrending list
        
        # This is the official symbol list. It contains the QC
        # symbol object 
        #symbols = [x.symbol for x in ema_uptrend if x.symbol.Value in sorted_by_metric_symbols]
        symbols = [x for x in sorted_by_metric["symbol"]]
        self.Debug(f"[{self.Time}] number of selected symbols: {len(symbols)}")
        #if len(symbols) > 0:
        #    self.Debug(f"{self.Time} type: {type(symbols)} |\n{symbols}")

        # --------------------------------------------------
        # assign symbols to global parameter so rebalancing
        #   can be handled in OnSecuritiesChanged() function
        self.symbols = symbols 
        return symbols # we need to return only the symbol objects

    #======================================================================
    # scheduled rebalance
    
    def update_portfolio_value(self):
        """
        Update portfolio equity value tracked by self.equity_arr by appending portfolio value to list
        """
        self.equity_arr.append(self.Portfolio.TotalPortfolioValue)
        return
    
    def check_current_weight(self, symbol):
        """
        Check current symbol portfolio weight

        :param symbol: str
        :return current_weight: float
        """
        P = self.Portfolio
        current_weight = float(P[symbol].HoldingsValue) / float(P.TotalPortfolioValue)
        return current_weight    
    
    def rebalance(self):
        
        if self.symbols is None:
            return
        
        self.Debug(f"[{self.Time}] init rebalance")
        
        # columns = []
        prices_list = []
        for k,v in self.averages.items():
            tmp_hist = self.averages[k].Window
            prices_list.append(tmp_hist)
            
        self.Log(f"{self.Time} series list:\n{str(prices_list[0])}")    
        self.Log(f"{self.Time} series list:\n{str(prices_list)}")            
        prices = pd.concat(prices_list, axis=1).dropna()
        
        returns = np.log(prices).diff().dropna()
        self.Log(f"{self.Time} returns:\n{returns.head()}")
        
        #self.update_prices()  # update prices
        self.update_portfolio_value()  # update portfolio value
        
        for sec in self.symbols:
            # get current weights
            current_weight = self.check_current_weight(sec)
            target_weight = 1/len(self.symbols) # * returns[sec].iloc[-1]

            # if current weights outside of tolerance send new orders
            tol = self.TOLERANCE * target_weight
            lower_bound = target_weight - tol
            upper_bound = target_weight + tol

            if (current_weight < lower_bound) or (current_weight > upper_bound):
                self.SetHoldings(sec, target_weight)
                
        return

    #======================================================================
    # this event fires whenever we have changes to our universe
    
    def OnSecuritiesChanged(self, changes):
        
        # liquidate removed securities
        for security in changes.RemovedSecurities:
            if security.Invested:
                self.Liquidate(security.Symbol)
            
    def OnData(self, data):
        pass

#======================================================================

class SymbolData(object):
    """
    This is the symbol data class. I have updated it to be able to 
    take advantage of what the MMI indicator is used for and that is
    trending securities. In this example I implemented a simple
    short term EMA crossover strategy, but it is easy enough to
    change.
    """
    def __init__(self, qccontext, history):

        self.FastPeriod= qccontext.ema_fast_window #30
        self.SlowPeriod= qccontext.ema_slow_window #90
        #self.MainPeriod=90
        self.Window = history
        self.EndTime = history.index[-1]
        
        # ema crossover attributes
        self.is_uptrend = False
        #self.ema_fast = ExponentialMovingAverage(8)
        #self.ema_slow = ExponentialMovingAverage(21)
        self.tolerance = 1.0 # QC example is 1.01
        
        # other key attributes
        self.dollar_volume = None
        self.symbol = None
        self.qc = qccontext
        

    def Update(self, cf):
        """
        Update the indicators for symbol.
        
        This function updates the symbol and dollar volume attribute
        which are needed for filtering, and also updates the history
        dataframe and the EMA crossover indicators and attributes.
        """
        
        # assign symbol attribute and dollar volume
        self.symbol = cf.Symbol
        self.dollar_volume = cf.DollarVolume
        
        if self.EndTime >= cf.EndTime:
            return
        
        # Convert CoarseFundamental to pandas.Series
        to_append = pd.Series([cf.AdjustedPrice], name=str(cf.Symbol),
            index=pd.Index([cf.EndTime], name='time'))
            
        self.Window.drop(self.Window.index[0], inplace=True)
        self.Window = pd.concat([self.Window, to_append], ignore_index=True)
        self.EndTime = cf.EndTime

        """# after warmup update ema indicators with current values to determine
        # if they are uptrending
        if self.ema_fast.Update(cf.EndTime, cf.AdjustedPrice) and self.ema_slow.Update(cf.EndTime, cf.AdjustedPrice):
            fast = self.ema_fast.Current.Value
            slow = self.ema_slow.Current.Value
            # self.qc.Debug(f"[{self.qc.Time}] symbol: {cf.Symbol} | fast: {round(self.ema_fast.Current.Value,2)} | slow: {round(self.ema_slow.Current.Value,2)}")
            self.is_uptrend = fast > slow * self.tolerance"""
            
        
    @property
    def Metric(self):
        """
        This computes the MMI value for the security.
        """
        
        # create returns to feed indicator
        returns = self.Window.pct_change().dropna()
        
        # compute the metric
        x = returns.rolling(self.FastPeriod).mean().iloc[-1]
            
        return x # should be single numeric