from datetime import date, timedelta, datetime
from decimal import Decimal
import numpy as np
import pandas as pd
from scipy.stats import linregress
import decimal as d

class MomentumandStateofMarketFiltersAlgorithm(QCAlgorithm):

    def Initialize(self):
        
        # QC setup
        self.SetStartDate(2020, 1, 1)
        
        # Cash setup
        self.SetCash(100000)
        
        # Initialize SPY
        self.spy = self.AddEquity("SPY", Resolution.Daily).Symbol    
        
        # Add universe selection
        self.AddUniverse(self.Coarse, self.Fine)
        
        # Set universe resolution to daily
        self.UniverseSettings.Resolution = Resolution.Daily
        
        # Max stocks filtered from coarse
        self.num_coarse = 250
        
        # User input list of stocks
        self.manual_list = ['AAPL', 'ABBV', 'ABT', 'ACN', 'ADBE', 'AIG', 'AMGN', 'AMT', 'AMZN', 'AVGO', 
                            'AXP', 'BA', 'BAC', 'BIIB', 'BK', 'BKNG', 'BLK', 'BMY', 'BRK.B', 'C', 'CAT', 
                            'CHTR', 'CL', 'CMCSA', 'COF', 'COP', 'COST', 'CRM', 'CSCO', 'CVS', 'CVX', 'DD', 
                            'DHR', 'DIS', 'DOW', 'DUK', 'EMR', 'EXC', 'F', 'FB', 'FDX', 'GD', 'GE', 'GILD', 
                            'GM', 'GOOG', 'GOOGL', 'GS', 'HD', 'HON', 'IBM', 'INTC', 'JNJ', 'JPM', 'KHC', 'KO', 
                            'LIN', 'LLY', 'LMT', 'LOW', 'MA', 'MCD', 'MDLZ', 'MDT', 'MET', 'MMM', 'MO', 'MRK', 'MS', 
                            'MSFT', 'NEE', 'NFLX', 'NKE', 'NVDA', 'ORCL', 'PEP', 'PFE', 'PG', 'PM', 'PYPL', 'QCOM', 'RTX', 
                            'SBUX', 'SO', 'SPG', 'T', 'TGT', 'TMO', 'TMUS', 'TSLA', 'TXN', 'UNH', 'UNP', 'UPS', 'USB', 'V', 
                            'VZ', 'WBA', 'WFC', 'WMT', 'XOM']

        # Length of history call for position sizing
        self.back_period = 21 * 3 + 1    
        
        # Volume period for position sizing
        self.vol_period = 21   
        
        # Target volume for position sizing
        self.target_vol = 0.2
        
        # Leverage
        self.lev = 1.5
        
        # Delta
        self.delta = 0.05
        
        # Create an array out of volume period
        self.x = np.asarray(range(self.vol_period))
        
        # Lookback period for momentum percent
        self.lookback = 20*6
        
        # MOM dictionary
        self.mom = {}
        
        # Set automatic warmup
        self.AutomaticIndicatorWarmup = True
    
    # Coarse selection
    def Coarse(self, coarse):
        
        # Filter stocks by conditions:
        filtered = [x for x in coarse 
                    
                    # If ticker has fundamental data
                    if x.HasFundamentalData
                    
                    # And if ticker dollar volume of ticker greater than 1000000
                    and x.DollarVolume > 1000000
                    
                    # And if ticker price greater than 5
                    and x.Price > 5
                    
                    # Or ticker symbol in user input list
                    or x.Symbol.Value in self.manual_list]
        
        # Sort by dollar volume
        sortedStocks = sorted(filtered, key = lambda x: x.DollarVolume, reverse = True)
        
        # Return number of stocks that user set 
        return [x.Symbol for x in sortedStocks][:self.num_coarse]
    
    # Fine selection
    def Fine(self, fine):
        
        # Return symbol if:
        return [x.Symbol for x in fine 
        
                # Stock is listed in NASDAQ
                if x.CompanyReference.PrimaryExchangeID == "NAS"
                
                # And Stock is listed in USA
                and x.CompanyReference.CountryId == "USA"
                
                # Or stock's symbol in user input list
                or x.Symbol.Value in self.manual_list]
    
    # On securities changed
    def OnSecuritiesChanged(self, changes):
        
        # Loop through stocks removed from algorithm 
        for stock in changes.RemovedSecurities:
            
            # Remove from indicators
            self.mom.pop(stock.Symbol, None)
            
            # Liquidate
            self.Liquidate(stock.Symbol)
        
        # Loop through added stocks
        for stock in changes.AddedSecurities:
                
            # Initialize momentum percent indicator for stock and put in self.mom dictionary
            self.mom[stock.Symbol] = self.MOMP(stock.Symbol, self.lookback)
    
    # Daily bars are received here
    def OnData(self, data):
        
        # Initialize temporary dictionary to store momp values
        temp_dict = {}
        
        # Loop through mom dictionary
        for i in self.mom:
            
            # Get current MOMP value
            current_MOMP = self.mom[i].Current.Value
            
            # Add to temporary dictionary
            temp_dict[i] = current_MOMP
        
        # Sort dictionary
        sorted_mom = sorted(temp_dict, key = temp_dict.get, reverse=True)
        
        # Long list contains top 20 stocks with highest momentum percent (getting stocks that have been having an upward trend)
        long_list = sorted_mom[:20]
        
        # Short list contains top 20 stocks with lowest momentum percent (getting stocks that have been having a downward trend)
        short_list = sorted_mom[-20:]
        
        # Check whether current SPY has positive momentum percent
        if self.mom[self.spy].Current.Value > 0:
        
            # If length of long list is greater than 0
            if len(long_list) > 0:
                
                # Initialize rebalancing with greater weight put on long
                self.rebalance(long_list, 1, 0.7)
                
            # If length of short list is greater than 0
            if len(short_list) > 0:
                
                # Initialize rebalancing with less weight put on short
                self.rebalance(short_list, -1, 0.3)
        
        # If SPY has negative momentum percent
        elif self.mom[self.spy].Current.Value < 0:
            
            # If length of long list is greater than 0
            if len(long_list) > 0:
                
                # Initialize rebalancing with less weight put on long
                self.rebalance(long_list, 1, 0.3)
                
            # If length of short list is greater than 0
            if len(short_list) > 0:
                
                # Initialize rebalancing with greater weight put on short
                self.rebalance(short_list, -1, 0.7)
    
    # Rebalance formula
    def rebalance(self, lists, sign, portfolio_weight):
        self.w = 1 / len(lists)
        try:
            pos_sizing = self.pos_sizing(lists, sign) 
        except Exception as e:
            msg = f'Exception: {e}'
            self.Log(msg)
            return
        tot_port = self.Portfolio.TotalPortfolioValue
        for symbol, info in pos_sizing.items():
            new_weight = info[0]
            yesterdayClose = info[1]
            security = self.Securities[symbol]
            quantity = security.Holdings.Quantity
            price = security.Price
            if price == 0: price = yesterdayClose
            curr_weight = quantity * price / tot_port
            shall_trade = abs(new_weight - curr_weight) > self.delta
            if shall_trade: 
                delta_shares = (sign * (int(new_weight * (tot_port * portfolio_weight) / price))) - quantity
                if delta_shares != 0:
                    self.MarketOnOpenOrder(symbol, delta_shares)
                    msg = f"{symbol} -- weight: {new_weight:.2f} (old weight was: {curr_weight:.2f}) -- last price: {price}"
    
    # Position sizing formula
    def pos_sizing(self, lists, sign):
        allPrices = self.History(lists, self.back_period, Resolution.Daily).close.unstack(level=0)
        pos = {}
        for symbol in lists:
            try:
                prices = allPrices[symbol]
                change = prices.pct_change().dropna()
                last = np.float(prices[-1])
                rsq = self.rsquared(self.x, prices[-self.vol_period:])
                alpha = min(0.5, np.exp(-10. * (1. - rsq)))
                vol = change.ewm(alpha=alpha).std() 
                ann_vol = np.float(vol.tail(1)) * np.sqrt(252)
                weight = (self.target_vol / ann_vol).clip(0.0, self.lev)  * self.w
                pos[symbol] =  (weight, last)
                msg = f"{symbol}: {pos[symbol][0]}, rsqr: {rsq}, alpha: {alpha}, ann_vol = {ann_vol}"
            except KeyError:
                pass
        return pos
    
    # Rsquared formula
    def rsquared(self, x, y):
        _, _, r_value, _, _ = linregress(x, y)
        return r_value**2