from scipy.stats import linregress
import numpy as np
import pandas as pd

class MultidimensionalVerticalCompensator(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2018, 1, 1)  # Set Start Date
        self.SetEndDate(2019, 12,31)
        self.SetCash(100000)  # Set Strategy Cash
        # self.AddEquity("SPY", Resolution.Minute)

        self.benchmark = 'SPY'
    
        #self.benchmark_obj = self.AddEquity(self.benchmark, Resolution.Hour)
        
        self.SetBenchmark(self.benchmark)
        
    
        
        # Strategy params
        self.window = 252
        self.filter_window = 200
        self.n_assets = 20
        self.mom_threshold = 60
        
        
        self.UniverseSettings.Resolution = Resolution.Daily
        
        self.AddUniverse(self.CoarseSelectionFunction, self.FineSelectionFunction)
        
        self.AddAlpha(ClenowAlphaModel())#ConstantAlphaModel(InsightType.Price, InsightDirection.Up, timedelta(minutes = 20), 0.025, None))
        
        self.SetPortfolioConstruction(EqualWeightingPortfolioConstructionModel())
        self.SetExecution(ImmediateExecutionModel())
        self.SetRiskManagement(NullRiskManagementModel())
        
        self.lastMonth = None
        
        # Flag to use the fine filter
        self.fine = True

        
        self.current_portfolio = []

    def OnData(self, data):
        '''OnData event is the primary entry point for your algorithm. Each new data point will be pumped in here.
            Arguments:
                data: Slice object keyed by symbol containing the stock data
        '''

        pass


    def OnSecuritiesChanged(self, changes):
        pass
 

    # def weight(self, security, atr):
    #     risk = float(self.Portfolio.TotalPortfolioValue)*0.0001
    #     weight = (((risk/atr) * float(self.Securities[security].Price))/float(self.Portfolio.TotalPortfolioValue)*100)
    #     return weight if not np.isnan(weight) else 0.0
        
            
    def gapper(self,security,period):
        
        security_data = self.History(security,period,Resolution.Daily)
        close_data = [float(data) for data in security_data['close']]
        
        return np.max(np.abs(np.diff(close_data))/close_data[:-1])>=0.15
        
    def moving_average_condition(self, security, period):

        security_data = self.History(security,period,Resolution.Daily)
        close_data = [float(data) for data in security_data['close']]
        
        return close_data[-1] > np.nanmean(close_data)

    # Get SP500
    def CoarseSelectionFunction(self, coarse):
        

        if self.Time.month == self.lastMonth:
            self.fine = False
            return Universe.Unchanged
        
        self.fine = True
        self.lastMonth = self.Time.month
            
        sortedByDollarVolume = sorted([x for x in coarse if x.HasFundamentalData and x.Volume > 0 and x.Price > 0],
                                        key= lambda x: x.DollarVolume, reverse=True)[:500]
                                        
        
        if len(sortedByDollarVolume) == 0:
            return Universe.Unchanged
            
        return [x.Symbol for x in sortedByDollarVolume]
        
        
    # Filter by momentum and trend
    def FineSelectionFunction(self, fine):
        
        if not self.fine:
            return Universe.Unchanged
        
        selection = []
        
        for asset in fine:
            slope = self._slope(asset.Symbol, self.window)
            
            
            trend_filter = self.History(asset.Symbol, self.filter_window, Resolution.Daily)
            trend_filter = trend_filter.close[-1] > trend_filter.close.mean()
            
            if trend_filter:
                selection.append(
                    (asset.Symbol, slope)    
                )
            
            
        selection = sorted(selection, key= lambda x: x[1], reverse=True)
        
        selected = [x[0] for x in selection[:self.n_assets] if x[1] > self.mom_threshold]
        teste = selected[0]
        
    
        
        filtered = []
        
        for stock in selected:
            isUpTrend = self.moving_average_condition(stock, 120) 
            isGapper = self.gapper(stock, 90)
            if isUpTrend and not isGapper:
                filtered.append(stock)
    
        return filtered
    
    def _slope(self, symbol, time_span):
        
        hist = self.History(symbol, time_span, Resolution.Daily)
        
        y = np.log(hist.close)
        x = range(len(y))
        
        slope, _, r_value, _, _ = linregress(x, y)
        
        annualized_slope = (np.power(np.exp(slope), 250) - 1) * 100
        annualized_slope = annualized_slope * (r_value ** 2)
        
        return annualized_slope
        
        
    def _atr(self, symbol, atr_window):
        
        data = self.History(symbol, self.window)
        
        h_minus_l = data.high - data.low
        h_minus_p_close = np.abs(data.high - data.close.shift(1))
        l_minus_p_close = np.abs(data.low - data.close.shift(1))
        
        
        tr = [max(x,y,z) for x,y,z in zip(h_minus_l, h_minus_p_close, l_minus_p_close)]
    
        atr = pd.Series(tr).rolling(atr_window).mean()
    
        return atr.values[-1]
        
        
        
class ClenowAlphaModel(AlphaModel):
    '''Alpha model that uses an EMA cross to create insights'''

    def __init__(self, resolution = Resolution.Daily):

        self.resolution = resolution
        self.predictionInterval = Time.Multiply(Extensions.ToTimeSpan(resolution), 20)
        self.stocks_to_trade = []
        self.removed = []
        self.lastMonth = None


    def Update(self, algorithm, data):
        '''Updates this alpha model with the latest data from the algorithm.
        This is called each time the algorithm receives data for subscribed securities
        Args:
            algorithm: The algorithm instance
            data: The new data available
        Returns:
            The new insights generated'''
            
        dt = datetime(algorithm.Time.year,algorithm.Time.month,algorithm.Time.day)
        
        same_month = self.lastMonth == algorithm.Time.month
        
        wednesday = dt.weekday() == 3
        
        trade_condition = (not same_month) and wednesday
        
        if not trade_condition:# or self.Securities[self.spy].Price < self.spy_200_sma.Current.Value:
            return []
        
        insights = []
        for symbol in self.stocks_to_trade:
            
            
            insight = Insight(
                symbol, self.predictionInterval, 
                InsightType.Price, InsightDirection.Up
            )
            
            insights.append(insight)

        self.lastMonth = algorithm.Time.month

        return insights
        

                
    def OnSecuritiesChanged(self, algorithm, changes):
        
        dt = datetime(algorithm.Time.year,algorithm.Time.month,algorithm.Time.day)
        if dt.weekday() != 3 and self.lastMonth == algorithm.Time.month:# or self.Securities[self.spy].Price < self.spy_200_sma.Current.Value:
            return 
        
        self.removed = [ x.Symbol for x in changes.RemovedSecurities ]
        for stock in self.removed:
            try:
                algorithm.current_portfolio.remove(stock)
                self.stocks_to_trade.remove(stock)
            except:
                pass
        
        self.stocks_to_trade += [stock.Symbol for stock in changes.AddedSecurities]
        
                
    def weight(self, algorithm, security, atr):
        risk = float(algorithm.Portfolio.TotalPortfolioValue)*0.0001
        weight = (((risk/atr) * float(self.Securities[security].Price))/float(self.Portfolio.TotalPortfolioValue)*100)
        return weight if not np.isnan(weight) else 0.0
        
            
    def gapper(self, algorithm, security,period):
        if not algorithm.Securities.ContainsKey(security):
            return 0
        security_data = algorithm.History(security,period,Resolution.Daily)
        close_data = [float(data) for data in security_data['close']]
        return np.max(np.abs(np.diff(close_data))/close_data[:-1])>=0.15
        
    def moving_average(self, algorithm, security, period):
        if not algorithm.Securities.ContainsKey(security):
            return 0
        security_data = algorithm.History(security,period,Resolution.Daily)
        close_data = [float(data) for data in security_data['close']]
        return np.nanmean(close_data)