#region imports
from AlgorithmImports import *
from QuantConnect.DataSource import *
import datetime as dt
import pandas as pd
import numpy as np
import math
from pandas.tseries.offsets import BDay
#endregion

"""
Intersection of ROC comparison using OUT_DAY approach by Vladimir

Inspired by Peter Guenther, Tentor Testivis, Dan Whitnable, Thomas Chang.
Further improved by ideas from Strongs and Tim Nilson
"""

# ------------------------------------------------------------------
# The original values. Later Optimization by TN (see below)
# STK = QQQ or QLD
# BND = TLT, UUP
# VOLA = 126; BASE_RET = 85; 
LEV = 1.0; 
# ------------------------------------------------------------------

# 21.4.2023 Optimal vola = 127; base_ret = 85

class DualMomentumInOut(QCAlgorithm):

    def Initialize(self):
        
        self.SetBrokerageModel(BrokerageName.QuantConnectBrokerage, AccountType.Margin)

        # First date for which JJC is available is 2018, 2, 2 
        # If the start date is earlier, we switch from the old pairs to the new pairs on that date
        # self.SetStartDate(2008, 1, 1) #original: 2008
        self.SetStartDate(2008, 7, 1) # colab start
        #self.SetEndDate(2021, 11, 17)

        self.cap = 100000
        self.SetCash(self.cap) 

        # Parameters
        #self.VOLA = int(self.GetParameter("volatility"))
        self.VOLA = 127
        #self.BASE_RET = int(self.GetParameter("base"))
        self.BASE_RET = 85

        # Switch from original pairs to new pairs with copper and turn on the switch to use UUP if TLT is trending down
        # If using start dates prior to 2018, use production_model = 1. Otherwise, use production_model = 4
        self.production_model = 4
        if self.LiveMode:
            self.pairs_model = self.production_model # 1 = Original 
        else:
            self.pairs_model = 1

        # self.STK  = self.AddEquity('QQQ', Resolution.Minute).Symbol #original: QLD
        self.STK  = self.AddEquity('QLD', Resolution.Minute).Symbol #original: QLD
        self.BND1 = self.AddEquity('TLT', Resolution.Minute).Symbol
        #self.BND2 = self.AddEquity('TMF', Resolution.Minute).Symbol
        self.UUP  = self.AddEquity('UUP', Resolution.Minute).Symbol # Dollar

        self.SLV = self.AddEquity('SLV', Resolution.Daily).Symbol # Silver 
        self.GLD = self.AddEquity('GLD', Resolution.Daily).Symbol # Gold
        self.XLI = self.AddEquity('XLI', Resolution.Daily).Symbol # Industrial
        self.XLU = self.AddEquity('XLU', Resolution.Daily).Symbol # Utilities
        self.DBB = self.AddEquity('DBB', Resolution.Daily).Symbol # Metals
        self.EEM = self.AddEquity('EEM', Resolution.Daily).Symbol # Emerging Markets
        
        #self.FXA = self.AddEquity('FXA', Resolution.Daily).Symbol
        #self.FXF = self.AddEquity('FXF', Resolution.Daily).Symbol

        self.MKT = self.AddEquity('SPY', Resolution.Daily).Symbol # S&P500
        #self.MKT = self.AddEquity('QQQ', Resolution.Daily).Symbol # S&P500

        # BNC is benchmark for equity plot. Use either self.MKT or self.STK
        self.BNC = self.STK

        # added UUP as an alternative to BND
        self.ASSETS = [self.STK, self.BND1, self.UUP]

        if self.pairs_model == 4:
            self.CPR = self.AddEquity('JJC', Resolution.Daily).Symbol # Copper
            self.pairs = [self.SLV, self.GLD, self.XLI, self.XLU, self.EEM, self.MKT, self.DBB, self.BND1, self.UUP, self.CPR]
        else:
            self.pairs = [self.SLV, self.GLD, self.XLI, self.XLU, self.EEM, self.MKT, self.DBB, self.BND1, self.UUP]
        
        self.bull = 1        
        self.count = 0 
        self.outday = 0
        self.wait_days = 0      
        self.wt = {}
        self.real_wt = {}
        self.mkt = []
        self.pairs_plot = {} 
        
        # self.SetWarmUp(timedelta(350))
        self.history_was_updated = False # need this when we switch over to JJC
        self.email_lines = [] # store email strings for daily email 
        self.email_sent = False 

        # Add two smas of TLT to test for TLT trends 
        self.bnd_sma_short = self.SMA(self.BND1, 40, Resolution.Daily)
        self.bnd_sma_long  = self.SMA(self.BND1, 80, Resolution.Daily)
       
        self.Schedule.On(self.DateRules.EveryDay(), self.TimeRules.AfterMarketOpen(self.MKT, 100), self.daily_check)
        
        self.symbols = [self.MKT] + self.pairs + [self.BND1, self.UUP] 
        self.build_history(self.symbols)

    def OnSecuritiesChanged(self, changes:SecurityChanges) -> None:
        for security in changes.AddedSecurities:
            security.SetDataNormalizationMode(DataNormalizationMode.Raw)
            security.SetFeeModel(CustomFeeModel())

    def consolidation_handler(self, sender, consolidated):
        self.history.loc[consolidated.EndTime, consolidated.Symbol] = consolidated.Close
        self.history = self.history.iloc[-(self.VOLA + 1):] 
        
    def build_history(self, symbols):
        symbols = list(set(symbols)) # to remove duplicates 
        for symbol in symbols:
            self.consolidator = TradeBarConsolidator(timedelta(days=1))
            self.consolidator.DataConsolidated += self.consolidation_handler
            self.SubscriptionManager.AddConsolidator(symbol, self.consolidator)
    
        self.history = self.History(symbols, self.VOLA + 1, Resolution.Daily)
        if self.history.empty or 'close' not in self.history.columns:
            return
        self.history = self.history['close'].unstack(level=0).dropna()

    def daily_check(self):
        self.email_lines = []

        # if time is past 2018, 2, 1, add JJC (Copper)
        jjc_start_date = dt.datetime(2018, 2, 1, 0, 0, 0)
        # if self.Time >= jjc_start_date and not self.history_was_updated:
        if self.Time >= (jjc_start_date + BDay(self.VOLA+1)) and not self.history_was_updated:
            if self.pairs_model != 4:
                self.CPR = self.AddEquity('JJC', Resolution.Daily).Symbol # Copper    
            self.pairs_model = self.production_model 
            # update the history 
            self.symbols = self.symbols + [self.CPR]
            self.build_history(self.symbols)
            self.history_was_updated = True

        # pairs switch based on TLT vs MKT
        #r30 = self.history.pct_change(30).iloc[-1]
        #switch_on = False 
        #if switch_on and (r30[self.BND1] - r30[self.MKT] > 0):
        #    self.pairs_model = 1
        #else:
        #    self.pairs_model = 5
        #    if self.Time >= jjc_start_date:
        #        self.pairs_model = 4

        vola = self.history[[self.MKT]].pct_change().std() * np.sqrt(252)
        self.wait_days = int(vola * self.BASE_RET)
        period = int((1.0 - vola) * self.BASE_RET)        
        r = self.history.pct_change(period).iloc[-1]
  
        # Below we can test various pairs. Only 1 and 4 are now used. 
        if self.pairs_model == 1:
            exit = ((r[self.SLV] < r[self.GLD]) and (r[self.XLI] < r[self.XLU]) and (r[self.DBB] < r[self.UUP]))
        elif self.pairs_model == 4:
            exit = ((r[self.SLV] < r[self.GLD]) and (r[self.XLI] < r[self.XLU]) and (r[self.EEM] < r[self.MKT]) and (r[self.CPR] < r[self.GLD]))
            self.email_lines.append(f"C_G {round(r[self.CPR], 3)} {round(r[self.GLD], 3)} {r[self.CPR] < r[self.GLD]}")
            self.pairs_plot['C_G'] = (r[self.CPR] - r[self.GLD]) * -1
        
        # for emailed info and end of algo reporting
        self.email_lines.append(f"S_G {round(r[self.SLV], 3)} {round(r[self.GLD], 3)} {r[self.SLV] < r[self.GLD]}")
        self.email_lines.append(f"I_U {round(r[self.XLI], 3)} {round(r[self.XLU], 3)} {r[self.XLI] < r[self.XLU]}")
        self.email_lines.append(f"E_M {round(r[self.EEM], 3)} {round(r[self.MKT], 3)} {r[self.EEM] < r[self.MKT]}")

        # For pairs plot 
        self.pairs_plot['S_G'] = (r[self.SLV] - r[self.GLD]) * -1
        self.pairs_plot['I_U'] = (r[self.XLI] - r[self.XLU]) * -1
        self.pairs_plot['E_M'] = (r[self.EEM] - r[self.MKT]) * -1
        self.pairs_plot['D_U'] = (r[self.DBB] - r[self.UUP]) * -1
        #self.pairs_plot['DBC_M'] = (r[self.DBC] - r[self.MKT]) * -1
        #self.pairs_plot['Vola'] = vola

        wait_text = ""
        if exit:
            self.bull = False
            self.outday = self.count
        if self.count >= self.outday + self.wait_days:
            self.bull = True
            wait_text = "We are RISK ON"

        if not self.bull and (self.count <= self.outday + self.wait_days):
            wait_text = "Wait: " + str((self.outday + self.wait_days) - self.count)

        self.email_lines.insert(0, f"Bull: {self.bull}, Exit: {exit}, {wait_text}")
        self.count += 1

        if not self.bull:

            # choose the out holding with the highest rate of return
            out_holding = self.BND1 # default 
           
            out_holdings = { self.BND1: r[self.BND1], self.UUP: r[self.UUP] } 
            out_holding  = max(out_holdings, key=out_holdings.get)
            
            # Then, switch to UUP if BND1 is trending down (improves the above) 
            if self.bnd_sma_short.Current.Value < self.bnd_sma_long.Current.Value:
                out_holding = self.UUP

            # if out_holding is negative, go into cash (only a very small improvement)
            if r[out_holding] < 0:
                out_holding = None

            self.email_lines.append(f"HOLD -> {out_holding}")
            
            for sec in self.ASSETS:    
                self.wt[sec] = LEV if sec is out_holding else 0 
            self.trade() 

        elif self.bull:
            
            self.email_lines.append(f"HOLD -> {self.STK}")

            for sec in self.ASSETS:
                self.wt[sec] = LEV if sec is self.STK else 0 
            self.trade()  

        # Send email with insights
        today = dt.datetime.now().replace(hour=0, minute=0, second=0, microsecond=0)
        if self.LiveMode and self.Time >= today:
            pass
            #self.Notify.Email("you@gmail.com", "IN/OUT Insights", "\n".join(self.email_lines))
            #self.Notify.Email("you@gmail.com", "IN/OUT Insights", "\n".join(self.email_lines))


    def trade(self):
        if self.IsWarmingUp:
            return

        # Liquidate first, then set new holdings. This avoids 'not enough funds' warnings.
        for sec, weight in self.wt.items():
            if weight == 0 and self.Portfolio[sec].IsLong:
                self.Liquidate(sec)

        for sec, weight in self.wt.items():
            if weight > 0 and not self.Portfolio[sec].Invested:
                self.SetHoldings(sec, weight)
            
                    
    def OnEndOfDay(self): 
        if self.IsWarmingUp:
            return
        
        if not self.LiveMode:
            mkt_price = self.Securities[self.BNC].Close

            # the below fixes the divide by zero error in the MKT plot
            if mkt_price > 0 and mkt_price is not None:
                self.mkt.append(mkt_price)
        
            if len(self.mkt) >= 2 and not self.IsWarmingUp:
                mkt_perf = self.mkt[-1] / self.mkt[0] * self.cap
                self.Plot('Strategy Equity', self.BNC, mkt_perf)
        
        account_leverage = self.Portfolio.TotalHoldingsValue / self.Portfolio.TotalPortfolioValue
        self.Plot('Holdings', 'leverage', round(account_leverage, 1))
        for sec, weight in self.wt.items(): 
            self.real_wt[sec] = round(self.ActiveSecurities[sec].Holdings.Quantity * self.Securities[sec].Price / self.Portfolio.TotalPortfolioValue,4)
            self.Plot('Holdings', self.Securities[sec].Symbol, round(self.real_wt[sec], 3))

        # Pairs chart
        for key, value in self.pairs_plot.items():
            if math.isnan(value):
                continue
            self.Plot('Pairs', key, value)

        # Pairs Model
        self.Plot('Pairs Model', "Pairs", self.pairs_model)

    # def OnEndOfAlgorithm(self):
    #     for line in self.email_lines:
    #         self.Debug(line)

class CustomFeeModel(FeeModel):
    def GetOrderFee(self, parameters):
        fee = 0.0 #parameters.Security.Price * parameters.Order.AbsoluteQuantity * 0.00005
        return OrderFee(CashAmount(fee, "USD"))