import numpy as np
import pandas as pd
import statsmodels.api as sm
from QuantConnect.Data.UniverseSelection import *
from scipy.stats import linregress
from datetime import datetime, timedelta

class Momentum(QCAlgorithm):

    def Initialize(self):
        self.reb1 = 1 # set the flag for momentum stock rebalancement
        self.initial = 0
        self.scale_equities = 1
        self.target_vol = 0.125
        self.num_coarse = 500 # Number of stocks to pass CoarseSelection process
        self.num_fine = 20 # Number of stocks to long
        self.SetStartDate(2017, 1, 1) # Set Start Date
        self.SetEndDate(datetime.now()) # Set End Date
        self.SetCash(15000) # Set Strategy Cash
        self.AddUniverse(self.CoarseSelectionFunction,self.FineSelectionFunction)
        self.spy = self.AddEquity("SPY", Resolution.Daily).Symbol
        self.gld = self.AddEquity("GLD", Resolution.Minute).Symbol # gold hedge
        self.iei = self.AddEquity("IEI", Resolution.Minute).Symbol # bond hedge
        self.tlt = self.AddEquity("TLT", Resolution.Minute).Symbol # long term bond hedge
        self.hedge = [self.gld, self.iei, self.tlt]
        self.Schedule.On(self.DateRules.MonthStart(self.spy), self.TimeRules.AfterMarketOpen(self.spy,5), Action(self.rebalance))
        self.SetSecurityInitializer(self.CustomSecurityInitializer)
        self.SetBrokerageModel(BrokerageName.InteractiveBrokersBrokerage, AccountType.Cash)
       
    def CustomSecurityInitializer(self, security):
        security.SetDataNormalizationMode(DataNormalizationMode.SplitAdjusted)

    def CoarseSelectionFunction(self, coarse):
    # if the rebalance flag is not 1, return null list to save time.
        if self.reb1 != 1:
            return Universe.Unchanged
        
        # make universe selection once a month
        sortedByDollarVolume = sorted(coarse, key=lambda x: x.DollarVolume, reverse=True)  
        filtered = [x.Symbol for x in sortedByDollarVolume if x.HasFundamentalData]
        
        # filtered down to the 500 most liquid stocks
        return filtered[:self.num_coarse]
        
    def FineSelectionFunction(self, fine):
        # return null list if it's not time to rebalance
        if self.reb1 != 1:
            return Universe.Unchanged
            
        # drop counter (will update back to 1 after rebalancement has occurred)
        self.reb1 = 0
        
        # create dictionaries to store the indicator values
        stock_filter = {}
        
        # filter by market cap 
        self.market_cap = [x for x in fine if x.MarketCap > 2e9]
        
        # prepare data
        hist_spy = self.History(self.spy, timedelta(days=760), Resolution.Daily).droplevel(level=0)
        hist_spy.rename(columns={"close":"spy close"}, inplace=True)
        
        # we now want to calculate the monthly market returns
        SPY_monthly_returns = hist_spy["spy close"].resample("M").ffill().pct_change().dropna()
        
        # sort the list by their price momentum 
        for security in self.market_cap: 
            
            hist_stock = self.History(security.Symbol, timedelta(days=760), Resolution.Daily)
            
            if hist_stock.index.nlevels > 1: 
                hist_stock = hist_stock.droplevel(level=0)
            
            if "close" in hist_stock.columns and "open" in hist_stock.columns and len(hist_stock) > 520:
                
                hist_stock.rename(columns={"close":"stock close"}, inplace=True)
                # we now want to calculate the monthly tesla returns
                stock_monthly_returns = hist_stock["stock close"].resample("M").ffill().pct_change().dropna()
                
                if len(stock_monthly_returns) == len(SPY_monthly_returns): 
                    
                    df = pd.concat([stock_monthly_returns, SPY_monthly_returns], axis=1)
                    df.rename(columns={"stock close":"stock monthly returns","spy close":"spy monthly returns"}, inplace=True)
                    
                    Y = df["stock monthly returns"]
                    X = df["spy monthly returns"]
                    X = sm.add_constant(X)
                    model = sm.OLS(Y,X).fit()
                    residual_values = model.resid # residual values
                    end = date.today()
                    start = date(end.year-1, end.month, end.day)
                    data_12mth = residual_values.loc[start:end]
                    x = np.sum(data_12mth)
                    y = np.std(data_12mth)
                    resid_mom = x/y
                    
                    # we now have a dictionary storing the values 
                    stock_filter[security.Symbol] = resid_mom
                    
        # we only want the highest values for the coeff
        self.sortedLong = sorted(stock_filter.items(), key=lambda d:d[1],reverse=True)
        sorted_symbolLong = [x[0] for x in self.sortedLong]

        # long the top 20
        self.long = sorted_symbolLong[:self.num_fine]
        return self.long
        
    def OnData(self, data):
        pass
    
    def rebalance(self):
        
        if self.initial == 0: 
            for i in self.long: 
                self.SetHoldings(i, self.scale_equities/self.num_fine)
                
            self.SetHoldings(self.iei, (1-self.scale_equities)/2)
            self.SetHoldings(self.gld, (1-self.scale_equities)/4)
            self.SetHoldings(self.tlt, (1-self.scale_equities)/4)

            self.initial += 1
            
        else: 
            
            # volatility scaling
            # we first have to calculate the correlation matrix of the various stocks in the portfolio
            
            # initialise
            n = 0
            
            # count how many symbols in portfolio
            invested = [x.Symbol.Value for x in self.Portfolio.Values if x.Invested]
            holdings = len(invested)

            for i in self.Portfolio.Values: 
                if i.Invested: 
                    hist = self.History(i.Symbol, timedelta(days=30), Resolution.Daily)
                    hist = hist.drop(columns=["high","low","open","volume"])
                    if hist.index.nlevels > 1: 
                        hist = hist.droplevel(level=0)
                    stock_daily_ret = hist["close"].pct_change().dropna()
                        
                    # first loop
                    if n == 0: 
                        cov_matrix = stock_daily_ret
                        n = n+1
                        
                    else: 
                        cov_matrix = pd.concat([stock_daily_ret, cov_matrix], axis=1)
                            
            # so as to annualise the portfolio volatility
            cov_annual = cov_matrix.cov()*252
            weights = np.full(holdings, 1/holdings).reshape(holdings, 1)
            port_variance = np.dot(weights.T, np.dot(cov_annual, weights))
            port_vol = np.sqrt(port_variance)

            # target annual volatility is set 
            scaling = round((self.target_vol/port_vol).item(), 1)
            
            # dynamic scaling according to volatility - deterministic weights; capped at 1 i.e. no leverage
            self.scale_equities = min(scaling, 1)
            
            if self.scale_equities < 0.4: 
                self.scale_equities = 0.4
                
            self.scale_hedge = 1 - self.scale_equities
            
            # to update the momentum stocks only monthly
            # this removes stocks no longer on long list
            for i in self.Portfolio.Values: 
                if i.Invested and i.Symbol not in self.long and i.Symbol not in self.hedge: 
                    self.Liquidate(i.Symbol)
                    
            # monthly rebalancement of momentum stocks
            for i in self.long:
                self.SetHoldings(i, self.scale_equities/self.num_fine)
                
            self.SetHoldings(self.iei, self.scale_hedge/2)
            self.SetHoldings(self.gld, self.scale_hedge/4)
            self.SetHoldings(self.tlt, self.scale_hedge/4)
                    
            self.Log("Invested:" + str(invested))
            self.Log(self.scale_equities)
            self.Log("Holdings:" + str(holdings))

            self.reb1 = 1