import numpy as np
import pandas as pd
import statsmodels.api as sm
from AlgorithmImports import *

class PairsTradingStrategy(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2023, 12, 1)
        self.SetEndDate(2024, 12, 1)
        self.SetCash(1_000_000)  # $1M initial capital
        
        self.lookback = 20  # Lookback period for z-score calculation
        self.entry_threshold = 2
        self.exit_threshold = 0
        self.risk_threshold = 3
        self.pair = ["NVDA", "MSFT"]  # Replace with the selected pair with lowest ADF
        
        # Add securities
        self.symbols = [self.AddEquity(ticker, Resolution.Daily).Symbol for ticker in self.pair]
        
        # Store historical data
        self.history = {sym: [] for sym in self.symbols}
        self.portfolio_state = 0  # 1: long, -1: short, 0: no position
        
        # Track current hedge ratio
        self.hedge_ratio = None
    
    def OnData(self, data):
        if not all(symbol in data for symbol in self.symbols):
            return
        
        # Store latest closing prices
        for symbol in self.symbols:
            self.history[symbol].append(data[symbol].Close)
            if len(self.history[symbol]) > self.lookback:
                self.history[symbol].pop(0)  # Maintain rolling window
        
        if len(self.history[self.symbols[0]]) < self.lookback:
            return  # Wait until we have enough data
        
        # Perform linear regression on log prices
        log_prices_P = np.log(self.history[self.symbols[0]])
        log_prices_Q = np.log(self.history[self.symbols[1]])
        
        model = sm.OLS(log_prices_P, sm.add_constant(log_prices_Q)).fit()
        self.hedge_ratio = model.params[1]  # Beta coefficient
        
        # Compute spread and z-score
        spread = log_prices_P - self.hedge_ratio * log_prices_Q
        mean_spread = np.mean(spread)
        std_spread = np.std(spread)
        zscore = (spread[-1] - mean_spread) / std_spread if std_spread != 0 else 0
        
        self.wt1 = 1/(1+self.hedge_ratio)
        self.wt2 = self.hedge_ratio/(1+self.hedge_ratio)


        # **Risk Management: Liquidate all positions if |zscore| > 3**
        if abs(zscore) > self.risk_threshold:
            self.Liquidate()
            self.portfolio_state = 0
            return


        # Entry conditions
        if self.portfolio_state == 0:
            if zscore < -self.entry_threshold:
                # Long P, short Q
                self.SetHoldings(self.symbols[0], self.wt1)
                self.SetHoldings(self.symbols[1], -self.wt2)
                self.portfolio_state = 1
            
            elif zscore > self.entry_threshold:
                # Short P, long Q
                self.SetHoldings(self.symbols[0], -self.wt1)
                self.SetHoldings(self.symbols[1], self.wt2)
                self.portfolio_state = -1
        
        # Exit conditions
        elif self.portfolio_state == 1 and zscore > self.exit_threshold:
            self.Liquidate()  # Close all positions
            self.portfolio_state = 0
        elif self.portfolio_state == -1 and zscore < self.exit_threshold:
            self.Liquidate()  # Close all positions
            self.portfolio_state = 0