from AlgorithmImports import *
from itertools import combinations, permutations
from statsmodels.tsa.stattools import adfuller

class PairsTradingADF(QCAlgorithm):
    def Initialize(self):
        self.SetStartDate(2023, 12, 1)
        self.SetEndDate(2024, 12, 1)
        self.SetCash(100000)
        
        # 3 tech stocks (top by market cap)
        self.symbols = ["AAPL", "MSFT", "NVDA"]
        for symbol in self.symbols:
            self.AddEquity(symbol, Resolution.Daily)
        
        # Generate 6 UNIQUE PAIRS (using combinations)
        self.stock_pairs = list(combinations(self.symbols, 2))  
        self.stock_pairs = list(permutations(self.symbols, 2))  

    def OnEndOfAlgorithm(self):
        # Get 1 year of CLOSE prices
        history = self.History(self.symbols, 252, Resolution.Daily)
        if history.empty:
            self.Debug("No historical data.")
            return
        
        # Compute daily returns
        closes = history['close'].unstack(level=0)
        returns = closes.pct_change().dropna()
        
        # Calculate ADF for each pair
        self.pair_adfs = {}
        for pair in self.stock_pairs:
            stock1, stock2 = pair
            if stock1 in returns.columns and stock2 in returns.columns:
                spread = returns[stock1] - returns[stock2]  # return spread
                adf_result = adfuller(spread.dropna())
                
                self.pair_adfs[pair] = {
                    "ADF Stat": adf_result[0],
                    "P-Value": adf_result[1],
                    "5% Critical": adf_result[4]['5%']  # add critical value
                }
        
        # Print all results line-by-line
        self.Debug("ADF Results (All Pairs):")
        for pair, stats in self.pair_adfs.items():
            self.Debug(f"{pair}: ADF={stats['ADF Stat']:.4f} | "
                      f"P-Val={stats['P-Value']:.4f} | Critical_5%={stats['5% Critical']:.4f}")
        
        # Extreme pairs
        if self.pair_adfs:
            min_pair = min(self.pair_adfs, key=lambda k: self.pair_adfs[k]["ADF Stat"])
            max_pair = max(self.pair_adfs, key=lambda k: self.pair_adfs[k]["ADF Stat"])
            self.Debug(f"\nLOWEST ADF: {min_pair} ({self.pair_adfs[min_pair]['ADF Stat']:.4f})")
            self.Debug(f"HIGHEST ADF: {max_pair} ({self.pair_adfs[max_pair]['ADF Stat']:.4f})")