# region imports
from AlgorithmImports import *
# endregion

class PreciousMetalsPairsTradingAlgorithm(QCAlgorithm):
    """Pairs trading strategy using gold (GLD) and silver (SLV) ratio mean reversion."""

    def initialize(self):
        self.set_start_date(2023, 1, 1)
        self.set_cash(100000)
        
        # Add precious metals ETFs
        self._gold = self.add_equity("GLD", Resolution.DAILY).symbol
        self._silver = self.add_equity("SLV", Resolution.DAILY).symbol
        
        # Strategy parameters
        self._lookback_period = 60  # Days for calculating mean and std dev
        self._entry_threshold = 2.0  # Z-score threshold for entry
        self._exit_threshold = 0.5   # Z-score threshold for exit
        
        # Rolling window to store gold/silver ratio
        self._ratio_window = RollingWindow[float](self._lookback_period)
        
        # Track position state
        self._position_side = 0  # 0 = flat, 1 = long gold/short silver, -1 = short gold/long silver
        
        # Warm up algorithm
        self.set_warm_up(self._lookback_period, Resolution.DAILY)
        
        # Schedule rebalancing daily after market open
        self.schedule.on(self.date_rules.every_day(self._gold),
                        self.time_rules.after_market_open(self._gold, 30),
                        self._rebalance)

    def _rebalance(self):
        """Check ratio and execute pairs trading logic."""
        # Get current prices
        gold_price = self.securities[self._gold].price
        silver_price = self.securities[self._silver].price
        
        if gold_price == 0 or silver_price == 0:
            return
        
        # Calculate current gold/silver ratio
        current_ratio = gold_price / silver_price
        
        # Add to rolling window
        self._ratio_window.add(current_ratio)
        
        # Need full window before trading
        if not self._ratio_window.is_ready:
            return
        
        # Calculate z-score
        ratio_values = [self._ratio_window[i] for i in range(self._ratio_window.count)]
        mean_ratio = np.mean(ratio_values)
        std_ratio = np.std(ratio_values)
        
        if std_ratio == 0:
            return
        
        z_score = (current_ratio - mean_ratio) / std_ratio
        
        self.plot("Ratio", "Gold/Silver Ratio", current_ratio)
        self.plot("Ratio", "Mean Ratio", mean_ratio)
        self.plot("Statistics", "Z-Score", z_score)
        
        # Entry signals
        if self._position_side == 0:
            if z_score > self._entry_threshold:
                # Ratio is high: Gold overvalued, Silver undervalued
                # Short Gold, Long Silver
                self._enter_trade(-1)
            elif z_score < -self._entry_threshold:
                # Ratio is low: Gold undervalued, Silver overvalued
                # Long Gold, Short Silver
                self._enter_trade(1)
        
        # Exit signals
        elif self._position_side != 0:
            if abs(z_score) < self._exit_threshold:
                # Ratio reverted to mean - exit positions
                self._exit_trade()
            elif (self._position_side == 1 and z_score > self._entry_threshold) or \
                 (self._position_side == -1 and z_score < -self._entry_threshold):
                # Stop loss: ratio moved further against us
                self._exit_trade()
    
    def _enter_trade(self, side: int):
        """Enter pairs trade. Side: 1 = long gold/short silver, -1 = short gold/long silver."""
        self._position_side = side
        
        if side == 1:
            # Long gold, short silver
            self.set_holdings(self._gold, 0.5)
            self.set_holdings(self._silver, -0.5)
            self.debug(f"Entered Long Gold / Short Silver at ratio {self.securities[self._gold].price / self.securities[self._silver].price:.2f}")
        else:
            # Short gold, long silver
            self.set_holdings(self._gold, -0.5)
            self.set_holdings(self._silver, 0.5)
            self.debug(f"Entered Short Gold / Long Silver at ratio {self.securities[self._gold].price / self.securities[self._silver].price:.2f}")
    
    def _exit_trade(self):
        """Exit pairs trade and flatten positions."""
        self.liquidate()
        self.debug(f"Exited pairs trade at ratio {self.securities[self._gold].price / self.securities[self._silver].price:.2f}")
        self._position_side = 0