#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, 4, 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 Gold data
        future_gold = self.add_future(Futures.Metals.GOLD, resolution = Resolution.HOUR) 
        future_gold.set_filter(0, 180)
        self.future_gold_symbol = future_gold.symbol

        self.first_gold_contract = None
        self.second_gold_contract = None
        self.third_gold_contract = None

        self.first_gold_expiry = None
        self.second_gold_expiry = None
        self.third_gold_expiry = None


        # # Requesting Crude Oil data
        future_CL = self.add_future(Futures.Energy.CRUDE_OIL_WTI, resolution = Resolution.HOUR) 
        future_CL.set_filter(0, 180)
        self.future_CL_symbol = future_CL.symbol

        self.first_CL_contract = None
        self.second_CL_contract = None
        self.third_CL_contract = None

        self.first_CL_expiry = None
        self.second_CL_expiry = None
        self.third_CL_expiry = None

        # # Requesting Y_10_TREASURY_NOTE data
        # future_BTC = self.add_future(Futures.Currencies.BTC, resolution = Resolution.HOUR) 
        # future_BTC.set_filter(0, 180)
        # self.future_BTC_symbol = future_BTC.symbol

        # self.first_BTC_contract = None
        # self.second_BTC_contract = None
        # self.third_BTC_contract = None

        # self.first_BTC_expiry = None
        # self.second_BTC_expiry = None
        # self.third_BTC_expiry = None

        # self.trade_signal = False


        # Requesting  data
        future_eur = self.add_future(Futures.Currencies.EUR, resolution = Resolution.HOUR) 
        future_eur.set_filter(0, 180)
        self.future_eur_symbol = future_eur.symbol

        self.first_eur_contract = None
        self.second_eur_contract = None
        self.third_eur_contract = None

        self.first_eur_expiry = None
        self.second_eur_expiry = None
        self.third_eur_expiry = None


        # 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_es = self.add_future(Futures.Indices.VIX, resolution = Resolution.HOUR, extended_market_hours = True) 
        self.set_brokerage_model(BrokerageName.INTERACTIVE_BROKERS_BROKERAGE, AccountType.MARGIN)
        future_es.set_filter(0, 180)
        self.future_es_symbol = future_es.symbol

        self.first_es_contract = None
        self.second_es_contract = None
        self.third_es_contract = None

        self.first_es_expiry = None
        self.second_es_expiry = None
        self.third_es_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

    def stats(self):
        # Request Historical Data
        df_es1 = self.History(self.first_es_contract.symbol, timedelta(self.lookback), Resolution.HOUR).rename(columns = {'close':'first'})
        df_es2 = self.History(self.second_es_contract.symbol, timedelta(self.lookback), Resolution.HOUR).rename(columns = {'close':'second'})
        # df_Gold3 = self.History(self.third_gold_contract.symbol,timedelta(self.lookback), Resolution.HOUR).rename(columns = {'close':'third'})

        df_merge = pd.merge(df_es1, df_es2, on = ['time'], how = 'inner')
        # df_Gold1 = df_Gold1["close"]
        # df_Gold2 = df_Gold2["close"]
        # df_Gold3 = df_Gold3["close"]
        # self.debug(f"{len(df_Gold1)}, {len(df_Gold2)}")

        es1_log = np.array(df_merge['first'].apply(lambda x: math.log(x))) 
        es2_log = np.array(df_merge['second'].apply(lambda x: math.log(x))) 
        # Gold3_log = np.array(df_Gold3.apply(lambda x: math.log(x))) 
        # self.debug(f"{len(Gold1_log)}, {len(Gold2_log)}")
        

        # 1st & 2nd

        # spread_series = df_merge['second'] - df_merge['first']
        # mean = spread_series.mean()
        # sigma = spread_series.std()
        # last_spread = spread_series[-1]
    
        X1 = sm.add_constant(es1_log)
        Y1 = es2_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:
        # self.debug(f"{self.time}: self.Rollover is {self.roll_signal}, first expiry is {self.first_gold_expiry}")

        # If backwardation




        # 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_es_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_es_contract = [contract for contract in contracts if contract.expiry == e[0]][0]
                        self.second_es_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_es_expiry = e[0]
                        self.second_es_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_es_contract.symbol].Price
                        self.note2_price[self.time] = self.Securities[self.second_es_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_es_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_es_expiry.date() - self.time.date()).days > self.rollover_days):
                            
                            self.wt_1 = 1/(1+stats[2])
                            self.wt_2 = 1 - self.wt_1
                            

                            # if self.Securities[self.first_es_contract.symbol].Price >= self.Securities[self.second_es_contract.symbol].Price:
                            #     self.backwardation == True
                                
                            #     self.set_holdings(self.first_es_contract.symbol, -self.wt_1, tag = f'spread  = mean + {round(n,2)}*sigma')
                            #     self.set_holdings(self.second_es_contract.symbol, -self.wt_2, tag = f'spread = mean + {round(n,2)}*sigma')                


                            # if stats[3]<0:
                            if  last_spread > mean + 0.9*sigma:
                                n = (last_spread-mean)/sigma
                                self.set_holdings(self.first_es_contract.symbol, self.wt_1, tag = f'spread  = mean + {round(n,2)}*sigma')
                                self.set_holdings(self.second_es_contract.symbol, -self.wt_2, tag = f'spread = mean + {round(n,2)}*sigma')
                                # self.set_holdings(self.first_es_contract.symbol, 04)
                                # self.set_holdings(self.second_es_contract.symbol,  -0.4)
                                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.85*sigma:
                                n = abs((last_spread-mean)/sigma)
                                self.set_holdings(self.first_es_contract.symbol, -self.wt_1, tag = f'spread < mean - {round(n,2)}*sigma')
                                self.set_holdings(self.second_es_contract.symbol, self.wt_2, tag = f'spread < mean - {round(n,2)}*sigma')
                                # self.set_holdings(self.first_es_contract.symbol, -0.4)
                                # self.set_holdings(self.second_es_contract.symbol,  0.4)
                                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_es_contract.symbol].Price
                self.note2_price[self.time] = self.Securities[self.second_es_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}')     


                if ((self.first_es_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
                    
                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 (last_spread < mean + 0 * sigma and self.large_diff == True)or (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.debug(f"exit position: z score is {stats[1][-1]}")
                    

                # roll over





        else:
            # chain = slice.futures_chains.get(self.future_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)]
            #     # expiry = e[0]
            #     self.first_gold_contract = [contract for contract in contracts if contract.expiry == e[0]][0]
            #     self.second_gold_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_gold_expiry = e[0]
            #     self.second_gold_expiry = e[1]

            stats = self.stats()
            self.zscore_df[self.time] = stats[1][-1]
            self.note1_price[self.time] = self.Securities[self.first_es_contract.symbol].Price
            self.note2_price[self.time] = self.Securities[self.second_es_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_es_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 on_securities_changed(self, changes: SecurityChanges) -> None:
    #     for security in changes.added_securities:
    #         # Historical data
    #         history = self.history(security.symbol, 10, Resolution.MINUTE)
    #         self.debug(f"We got {len(history)} from our history request for {security.symbol}")


    def OnOrderEvent(self, orderEvent):

        
        if orderEvent.Status != OrderStatus.Filled:
            return
        

        # Webhook Notification    
        symbol = orderEvent.symbol
        price = orderEvent.FillPrice
        quantity = orderEvent.quantity
        # self.debug(f"SP500 Enhanced-Indexing Paper order update] \nSymbol: {symbol} \nPrice: {price} \nQuantity: {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):

        # self.Log(f'{orderEvent.OrderId}--{orderEvent.Status}--{orderEvent.quantity}')
        
        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_es_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_es_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')

# 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