# region imports
from AlgorithmImports import *
import numpy as np
import pandas as pd
import math
import statsmodels.api as sm
from pandas.tseries.offsets import BDay
from pykalman import KalmanFilter
from statsmodels.tsa.stattools import coint, adfuller
# endregion

from config import general_setting


class BasicTemplateFuturesAlgorithm(QCAlgorithm):
    def Initialize(self):
        self.debug("start calendar spread algo")
        self.SetStartDate(2023, 10, 8)
        self.SetCash(1000000)

        self.universe_settings.resolution = Resolution.MINUTE

        # lookback frequency settings
        self.lookback = general_setting['lookback']
        self.lookback_RESOLUTION = general_setting['lookback_RESOLUTION'] 

        self.enter = general_setting["enter_level"]
        self.exit = general_setting["exit_level"]

        # Subscribe and set our expiry filter for the futures chain
        future1 = self.AddFuture(Futures.Metals.GOLD, resolution=Resolution.MINUTE)
        future1.SetFilter(timedelta(0), timedelta(365))

        # benchmark = self.AddEquity("SPY")
        # self.SetBenchmark(benchmark.Symbol)

        seeder = FuncSecuritySeeder(self.GetLastKnownPrices)
        self.SetSecurityInitializer(lambda security: seeder.SeedSecurity(security))
        
        self.gold1_contract = None
        self.gold2_contract = None
        self.gold3_contract = None

        self.minute_counter = 0
        self.Schedule.On(self.date_rules.every_day(), self.TimeRules.At(18,0), self.reset_minute_counter)  # Check Take profit and STOP LOSS every minute


    def reset_minute_counter(self):
        self.minute_counter = 0


    def stats(self, symbols, method="Regression"):
        # lookback here refers to market hour, whereas additional extended-market-hour data are also included.
        if self.lookback_RESOLUTION == "MINUTE":
            df_Gold1 = self.History(symbols[0], self.lookback, Resolution.MINUTE)
            df_Gold2 = self.History(symbols[1], self.lookback, Resolution.MINUTE)
            df_Gold3 = self.History(symbols[2], self.lookback, Resolution.MINUTE)
        elif self.lookback_RESOLUTION == "HOUR":
            df_Gold1 = self.History(symbols[0], self.lookback, Resolution.HOUR)
            df_Gold2 = self.History(symbols[1], self.lookback, Resolution.HOUR)
            df_Gold3 = self.History(symbols[2], self.lookback, Resolution.HOUR)
        else:
            df_Gold1 = self.History(symbols[0], self.lookback, Resolution.DAILY)
            df_Gold2 = self.History(symbols[1], self.lookback, Resolution.DAILY)
            df_Gold3 = self.History(symbols[2], self.lookback, Resolution.DAILY)
        
        if df_Gold1.empty or df_Gold2.empty:
            return 0

        df_Gold1 = df_Gold1["close"]
        df_Gold2 = df_Gold2["close"]
        df_Gold3 = df_Gold3["close"]

        Gold1_log = np.array(df_Gold1.apply(lambda x: math.log(x))) 
        Gold2_log = np.array(df_Gold2.apply(lambda x: math.log(x))) 
        Gold3_log = np.array(df_Gold3.apply(lambda x: math.log(x))) 


        # Gold1 & Gold2 Regression and ADF test
        X1 = sm.add_constant(Gold1_log)
        Y1 = Gold2_log
        model1 = sm.OLS(Y1, X1)
        results1 = model1.fit()
        sigma1 = math.sqrt(results1.mse_resid)
        slope1 = results1.params[1]
        intercept1 = results1.params[0]
        res1 = results1.resid
        zscore1 = res1/sigma1

        adf1 = adfuller(res1)

        p_value1 = adf1[1]
        test_passed1 = p_value1 <= general_setting['p_value_threshold']
        self.debug(f"p value is {p_value1}")
        # p 越小越显著

        # Gold1 & Gold3 Regression and ADF test
        X2 = sm.add_constant(Gold1_log)
        Y2 = Gold3_log
        model2 = sm.OLS(Y2, X2)
        results2 = model2.fit()
        sigma2 = math.sqrt(results2.mse_resid)
        slope2 = results2.params[1]
        intercept2 = results2.params[0]
        res2 = results2.resid
        zscore2 = res2/sigma2

        adf2 = adfuller(res2)

        p_value2 = adf2[1]
        test_passed2 = p_value2 <= general_setting['p_value_threshold']

        
        # Gold1 & Gold3 Regression and ADF test
        X3 = sm.add_constant(Gold2_log)
        Y3 = Gold3_log
        model3 = sm.OLS(Y3, X3)
        results3 = model3.fit()
        sigma3 = math.sqrt(results3.mse_resid)
        slope3 = results3.params[1]
        intercept3 = results3.params[0]
        res3 = results3.resid
        zscore3 = res3/sigma3

        adf3 = adfuller(res3)

        p_value3 = adf3[1]
        test_passed3 = p_value3 <= general_setting['p_value_threshold']


        # Kalman Filtering to get parameters
        if method == "Kalman_Filter":
            obs_mat = sm.add_constant(Gold1_log, prepend=False)[:, np.newaxis]
            trans_cov = 1e-5 / (1 - 1e-5) * np.eye(2)
            kf = KalmanFilter(n_dim_obs=1, n_dim_state=2,
                  initial_state_mean=np.ones(2),
                  initial_state_covariance=np.ones((2, 2)),
                  transition_matrices=np.eye(2),
                  observation_matrices=obs_mat,
                  observation_covariance=0.5,
                  transition_covariance=0.000001 * np.eye(2))
        
            state_means, state_covs = kf.filter(Gold2_log)
            slope = state_means[:, 0][-1]
            intercept = state_means[:, 1][-1]

        self.printed = True
        return [test_passed1, zscore1, slope1]


    def OnData(self,slice):

        for chain in slice.FutureChains:
            contracts = list(filter(lambda x: x.Expiry > self.Time + timedelta(90), chain.Value))
            if len(contracts) == 0: 
                continue
            front1 = sorted(contracts, key = lambda x: x.Expiry)[0]
            front2 = sorted(contracts, key = lambda x: x.Expiry)[1]
            front3 = sorted(contracts, key = lambda x: x.Expiry)[2]

            self.Debug (" Expiry " + str(front3.Expiry) + " - " + str(front3.Symbol))
            self.gold1_contract = front1.Symbol
            self.gold2_contract = front2.Symbol
            self.gold3_contract = front3.Symbol

#region imports
from AlgorithmImports import *
#endregion


general_setting = {
    "lookback": 100,
    "lookback_RESOLUTION": "HOUR",

    "ratio_method": "Regression",

    "Take_Profit_pct": 0.3,
    "Stop_Loss_pct": 0.08,

    "p_value_threshold_entry": 0.0001,
    "p_value_threshold_exit": 0.00001,
    "rollover_days": 2, 

}

from AlgorithmImports import *
from QuantConnect.DataSource import *

from config import general_setting
import pickle

import numpy as np
import pandas as pd
import math
import statsmodels.api as sm
from pandas.tseries.offsets import BDay
from pykalman import KalmanFilter
from statsmodels.tsa.stattools import coint, adfuller

class CalendarSpread(QCAlgorithm):

    def initialize(self) -> None:
        self.SetTimeZone(TimeZones.NEW_YORK)

        self.set_start_date(2024, 1, 1) 
        # self.set_end_date(2024,9,10)
        self.set_cash(1000000) 
        self.universe_settings.asynchronous = True

        self.zscore_df = {}
        self.note1_price = {}
        self.note2_price = {}


        # Requesting  data
        # Futures.Currencies.EUR
        # Futures.Currencies.MICRO_EUR
        # Futures.Financials.Y_2_TREASURY_NOTE
        # Futures.Financials.Y_5_TREASURY_NOTE
        # Futures.Indices.MICRO_NASDAQ_100_E_MINI
        # Futures.Indices.SP_500_E_MINI
        # Futures.Indices.VIX

        future_vix = self.add_future(Futures.Indices.VIX, resolution = Resolution.HOUR, extended_market_hours = True) 
        self.set_brokerage_model(BrokerageName.INTERACTIVE_BROKERS_BROKERAGE, AccountType.MARGIN)
        future_vix.set_filter(0, 180)
        self.future_vix_symbol = future_vix.symbol

        self.first_vix_contract = None
        self.second_vix_contract = None
        self.third_vix_contract = None

        self.first_vix_expiry = None
        self.second_vix_expiry = None
        self.third_vix_expiry = None

        self.lookback  = general_setting['lookback']
        self.p_threshold_entry = general_setting['p_value_threshold_entry']
        self.p_threshold_exit = general_setting['p_value_threshold_exit']
        self.rollover_days = general_setting['rollover_days']

        self.wt_1 = None
        self.wt_2 = None

        self.roll_signal = False
        self.Margin_Call = False

        self.prev_cap = None

        self.large_diff = None
        self.backwardation = False

        self.diversion = None

    def stats(self):
        # Request Historical Data
        df_vix1 = self.History(self.first_vix_contract.symbol, timedelta(self.lookback), Resolution.HOUR).rename(columns = {'close':'first'})
        df_vix2 = self.History(self.second_vix_contract.symbol, timedelta(self.lookback), Resolution.HOUR).rename(columns = {'close':'second'})
        # df_vix3 = self.History(self.third_vix_contract.symbol,timedelta(self.lookback), Resolution.HOUR).rename(columns = {'close':'third'})

        df_merge = pd.merge(df_vix1, df_vix2, on = ['time'], how = 'inner')

        vix1_log = np.array(df_merge['first'].apply(lambda x: math.log(x))) 
        vix2_log = np.array(df_merge['second'].apply(lambda x: math.log(x))) 
        # vix3_log = np.array(df_Gold3.apply(lambda x: math.log(x))) 

        # 1st & 2nd

    
        X1 = sm.add_constant(vix1_log)
        Y1 = vix2_log
        model1 = sm.OLS(Y1, X1)
        results1 = model1.fit()
        sigma1 = math.sqrt(results1.mse_resid)
        slope1 = results1.params[1]
        intercept1 = results1.params[0]
        res1 = results1.resid
        zscore1 = res1/sigma1

        adf1 = adfuller(res1)

        p_value1 = adf1[1]

        # spread = res1[len(res1)-1]
        df_merge['spread'] = df_merge['second'] - df_merge['first']
        
        spread = np.array(df_merge['spread'])
        # test_passed1 = p_value1 <= self.p_threshold
        # self.debug(f"p value is {p_value1}")

        return [p_value1, zscore1, slope1, spread]

    
    def on_data(self, slice: Slice) -> None:


        # Entry signal
        # if self.time.minute == 0 or self.time.minute ==10 or self.time.minute == 20 or self.time.minute==30 or self.time.minute == 40 or self.time.minute == 50:
        if self.roll_signal == False:
            if not self.portfolio.Invested:

                chain = slice.futures_chains.get(self.future_vix_symbol)
                if chain:
                    contracts = [i for i in chain ]
                    
                    e = [i.expiry for i in contracts]
                    e = sorted(list(set(sorted(e, reverse = True))))
                    # e = [i.expiry for i in contracts if i.expiry- self.Time> timedelta(5)]
                    # self.debug(f"the first contract is {e[0]}, the length of e is {len(e)}")
                    # expiry = e[0]
                    
                    try:
                        self.first_vix_contract = [contract for contract in contracts if contract.expiry == e[0]][0]
                        self.second_vix_contract = [contract for contract in contracts if contract.expiry == e[1]][0]
                        # self.third_gold_contract = [contract for contract in contracts if contract.expiry == e[2]][0]
                        self.first_vix_expiry = e[0]
                        self.second_vix_expiry = e[1]
                        # self.third_gold_expiry = e[2]

                        stats = self.stats()
                        self.zscore_df[self.time] = stats[1][-1]
                        self.note1_price[self.time] = self.Securities[self.first_vix_contract.symbol].Price
                        self.note2_price[self.time] = self.Securities[self.second_vix_contract.symbol].Price
                        sigma = stats[3].std()
                        mean = stats[3].mean()
                        last_spread = stats[3][-1]
                        self.debug(f'mean is {mean}, sigma is {sigma}, last_spread is {last_spread}')                        
                        # self.plot('z_score_plot','z_score',stats[1][-1] )
                        # self.plot('p_value_plot','p_value', stats[0])
                        # self.plot('p_value_plot','p_value', stats[0] )
                        # self.plot('spread_plot','spread', stats[3] )
                        # if (self.first_vix_expiry.date() - self.time.date()).days > self.rollover_day:
                        self.trade_signal = True
                        
                        # else:
                        #     self.trade_signal = False




                        if self.trade_signal and ((self.first_vix_expiry.date() - self.time.date()).days > self.rollover_days):
                            
                            self.wt_1 = 1/(1+stats[2])
                            self.wt_2 = 1 - self.wt_1
                                       


                            # if stats[3]<0:
                            if  last_spread > mean + 1.15*sigma and (last_spread < mean + 1.78*sigma):
                                n = (last_spread-mean)/sigma
                                self.set_holdings(self.first_vix_contract.symbol, -self.wt_1, tag = f'spread  = mean + {round(n,2)}*sigma (diversion)')
                                self.set_holdings(self.second_vix_contract.symbol,  self.wt_2, tag = f'spread = mean + {round(n,2)}*sigma  (diversion)')
                                self.prev_cap = self.portfolio.total_portfolio_value
                                self.large_diff = True

                            if  (last_spread > mean + 1.78*sigma):
                                n = (last_spread-mean)/sigma
                                self.set_holdings(self.first_vix_contract.symbol, self.wt_1, tag = f'spread  = mean + {round(n,2)}*sigma (mean reversion)')
                                self.set_holdings(self.second_vix_contract.symbol,  -self.wt_2, tag = f'spread = mean + {round(n,2)}*sigma  (mean reversion)')
                                self.prev_cap = self.portfolio.total_portfolio_value
                                self.large_diff = True
                                # self.debug(f"enter position: z score is {stats[1][-1]}")

                            elif last_spread < mean - 0.94*sigma and last_spread > mean - 1.73*sigma:
                                n = abs((last_spread-mean)/sigma)
                                self.set_holdings(self.first_vix_contract.symbol, self.wt_1, tag = f'spread < mean - {round(n,2)}*sigma (diversion)')
                                self.set_holdings(self.second_vix_contract.symbol, -self.wt_2, tag = f'spread < mean - {round(n,2)}*sigma (diversion)')
                                self.prev_cap = self.portfolio.total_portfolio_value
                                self.large_diff = False
                                # self.debug(f"enter position: z score is {stats[1][-1]}")
                                self.diversion = True

                            elif last_spread < mean - 1.73*sigma:
                                n = abs((last_spread-mean)/sigma)
                                self.set_holdings(self.first_vix_contract.symbol, -self.wt_1, tag = f'spread < mean - {round(n,2)}*sigma (mean reversion)')
                                self.set_holdings(self.second_vix_contract.symbol, self.wt_2, tag = f'spread < mean - {round(n,2)}*sigma (mean reversion)')
                                self.prev_cap = self.portfolio.total_portfolio_value
                                self.large_diff = False
                                # self.debug(f"enter position: z score is {stats[1][-1]}")

                            self.trade_signal = False

                    except:
                        return

            else:
                # exit signal
                stats = self.stats()
                self.zscore_df[self.time] = stats[1][-1]
                self.note1_price[self.time] = self.Securities[self.first_vix_contract.symbol].Price
                self.note2_price[self.time] = self.Securities[self.second_vix_contract.symbol].Price
                sigma = stats[3].std()
                mean = stats[3].mean()
                last_spread = stats[3][-1]
                # self.plot('p_value_plot','p_value', stats[0])
                # self.plot('z_score_plot','z_score',stats[1][-1] )
                # self.plot('spread_plot','spread', stats[3] )
                # self.debug(f'mean is {mean}, sigma is {sigma}, last_spread is {last_spread}')     

                # Roll over
                if ((self.first_vix_expiry.date() - self.time.date()).days <= self.rollover_days):
                    self.roll_signal = True
                    if self.portfolio.total_portfolio_value>= self.prev_cap:
                        self.liquidate(tag = 'rollover; Win')
                    else:
                        self.liquidate(tag = 'rollover; Loss')
                    self.prev_cap = None
                    self.large_diff = None

                    return
                    
                # Take Profit / Stop Loss    
                # if self.prev_cap :
                #     if self.portfolio.total_portfolio_value> 1.1 * self.prev_cap:
                #         self.liquidate(tag = 'Take Profit')
                #         self.prev_cap = None
                #         self.large_diff = None
                #         return

                #     elif self.portfolio.total_portfolio_value< 0.93 * self.prev_cap:
                #         self.liquidate(tag = 'Stop Loss')
                #         self.prev_cap = None
                #         self.large_diff = None
                #         return

                if self.diversion == True:

                    if (last_spread > mean + 1.78 * sigma and self.large_diff == True):
                        if self.portfolio.total_portfolio_value>= self.prev_cap:
                            self.liquidate(tag = 'Diversion; Win')
                        else:
                            self.liquidate(tag = 'Diversion; Loss')

                        stats = self.stats()
                        self.zscore_df[self.time] = stats[1][-1]
                        self.note1_price[self.time] = self.Securities[self.first_vix_contract.symbol].Price
                        self.note2_price[self.time] = self.Securities[self.second_vix_contract.symbol].Price
                        sigma = stats[3].std()
                        mean = stats[3].mean()
                        last_spread = stats[3][-1]
                        n = (last_spread-mean)/sigma
                        self.set_holdings(self.first_vix_contract.symbol, self.wt_1, tag = f'spread  = mean + {round(n,2)}*sigma (mean_revesion)')
                        self.set_holdings(self.second_vix_contract.symbol,  -self.wt_2, tag = f'spread = mean + {round(n,2)}*sigma (mean_reversion)')
                        self.prev_cap = self.portfolio.total_portfolio_value

                        self.large_diff = True
                        # self.debug(f"exit position: z score is {stats[1][-1]}")
                        self.diversion = False
                    elif (last_spread < mean - 1.73*sigma and self.large_diff == False):
                        if self.portfolio.total_portfolio_value>= self.prev_cap:
                            self.liquidate(tag = 'Diversion; Win')
                        else:
                            self.liquidate(tag = 'Diversion; Loss')

                        stats = self.stats()
                        self.zscore_df[self.time] = stats[1][-1]
                        self.note1_price[self.time] = self.Securities[self.first_vix_contract.symbol].Price
                        self.note2_price[self.time] = self.Securities[self.second_vix_contract.symbol].Price
                        sigma = stats[3].std()
                        mean = stats[3].mean()
                        last_spread = stats[3][-1]
                        n = (last_spread-mean)/sigma
                        self.set_holdings(self.first_vix_contract.symbol, -self.wt_1, tag = f'spread  = mean - {abs(round(n,2))}*sigma (mean_revesion)')
                        self.set_holdings(self.second_vix_contract.symbol,  self.wt_2, tag = f'spread = mean - {abs(round(n,2))}*sigma (mean_reversion)')
                        self.prev_cap = self.portfolio.total_portfolio_value
                        self.large_diff = False
                        self.diversion = False
                        # self.debug(f"exit position: z score is {stats[1][-1]}")

                    # elif : 
                    #     if self.portfolio.total_portfolio_value>= self.prev_cap:
                    #         self.liquidate(tag = 'Diversion; Win')
                    #     else:
                    #         self.liquidate(tag = 'Diversion; Loss')

                    #     stats = self.stats()
                    #     self.zscore_df[self.time] = stats[1][-1]
                    #     self.note1_price[self.time] = self.Securities[self.first_vix_contract.symbol].Price
                    #     self.note2_price[self.time] = self.Securities[self.second_vix_contract.symbol].Price
                    #     sigma = stats[3].std()
                    #     mean = stats[3].mean()
                    #     last_spread = stats[3][-1]
                    #     n = (last_spread-mean)/sigma
                    #     self.set_holdings(self.first_vix_contract.symbol, self.wt_1, tag = f'spread  = mean + {round(n,2)}*sigma (mean_revesion)')
                    #     self.set_holdings(self.second_vix_contract.symbol,  -self.wt_2, tag = f'spread = mean + {round(n,2)}*sigma (mean_reversion)')
                    #     self.prev_cap = self.portfolio.total_portfolio_value

                    #     self.large_diff = True
                    #     # self.debug(f"exit position: z score is {stats[1][-1]}")
                    #     self.diversion = False

                else:
                    if (last_spread < mean -0 * sigma and self.large_diff == True):
                        if self.portfolio.total_portfolio_value>= self.prev_cap:
                            self.liquidate(tag = 'Mean Reversion; Win')
                        else:
                            self.liquidate(tag = 'Mean Reversion; Loss')
                        
                        
                        self.diversion = None
                        self.prev_cap = None
                        self.large_diff = None


                        # self.debug(f"exit position: z score is {stats[1][-1]}")
                        
                    elif (last_spread > mean + 0*sigma and self.large_diff == False):
                        if self.portfolio.total_portfolio_value>= self.prev_cap:
                            self.liquidate(tag = 'Mean Reversion; Win')
                        else:
                            self.liquidate(tag = 'Mean Reversion; Loss')

                        self.prev_cap = None
                        self.large_diff = None
                        self.diversion = None




        else:

            stats = self.stats()
            self.zscore_df[self.time] = stats[1][-1]
            self.note1_price[self.time] = self.Securities[self.first_vix_contract.symbol].Price
            self.note2_price[self.time] = self.Securities[self.second_vix_contract.symbol].Price

            # self.plot('z_score_plot','z_score',stats[1][-1] )            
            # self.plot('p_value_plot','p_value', stats[0])

            if self.first_vix_expiry.date() < self.time.date():
                self.roll_signal = False


        # if self.zscore_df:
        #     df = pd.DataFrame.from_dict(self.zscore_df, orient='index',columns=['zscore'])
        #     file_name = 'CalendarSpread/zscore_df'
        #     self.object_store.SaveBytes(file_name, pickle.dumps(df))


        # if self.note1_price:
        #     df = pd.DataFrame.from_dict(self.note1_price, orient='index',columns=['price1'])
        #     file_name = 'CalendarSpread/note1_df'
        #     self.object_store.SaveBytes(file_name, pickle.dumps(df))

        # if self.note2_price:
        #     df = pd.DataFrame.from_dict(self.note2_price, orient='index',columns=['price2'])
        #     file_name = 'CalendarSpread/note2_df'
        #     self.object_store.SaveBytes(file_name, pickle.dumps(df))




    def OnOrderEvent(self, orderEvent):

        
        if orderEvent.Status != OrderStatus.Filled:
            return
        

        # Webhook Notification    
        symbol = orderEvent.symbol
        price = orderEvent.FillPrice
        quantity = orderEvent.quantity
        a = { "text": f"[Calendar Arbitrage Paper order update] \nSymbol: {symbol} \nPrice: {price} \nQuantity: {quantity}" }
        payload = json.dumps(a)
        self.notify.web("https://hooks.slack.com/services/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", payload)


    def on_margin_call(self, requests):
        self.debug('Margin Call is coming')
        self.Margin_Call =  True

        a = { "text": f"[Calendar Spread Margin Call update]Margin Call is coming" }
        payload = json.dumps(a)
        self.notify.web("https://hooks.slack.com/services/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", payload)

        return requests

    def OnOrderEvent(self, orderEvent):

        
        if orderEvent.Status != OrderStatus.Filled:
            return

        if self.Margin_Call:
            qty = orderEvent.quantity
            symbol = orderEvent.symbol
            
            self.Margin_Call = False
            self.debug(f'Hit margin call, the qty is {qty}')

            if symbol == self.first_vix_contract.symbol:
                self.debug(f'if come here, symbol is {symbol}, qty is {qty}')
                self.market_order(self.second_es_contract.symbol, -qty)

            if symbol == self.second_vix_contract.symbol:
                self.debug(f'if come here, symbol is {symbol}, qty is {qty}')
                self.market_order(self.first_es_contract.symbol, -qty)

            # self.liquidate(tag = 'margin call')