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)
        self.SetCash(100000)
        self.SetBrokerageModel(BrokerageName.InteractiveBrokersBrokerage,
                       AccountType.Margin)
        self.SetBenchmark('SPY')
        self.spy = self.AddEquity("SPY", Resolution.Daily).Symbol        
        self.AddUniverse(self.Coarse, self.Fine)
        self.UniverseSettings.Resolution = Resolution.Daily
        
        #User input area:
        self.ema_fast_span = 2                             #Fast_ema_period
        self.ema_slow_span = 5                             #Slow ema period
        self.spy_fast_span = 1                             #SPY ema period
        self.spy_slow_span = 3                             #SPY ema period
        self.manual_list = []
        self.portfolio_sizing = True                       #If set to false, algorithm will not run buy signal if invested in any equities long or short
                            
        #Variables
        self.symbols = []
        self.long_list = []
        self.short_list = []
        self.holdings = {}
        self.ema_fast = {}
        self.ema_slow = {}
        self.prices = {}
        self.spy_ema_fast = {}
        self.spy_ema_slow = {}
        
        self.back_period = 10     # 3 months
        self.vol_period = 2    # days for calc vol
        self.target_vol = 0.2
        self.lev = 1          # max lev from ratio targ_vol / real_vol
        self.delta = 0.05       # min rebalancing
        
        self.x = np.asarray(range(self.vol_period))
        self.spy_fast_or_slow = False
        self.lookback = 1
        self.num_coarse = 100
        self.long = []
        self.short = []
        self.mom = {}
    
    def Coarse(self, coarse):
        filtered = [x for x in coarse if x.HasFundamentalData
                    and x.DollarVolume > 1000000
                    and (float(x.Price) > 5)
                    or x.Symbol.Value in self.manual_list
                    or x.Symbol.Value in self.holdings]
        sortedStocks = sorted(filtered, key = lambda x: x.DollarVolume, reverse = True)
        return [x.Symbol for x in sortedStocks][:self.num_coarse]

    def Fine(self, fine):
        return [x.Symbol for x in fine 
                if x.CompanyReference.PrimaryExchangeID == "NAS"
                and x.CompanyReference.CountryId == "USA"
                or x.Symbol.Value in self.manual_list
                or x.Symbol.Value in self.holdings]
    
    def OnSecuritiesChanged(self, changes):
        for stock in changes.RemovedSecurities:
            symbol = stock.Symbol
            self.Liquidate(symbol)
            if symbol in self.symbols:
                self.symbols.remove(symbol)
                self.ema_fast.pop(symbol, None)
                self.ema_slow.pop(symbol, None)
                self.prices.pop(symbol, None)
            if symbol in self.long:
                self.long.remove(symbol)
            if symbol in self.short:
                self.short.remove(symbol)
            if symbol in self.mom:
                self.mom.pop(symbol, None)
                
        self.history = self.History([stock.Symbol for stock in changes.AddedSecurities], 100, Resolution.Daily)
        for stock in changes.AddedSecurities:
            symbol = stock.Symbol
            if symbol in self.history.index:
                if symbol not in self.symbols:
                    close = self.history.loc[symbol]["close"].to_list()
                    self.symbols.append(symbol)
                    self.prices[symbol] = close
                    self.mom[symbol] = Momentum(self.lookback)
                    close_prices = {}
                    close_prices["close"] = close
                    ema = pd.DataFrame(close_prices , columns = ["close"])
                    ema["EMA_fast"] = ema["close"].ewm(span=self.ema_fast_span,min_periods=0,adjust=False,ignore_na=False).mean()
                    self.ema_fast[symbol] = ema["EMA_fast"].to_list()
                    ema["EMA_slow"] = ema["close"].ewm(span=self.ema_slow_span,min_periods=0,adjust=False,ignore_na=False).mean()
                    self.ema_slow[symbol] = ema["EMA_slow"].to_list()
            if symbol not in self.history.index:
                if symbol in self.long:
                    self.long.remove(symbol)
                if symbol in self.short:
                    self.short.remove(symbol)
        
        addedSymbols = [k for k,v in self.mom.items() if not v.IsReady]
        history = self.History(addedSymbols, 1 + self.lookback, Resolution.Daily)
        history = history.close.unstack(level=0)
        for symbol in addedSymbols:
            ticker = str(symbol)
            if ticker in history:
                for time, value in history[ticker].items():
                    item = IndicatorDataPoint(symbol, time, value)
                    self.mom[symbol].Update(item)
        
    def OnData(self, data):
        self.UpdateData(data)
        self.CheckEMA()
        self.CheckMom()
        self.CheckLiquidate()
        self.CheckSPY()
        self.Signal()
        self.JustBeforeMarketClose()
        self.Reset()
    
    def UpdateData(self, data):
        self.data = data
        self.tradebars = data.Bars
        for symbol in self.symbols:
            if not self.data.ContainsKey(symbol):
                continue
            if not self.data.Bars.ContainsKey(symbol):
                continue
            self.prices[symbol].append(self.tradebars[symbol].Close)
            close = self.prices[symbol]
            close_prices = {}
            close_prices["close"] = close
            ema = pd.DataFrame(close_prices , columns = ["close"])
            ema["EMA_fast"] = ema["close"].ewm(span=self.ema_fast_span,min_periods=0,adjust=False,ignore_na=False).mean()
            ema_fast = ema["EMA_fast"].to_list()
            self.ema_fast[symbol] = ema_fast
            ema["EMA_slow"] = ema["close"].ewm(span=self.ema_slow_span,min_periods=0,adjust=False,ignore_na=False).mean()
            ema_slow = ema["EMA_slow"].to_list()
            self.ema_slow[symbol] = ema_slow
        temp_list = self.prices[self.spy]
        close_prices = {}
        close_prices["close"] = temp_list
        ema = pd.DataFrame(close_prices , columns = ["close"])
        ema["EMA_fast"] = ema["close"].ewm(span=self.spy_fast_span,min_periods=0,adjust=False,ignore_na=False).mean()
        self.spy_ema_fast[self.spy] = ema["EMA_fast"].to_list()
        ema["EMA_slow"] = ema["close"].ewm(span=self.spy_slow_span,min_periods=0,adjust=False,ignore_na=False).mean()
        self.spy_ema_slow[self.spy] = ema["EMA_slow"].to_list()
        for symbol, mom in self.mom.items():
            if not self.data.Bars.ContainsKey(symbol):
                continue
            mom.Update(self.Time, self.tradebars[symbol].Close)
            
    def CheckEMA(self):
        for i in self.symbols:
            if i in self.ema_fast:
                if self.ema_fast[i][-1] > self.ema_slow[i][-1]:
                    self.long.append(i)
                elif self.ema_fast[i][-1] < self.ema_slow[i][-1]:
                    self.short.append(i)
                    
    def CheckMom(self):
        sorted_mom = sorted([k for k,v in self.mom.items() if v.IsReady],
            key=lambda x: self.mom[x].Current.Value, reverse=True)
        for i in sorted_mom:
            if i in self.long:
                self.long_list.append(i)
        for i in sorted_mom:
            if i in self.short:
                self.short_list.append(i)
        
        self.long_list = self.long_list[:20]
        self.short_list = self.short_list[-20:]
    
    def CheckLiquidate(self):
        self.holdings = {}
        for kvp in self.Portfolio:
            security_holding = kvp.Value
            if security_holding.Invested:
                symbol = security_holding.Symbol
                quantity = security_holding.Quantity
                self.holdings[symbol] = quantity
        if len(self.holdings) > 0:
            for i in self.holdings:
                if self.holdings[i] > 0:
                    if i in self.ema_slow:
                        if self.ema_slow[i][-1] > self.ema_fast[i][-1]:
                            self.Liquidate(i)
                else:
                    0
                if self.holdings[i] < 0:
                    if i in self.ema_slow:
                        if self.ema_fast[i][-1] > self.ema_slow[i][-1]:
                            self.Liquidate(i)
                else:
                    0
    
    def CheckSPY(self):
        if self.spy_ema_fast[self.spy][-1] > self.spy_ema_slow[self.spy][-1]:
            self.spy_fast_or_slow = True
        if self.spy_ema_fast[self.spy][-1] < self.spy_ema_slow[self.spy][-1]:
            self.spy_fast_or_slow = False
            
    def Signal(self):
        if self.long is None or self.short is None: return
        if len(self.holdings) > 0 and self.portfolio_sizing == True:
            if self.spy_fast_or_slow == True:
                if len(self.long_list) > 0:
                    self.rebalance(self.long_list, 1, 0.8)
                if len(self.short_list) > 0:
                    self.rebalance(self.short_list, -1, 0.2)
            elif self.spy_fast_or_slow == False:
                if len(self.short_list) > 0:
                    self.rebalance(self.short_list, -1, 0.8)
                if len(self.long_list) > 0:
                    self.rebalance(self.long_list, 1, 0.2)
        elif len(self.holdings) == 0 and self.portfolio_sizing == True:
            if self.spy_fast_or_slow == True:
                if len(self.long_list) > 0:
                    self.rebalance(self.long_list, 1, 0.8)
                if len(self.short_list) > 0:
                    self.rebalance(self.short_list, -1, 0.2)
            elif self.spy_fast_or_slow == False:
                if len(self.short_list) > 0:
                    self.rebalance(self.short_list, -1, 0.8)
                if len(self.long_list) > 0:
                    self.rebalance(self.long_list, 1, 0.2)
        elif len(self.holdings) > 0 and self.portfolio_sizing == False:
            0
        elif len(self.holdings) == 0 and self.portfolio_sizing == False:
            if self.spy_fast_or_slow == True:
                if len(self.long_list) > 0:
                    self.rebalance(self.long_list, 1, 0.8)
                if len(self.short_list) > 0:
                    self.rebalance(self.short_list, -1, 0.2)
            elif self.spy_fast_or_slow == False:
                if len(self.short_list) > 0:
                    self.rebalance(self.short_list, -1, 0.8)
                if len(self.long_list) > 0:
                    self.rebalance(self.long_list, 1, 0.2)
        else:
            0
                    
    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
                self.MarketOnOpenOrder(symbol, delta_shares)
                msg = f"{symbol} -- weight: {new_weight:.2f} (old weight was: {curr_weight:.2f}) -- last price: {price}"

    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.6, np.exp(-10. * (1. - rsq)))
                vol = change.ewm(alpha=alpha).std() 
                ann_vol = np.float(vol.tail(1)) * np.sqrt(100)
                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
    
    def Reset(self):
        self.long.clear()
        self.short.clear()
        self.long_list.clear()
        self.short_list.clear()
            
    def rsquared(self, x, y):
        _, _, r_value, _, _ = linregress(x, y)
        return r_value**2
    
    def OnMarginCallWarning(self):
        msg = f"{self.Time} : check warning margin call! Fast"
        self.Log(msg)

    def JustBeforeMarketClose(self):
        msg = f"End of day: {self.Time} \nPortfolio value is {self.Portfolio.TotalPortfolioValue:.2f} and Margin Remaining is: {self.Portfolio.MarginRemaining:.2f}  (Total Holdings Value: {self.Portfolio.TotalHoldingsValue:.2f})"
        self.Log(msg)

    def OnOrderEvent(self, orderEvent):
        order = self.Transactions.GetOrderById(orderEvent.OrderId)
        self.Log(f"{self.Time}: {order.Type}: {orderEvent}")
        
    def TimeIs(self, day, hour, minute):
        return self.Time.day == day and self.Time.hour == hour and self.Time.minute == minute