# Copyright 2023, MetaQuotes Ltd.
# https://www.mql5.com
"""
 Script demonstrates use of coint() from statsmodels
 to test symbols for cointegration
"""
# imports 
from statsmodels.tsa.stattools import coint
from itertools import combinations
from datetime import datetime
import MetaTrader5 as mt5
import pandas as pd
import numpy as np
import pytz

#initialize connection to mt5
if not mt5.initialize():
    print("initialize() failed ")
    mt5.shutdown()
 
#set up timezone infomation   
tz=pytz.timezone("Etc/UTC")

#use time zone to set correct date for history data extraction
startdate = datetime(2022,1,1,hour=0,minute=0,second=1,tzinfo=tz)

#list the symbols 
Symbols = ["FB_us","GOOG_us","MSFT_us","NFLX_us","NVDA_us","AAPL_us","TSLA_us"]

#set length of data history
num_bars = 250

#set up the shape of the data structure to store prices 
data = np.zeros((num_bars,len(Symbols)))
prices = pd.DataFrame(data,columns=Symbols)

#fill prices dataframe with close prices
for symbol in Symbols:
    prices[symbol]=[rate[4]  for rate in mt5.copy_rates_from(symbol,mt5.TIMEFRAME_D1,startdate,num_bars)]
    
#we donot need mt5 from here 
mt5.shutdown()

#generate pairs from Symbols list 
pairs = list(combinations(prices.columns,2))

#set our desired significance level, 0.01->99%, 0.05->95%, 0.1->90%
confidence_level = 0.1

#do the test for cointegration on each pair and print results
for pair in pairs:
    df=prices[list(pair)]
    adf_stat,pvalue,critvalues=coint(df.values[:,0],df.values[:,1])
    if pvalue < confidence_level:
        print(pair[0]," and ",pair[1], " are likely cointegrated")
    else:
        print(pair[0]," and ",pair[1], " are likely not cointegrated")