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.benchmark = 'SPY'
        
        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())
        
        pcm = EqualWeightingPortfolioConstructionModel(Resolution.Daily)
        
        self.SetPortfolioConstruction(pcm)
        
        self.SetExecution(ImmediateExecutionModel())
        self.SetRiskManagement(NullRiskManagementModel())
        
        self.lastMonth = None
        
        # Flag to use the fine filter
        self.fine = True
        
        self.current_portfolio = []

    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):
    
    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):
        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, timedelta(20), 
                InsightType.Price, InsightDirection.Up
            )
            
            insights.append(insight)

        self.lastMonth = algorithm.Time.month

        return insights
        
    def OnSecuritiesChanged(self, algorithm, changes):
        
        self.removed = [ x.Symbol for x in changes.RemovedSecurities ]
        for stock in self.removed:
            if stock in self.stocks_to_trade:
                self.stocks_to_trade.remove(stock)
        
        self.stocks_to_trade += [stock.Symbol for stock in changes.AddedSecurities]