import numpy as np
import pandas as pd
import operator

### <summary>
### Basic template algorithm simply initializes the date range and cash. This is a skeleton
### framework you can use for designing an algorithm.
### </summary>
class ResolveGEM(QCAlgorithm):
    '''Basic template algorithm simply initializes the date range and cash'''

    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(2005,1, 1)  #Set Start Date
        #self.SetEndDate(2013,10,11)    #Set End Date
        self.SetCash(100000)           #Set Strategy Cash
        
        self.equity = ['SPY', 'EFA']
                        
        self.bonds = ['TLT']
        self.symbols = self.equity + self.bonds
        
        self.leverage = 1

        for symbol in self.symbols:
            self.AddEquity(symbol)
            #self.Securities[symbol].SetDataNormalizationMode(DataNormalizationMode.TotalReturn)
        
        self.lookbacks = range(1, 19)
        self.ma_periods = range(2, 19)
        self.trend_style = ['SPY', 'Multi']
        self.momentum_style = ['Price', 'MA']
        
        columns = ['Momentum Style', 'Trend Style', 'Momentum Lookback', 'Trend Lookback', 'MA Period']
        
        # Create multiple strategies with different settings
        settings = []
        for trend in self.trend_style:
            for momentum in self.momentum_style:
                if momentum == 'Price':
                    for m_look in self.lookbacks:
                        for t_look in self.lookbacks:
                            settings.append([momentum, trend, m_look, t_look, 0])
                else:
                    for ma in self.ma_periods:
                        for ma_look in self.ma_periods:
                            settings.append([momentum, trend, 0, ma_look, ma])
            
        self.strategies = pd.DataFrame(settings, columns=columns)
        self.strategies.fillna(0)
        

        self.allocation = {}
        
        for sym in self.symbols:
            self.allocation[sym] = 0

        self.Schedule.On(self.DateRules.MonthEnd('SPY'), 
                        self.TimeRules.BeforeMarketClose('SPY', 10),
                        self.signals)
                        
                        
    def signals(self):
        
        h = self.History(self.equity, 252 * 2, Resolution.Daily).unstack(level=0).close
        
        for strategy in range(self.strategies.iloc[:,0].count()):
            
            if self.strategies.iloc[strategy]['Momentum Style'] == 'Price':
                
                spy_trend = h['SPY'].iloc[-1] / h['SPY'].iloc[-22 * self.strategies.iloc[strategy]['Trend Lookback']] - 1 > 0
                efa_trend = h['EFA'].iloc[-1] / h['EFA'].iloc[-22 * self.strategies.iloc[strategy]['Trend Lookback']] - 1 > 0
                
                # Are SPY OR EFA returns positive
                if spy_trend or efa_trend:
                    
                    # Calculate momentum as % change in price
                    ret = {}
                    for sym in self.equity:
                        ret[sym] = h[sym].iloc[-1] / h[sym].iloc[-22 * self.strategies.iloc[strategy]['Momentum Lookback']] -1
                       
                    # Select highest momentum equity if SPY > Moving Average
                    if self.strategies.iloc[strategy]['Trend Style'] == 'SPY' and spy_trend:
                        self.allocation[max(ret.items(), key=operator.itemgetter(1))[0]] += 1 
                    
                    # Select highest momentum equity if SPY and EFA > Moving Average       
                    elif self.strategies.iloc[strategy]['Trend Style'] == 'Multi'  and spy_trend and efa_trend:
                        self.allocation[max(ret.items(), key=operator.itemgetter(1))[0]] += 1 
                        
                    else:
                       self.allocation['TLT'] += 1 
                        
                else:
                    self.allocation['TLT'] += 1
            else:
                
                spy_trend = h['SPY'].iloc[-1] > h['SPY'].iloc[-22 * self.strategies.iloc[strategy]['Trend Lookback']:].mean()
                efa_trend = h['EFA'].iloc[-1] > h['EFA'].iloc[-22 * self.strategies.iloc[strategy]['Trend Lookback']:].mean()
                
                if spy_trend or efa_trend:
                    
                    # Calculate momentum as % difference to moving average
                    ret = {}
                    for sym in self.equity:
                        ret[sym] = h[sym].iloc[-1] / h[sym].iloc[-22 * self.strategies.iloc[strategy]['MA Period']:].mean() -1
                        
                    # Select highest momentum equity if SPY > Moving Average
                    if self.strategies.iloc[strategy]['Trend Style'] == 'SPY' and spy_trend:
                        self.allocation[max(ret.items(), key=operator.itemgetter(1))[0]] += 1 
                    
                    # Select highest momentum equity if SPY and EFA > Moving Average    
                    elif self.strategies.iloc[strategy]['Trend Style'] == 'Multi'  and spy_trend and efa_trend:
                        self.allocation[max(ret.items(), key=operator.itemgetter(1))[0]] += 1 
                    
                    else:
                       self.allocation['TLT'] += 1 
                    
                else:
                    self.allocation['TLT'] += 1
                
        self.trade()
    
    def trade(self):
        
        # Place sales first to ensure buying power is no exceeded in rebalance
        for s in self.symbols:
            current_weight = float((self.Portfolio[s].Quantity * self.Portfolio[s].Price) / self.Portfolio.TotalPortfolioValue)
            weight = self.allocation[s] / sum(self.allocation.values())
            if (weight * self.leverage) < current_weight:
                self.SetHoldings(s, weight * self.leverage )
        
        # Place purchases       
        for s in self.symbols:
            current_weight = float((self.Portfolio[s].Quantity * self.Portfolio[s].Price) / self.Portfolio.TotalPortfolioValue)
            weight = self.allocation[s] / sum(self.allocation.values())
            if (weight * self.leverage) > current_weight:
                self.SetHoldings(s, weight * self.leverage )
    
        # Reset allocations    
        for sym in self.symbols:
            self.allocation[sym] = 0