from AlgorithmImports import *
import numpy as np

class ComplexCrossoverStrategy(QCAlgorithm):
    """
    This strategy trades SPY based on a multi-factor confirmation model using
    the full polar representation (magnitude and phase) of complex numbers.

    - A 200-day SMA is used as a long-term trend filter.
    - Magnitudes (r) and phase angles (phi) of z1 and z2 are smoothed.
    - Entry requires a bullish crossover in BOTH magnitude and phase.
    - Exits are triggered by an ATR trailing stop-loss or a bearish magnitude crossover.
    """

    def Initialize(self):
        """Initialise the data and resolution required for the algorithm."""
        self.SetStartDate(2020, 1, 1)
        self.SetEndDate(2024, 12, 31)
        self.SetCash(100000)

        self.symbol = self.AddEquity("SPY", Resolution.Daily).Symbol

        # --- Indicator Configuration ---
        self.trend_period = 200
        self.smoothing_period = 5
        self.atr_period = 14
        self.atr_multiplier = 2.5

        # --- Trend and Signal Indicators ---
        self.trend_filter = self.SMA(self.symbol, self.trend_period)

        # --- UPGRADE: Indicators for Magnitude (r) ---
        self.smoothed_r1 = SimpleMovingAverage(self.smoothing_period)
        # FIX: Corrected the typo from SimpleMovingAverag to SimpleMovingAverage
        self.smoothed_r2 = SimpleMovingAverage(self.smoothing_period)

        # --- UPGRADE: Indicators for Phase Angle (phi) ---
        self.smoothed_phi1 = SimpleMovingAverage(self.smoothing_period)
        self.smoothed_phi2 = SimpleMovingAverage(self.smoothing_period)

        # --- Risk Management Indicators ---
        self.atr = self.ATR(self.symbol, self.atr_period)
        self.trailing_stop_price = 0

        # --- State variables for crossover detection ---
        self.prev_smoothed_r1 = None
        self.prev_smoothed_r2 = None

        # Set warm-up period for the longest indicator
        self.SetWarmUp(self.trend_filter.Period + self.smoothing_period)

    def OnData(self, data: Slice):
        """OnData event is the primary entry point for your algorithm."""
        if self.IsWarmingUp or not data.ContainsKey(self.symbol) or not self.atr.IsReady:
            return

        bar = data[self.symbol]
        if bar is None:
            return

        # --- Signal Calculation ---
        # 1. Construct complex numbers
        if bar.Close >= bar.Open:
            z1, z2 = (bar.Open + 1j * bar.Low, bar.Close + 1j * bar.High)
        else:
            z1, z2 = (bar.Open + 1j * bar.High, bar.Close + 1j * bar.Low)

        # 2. Update Magnitude (r) and Phase (phi) indicators
        self.smoothed_r1.Update(bar.EndTime, abs(z1))
        self.smoothed_r2.Update(bar.EndTime, abs(z2))
        self.smoothed_phi1.Update(bar.EndTime, np.angle(z1))
        self.smoothed_phi2.Update(bar.EndTime, np.angle(z2))

        # 3. Wait for all indicators to be ready
        if not all([self.trend_filter.IsReady, self.smoothed_r1.IsReady, self.smoothed_phi1.IsReady]):
            return

        # Get current smoothed values
        r1_smooth = self.smoothed_r1.Current.Value
        r2_smooth = self.smoothed_r2.Current.Value
        phi1_smooth = self.smoothed_phi1.Current.Value
        phi2_smooth = self.smoothed_phi2.Current.Value

        if self.prev_smoothed_r1 is None:
            self.prev_smoothed_r1 = r1_smooth
            self.prev_smoothed_r2 = r2_smooth
            return

        # --- Trading Logic ---
        is_invested = self.Portfolio[self.symbol].Invested
        price = self.Securities[self.symbol].Price

        # --- Risk Management: Trailing Stop-Loss ---
        if is_invested:
            new_stop_price = price - self.atr.Current.Value * self.atr_multiplier
            self.trailing_stop_price = max(self.trailing_stop_price, new_stop_price)
            if price <= self.trailing_stop_price:
                self.Debug(f"{self.Time} - ATR Trailing Stop Hit at {price:.2f}. Liquidating.")
                self.Liquidate(self.symbol)
                self.trailing_stop_price = 0
                return

        # --- Entry and Exit Signals ---
        is_uptrend = price > self.trend_filter.Current.Value

        # --- UPGRADE: Dual Confirmation Bullish Entry ---
        is_magnitude_bullish = self.prev_smoothed_r2 <= self.prev_smoothed_r1 and r2_smooth > r1_smooth
        is_phase_bullish = phi2_smooth > phi1_smooth

        if is_magnitude_bullish and is_phase_bullish:
            if is_uptrend and not is_invested:
                self.Debug(f"{self.Time} - Dual Confirmation: Bullish Crossover in Uptrend. Buying SPY.")
                self.SetHoldings(self.symbol, 1.0)
                self.trailing_stop_price = price - self.atr.Current.Value * self.atr_multiplier
                self.Debug(f"{self.Time} - Initial ATR Stop set at {self.trailing_stop_price:.2f}")

        # --- Bearish Magnitude Crossover Exit ---
        is_magnitude_bearish = self.prev_smoothed_r1 <= self.prev_smoothed_r2 and r1_smooth > r2_smooth
        if is_magnitude_bearish and is_invested:
            self.Debug(f"{self.Time} - Bearish Magnitude Crossover. Liquidating SPY.")
            self.Liquidate(self.symbol)
            self.trailing_stop_price = 0

        # Update previous values for the next iteration
        self.prev_smoothed_r1 = r1_smooth
        self.prev_smoothed_r2 = r2_smooth

    def OnEndOfAlgorithm(self):
        """Called at the end of the algorithm execution."""
        self.Log(f"Ending backtest. Final portfolio value: {self.Portfolio.TotalPortfolioValue}")