from QuantConnect.Data.UniverseSelection import *
import math
import numpy as np
import pandas as pd
from scipy import stats 


class FundamentalFactorAlgorithm(QCAlgorithm):
    def slope(self,ts):
        """
        Input: Price time series.
        Output: Annualized exponential regression slope, multipl
        """
        x = np.arange(len(ts))
        log_ts = np.log(ts)
        slope, intercept, r_value, p_value, std_err = stats.linregress(x, log_ts)
        annualized_slope = (np.power(np.exp(slope), 250) - 1) * 100
        return annualized_slope * (r_value ** 2)

    def Initialize(self):

        self.SetStartDate(2003, 1, 1)  #Set Start Date
        self.SetEndDate(2018, 3, 1)  #Set Start Date       
        self.SetCash(10000)            #Set Strategy Cash
    
        
        self.UniverseSettings.Resolution = Resolution.Daily
        self.AddUniverse(self.CoarseSelectionFunction, self.FineSelectionFunction)
        self.SetBrokerageModel(BrokerageName.InteractiveBrokersBrokerage, AccountType.Cash); 
        self.spy = self.AddEquity("SPY", Resolution.Minute).Symbol 
        self.holding_months = 1
        self.num_screener = 100
        self.num_stocks = 3
        self.formation_days = 200
        self.lowmom = False
        self.month_count = self.holding_months
        self.Schedule.On(self.DateRules.MonthStart("SPY"), self.TimeRules.At(0, 0), Action(self.monthly_rebalance))
        self.Schedule.On(self.DateRules.MonthStart("SPY"), self.TimeRules.At(10, 0), Action(self.rebalance))
        # rebalance the universe selection once a month
        self.rebalence_flag = 0
        # make sure to run the universe selection at the start of the algorithm even it's not the manth start
        self.first_month_trade_flag = 1
        self.trade_flag = 0 
        self.symbols = None
            
        # This version uses the average of two momentum slopes.
        # Want just one? Set them both to the same number.
        self.momentum_window = 60  # first momentum window.
        self.momentum_window2 = 90  # second momentum window
    
        # Limit minimum slope. Keep in mind that shorter momentum windows
        # yield more extreme slope numbers. Adjust one, and you may want
        # to adjust the other.
        self.minimum_momentum = 60  # momentum score cap
        self.exclude_days = 5
        self.size_method = 2
        self.use_bond_etf = True
        self.bond_etf = 'TLT'
 
    def CoarseSelectionFunction(self, coarse):
        if self.rebalence_flag or self.first_month_trade_flag:
            # drop stocks which have no fundamental data or have too low prices
            selected = [x for x in coarse if (x.HasFundamentalData) and (float(x.Price) > 5)]
            # rank the stocks by dollar volume 
            filtered = sorted(selected, key=lambda x: x.DollarVolume, reverse=True) 
    
            return [ x.Symbol for x in filtered[:200]]
        else:
            return self.symbols


    def FineSelectionFunction(self, fine):
      
        if self.rebalence_flag or self.first_month_trade_flag:
            try:
                filtered_fine = [x for x in fine if (x.ValuationRatios.EVToEBITDA > 0) 
                                                    and (x.EarningReports.BasicAverageShares.ThreeMonths > 0) 
                                                    and x.EarningReports.BasicAverageShares.ThreeMonths * (x.EarningReports.BasicEPS.TwelveMonths*x.ValuationRatios.PERatio) > 2e9]
            except:
                filtered_fine = [x for x in fine if (x.ValuationRatios.EVToEBITDA > 0) 
                                                and (x.EarningReports.BasicAverageShares.ThreeMonths > 0)] 

            top = sorted(filtered_fine, key = lambda x: x.ValuationRatios.EVToEBITDA, reverse=True)[:self.num_screener]
            self.symbols = [x.Symbol for x in top]
            
            self.rebalence_flag = 0
            self.first_month_trade_flag = 0
            self.trade_flag = 1
            return self.symbols
        else:
            return self.symbols
    
    def OnData(self, data):
        pass
    
    def monthly_rebalance(self):
        self.rebalence_flag = 1


    def inv_vola_calc(self,ts):
        """
        Input: Price time series.
        Output: Inverse exponential moving average standard deviation. 
        Purpose: Provides inverse vola for use in vola parity position sizing.
        """
        returns = np.log(ts).diff()
        stddev = returns.ewm(halflife=20, ignore_na=True, min_periods=0,
                         adjust=True).std(bias=False).dropna()
        return 1 / stddev.iloc[-1]
   

    def rebalance(self):
        #spy_hist = self.History([self.spy], 120, Resolution.Daily).loc[str(self.spy)]['close']
        #if self.Securities[self.spy].Price < spy_hist.mean():
           # for symbol in self.Portfolio.Keys:
               # if symbol.Value != "TLT":
                    #self.Liquidate()
            #self.AddEquity("TLT")
            #self.SetHoldings("TLT", 1)
            #return
        
        
        # Get data
        hist_window = max(self.momentum_window,
                      self.momentum_window2) + self.exclude_days
        if self.symbols is None: return
        stocks = self.symbols
        self.Debug("Stocks " + str(len(stocks)))
        #hist = self.History(context.security_list, "close", hist_window, "1d")
        hist = self.History(stocks, self.formation_days, Resolution.Daily)
        current = self.History(stocks, 1, Resolution.Minute)
        c_data = hist["close"].unstack(level=0)
        self.Debug("c_data " + str(c_data))
        data_end = -1 * (self.exclude_days + 1 )  # exclude most recent data

        momentum1_start = -1 * (self.momentum_window + self.exclude_days)
        momentum_hist1 = c_data[momentum1_start:data_end]
        momentum2_start = -1 * (self.momentum_window2 + self.exclude_days)
        momentum_hist2 = c_data[momentum2_start:data_end]
    
        # Calculate momentum scores for all stocks.
        momentum_list = momentum_hist1.apply(self.slope)  # Mom Window 1
        self.Debug("momentum_list " + str((momentum_list)))
        momentum_list2 =  momentum_hist2.apply(self.slope)   # Mom Window 2
        # Combine the lists and make average
        momentum_concat = pd.concat((momentum_list, momentum_list2))
        mom_by_row = momentum_concat.groupby(momentum_concat.index)
        mom_means = mom_by_row.mean()
        
        # Sort the momentum list, and we've got ourselves a ranking table.
        ranking_table = mom_means.sort_values(ascending=False)
        self.Debug("ranking_table " + str(len(ranking_table)))
        # Get the top X stocks, based on the setting above. Slice the dictionary.
        # These are the stocks we want to buy.
        buy_list = ranking_table[:self.num_stocks]
        self.Debug(" buy list " + str(len(buy_list)))
        final_buy_list = buy_list[buy_list > self.minimum_momentum] # those who passed minimum slope requirement
        self.Debug("Final buy list " + str(len(final_buy_list))) 
        # Calculate inverse volatility, for position size.
        inv_vola_table = c_data[buy_list.index].apply(self.inv_vola_calc)
        # sum inv.vola for all selected stocks.
        sum_inv_vola = np.sum(inv_vola_table)


        spy_hist_150 = self.History([self.spy], 150, Resolution.Daily).loc[str(self.spy)]['close']
        spy_hist_200 = self.History([self.spy], 200, Resolution.Daily).loc[str(self.spy)]['close']
        if self.Securities[self.spy].Price > (spy_hist_150.mean() + spy_hist_200.mean())/2:
            can_buy = True
        else:
            can_buy = False
       
    
    
        equity_weight = 0.0 # for keeping track of exposure to stocks
        
        # Sell positions no longer wanted.
        for security in self.Portfolio.Keys:
            if (security not in final_buy_list):
                if (security.Value != self.bond_etf):
                    # print 'selling %s' % security
                    #self.Debug("Sell Security " + str(security) )
                    self.SetHoldings(security, 0)
                    
        vola_target_weights = inv_vola_table / sum_inv_vola
 
        for security in final_buy_list.index:
            # allow rebalancing of existing, and new buys if can_buy, i.e. passed trend filter.
            if (security in self.Portfolio.Keys) or (can_buy): 
                if (self.size_method == 1):
                    weight = vola_target_weights[security]
                elif (self.size_method == 2):
                    weight = (0.99 / self.num_stocks)
             
                self.Debug("Number of stocks " + str(self.num_stocks) + " Rebalance security " + str(security) + " with weight " + str(weight)  )
                self.SetHoldings(security, weight)
                equity_weight += weight
        
           
    
        # Fill remaining portfolio with bond ETF
        etf_weight = max(1 - equity_weight, 0.0)
    
        print ('equity exposure should be %s ' % str(equity_weight))
    
        if (self.use_bond_etf):
            self.Debug("Buy ETF" )
            self.AddEquity("TLT")
            self.SetHoldings(self.bond_etf, etf_weight)
           


        #=========================
        #if self.symbols is None: return
        #chosen_df = self.calc_return(self.symbols)
        #chosen_df = chosen_df.iloc[:self.num_stocks]
        
        #self.existing_pos = 0
        #add_symbols = []
        #for symbol in self.Portfolio.Keys:
         #   if symbol.Value == 'SPY': continue
          #  if (symbol.Value not in chosen_df.index):  
           #     self.SetHoldings(symbol, 0)
            #elif (symbol.Value in chosen_df.index): 
             #   self.existing_pos += 1
            
        #weight = 0.99/len(chosen_df)
        #for symbol in chosen_df.index:
         #   self.AddEquity(symbol)
          #  self.Debug("Symbol " + str(symbol) + " Weight " + str(weight))
           # self.SetHoldings(symbol, weight)