#region imports
from AlgorithmImports import *
#endregion

import numpy as np
from collections import deque


"""
PAIR TRADING STRATEGY: AAPL VS MSFT

This strategy trades a relative-value pair between Apple and Microsoft.

The model calculates the ratio:

    AAPL price / MSFT price

It then computes a rolling z-score of that ratio. The z-score measures how far
the current ratio is from its recent average.

Trading logic:

1. If the z-score is high:
   AAPL is expensive relative to MSFT.
   The strategy enters a short spread:
       short AAPL
       long MSFT

2. If the z-score is low:
   AAPL is cheap relative to MSFT.
   The strategy enters a long spread:
       long AAPL
       short MSFT

3. If the z-score mean-reverts toward zero:
   The strategy exits both legs.

This is a pair-trading strategy, not a long-only equity strategy. The goal is to
profit from relative movement between AAPL and MSFT, not from broad market beta.
Therefore, the most appropriate pedagogical benchmark is cash.

AAPL and MSFT buy-and-hold curves are also plotted as references, but they are not
the primary benchmark.
"""


class PairsTradingZScoreAlgorithm(QCAlgorithm):

    def Initialize(self):

        # ------------------------------------------------------------
        # 1. BACKTEST SETTINGS
        # ------------------------------------------------------------
        self.SetStartDate(2024, 1, 1)
        self.SetEndDate(2026, 5, 5)

        self.initial_cash = 1000000
        self.SetCash(self.initial_cash)

        # ------------------------------------------------------------
        # 2. SYMBOLS
        # ------------------------------------------------------------
        self.symbol1 = self.AddEquity("AAPL", Resolution.Daily).Symbol
        self.symbol2 = self.AddEquity("MSFT", Resolution.Daily).Symbol

        # Cash is the correct benchmark for a market-neutral pair trade.
        self.SetBenchmark(lambda time: self.initial_cash)

        # ------------------------------------------------------------
        # 3. PARAMETERS
        # ------------------------------------------------------------
        self.window = 20
        self.entry_threshold = 1.5
        self.exit_threshold = 0.5

        # Gross spread exposure.
        # 0.30 means approximately 15% long one leg and 15% short the other.
        self.position_size = 0.30

        # ------------------------------------------------------------
        # 4. DATA STORAGE
        # ------------------------------------------------------------
        self.prices1 = deque(maxlen=self.window)
        self.prices2 = deque(maxlen=self.window)

        # 0 = no position
        # 1 = long spread: long AAPL, short MSFT
        # -1 = short spread: short AAPL, long MSFT
        self.current_position = 0

        # ------------------------------------------------------------
        # 5. REFERENCE PRICES
        # ------------------------------------------------------------
        self.initial_aapl_price = None
        self.initial_msft_price = None

        # ------------------------------------------------------------
        # 6. TRADE CONTROL
        # ------------------------------------------------------------
        self.last_trade_date = None

        # ------------------------------------------------------------
        # 7. LOAD HISTORY
        # ------------------------------------------------------------
        self.LoadHistory()

        # ------------------------------------------------------------
        # 8. SCHEDULE DAILY TRADING DECISION
        # ------------------------------------------------------------
        self.Schedule.On(
            self.DateRules.EveryDay(self.symbol1),
            self.TimeRules.AfterMarketOpen(self.symbol1, 30),
            self.TradeLogic
        )

    def LoadHistory(self):

        history = self.History(
            [self.symbol1, self.symbol2],
            self.window,
            Resolution.Daily
        )

        if history.empty:
            self.Debug("No history data available.")
            return

        try:
            hist1 = history.loc[self.symbol1]["close"]
            hist2 = history.loc[self.symbol2]["close"]

            for p1, p2 in zip(hist1, hist2):
                self.prices1.append(float(p1))
                self.prices2.append(float(p2))

            self.Debug(
                "History loaded: "
                + str(len(self.prices1))
                + " AAPL prices and "
                + str(len(self.prices2))
                + " MSFT prices"
            )

        except:
            self.Debug("History loading error: missing symbol data.")

    def TradeLogic(self):

        # ------------------------------------------------------------
        # 1. PREVENT MULTIPLE DECISIONS SAME DAY
        # ------------------------------------------------------------
        if self.last_trade_date == self.Time.date():
            return

        # ------------------------------------------------------------
        # 2. CHECK DATA
        # ------------------------------------------------------------
        if not self.Securities[self.symbol1].HasData:
            return

        if not self.Securities[self.symbol2].HasData:
            return

        price1 = self.Securities[self.symbol1].Price
        price2 = self.Securities[self.symbol2].Price

        if price1 <= 0 or price2 <= 0:
            return

        # ------------------------------------------------------------
        # 3. UPDATE ROLLING WINDOW
        # ------------------------------------------------------------
        self.prices1.append(float(price1))
        self.prices2.append(float(price2))

        if len(self.prices1) < self.window:
            return

        if len(self.prices2) < self.window:
            return

        # ------------------------------------------------------------
        # 4. CALCULATE RATIO AND Z-SCORE
        # ------------------------------------------------------------
        ratio_series = np.array(self.prices1) / np.array(self.prices2)

        ratio_mean = np.mean(ratio_series)
        ratio_std = np.std(ratio_series)

        if ratio_std <= 0:
            return

        current_ratio = price1 / price2
        z_score = (current_ratio - ratio_mean) / ratio_std

        # ------------------------------------------------------------
        # 5. TRADING LOGIC
        # ------------------------------------------------------------
        if self.current_position == 0:

            if z_score > self.entry_threshold:
                self.EnterShortSpread()

            elif z_score < -self.entry_threshold:
                self.EnterLongSpread()

        elif self.current_position == 1:

            # Long spread exits when z-score mean-reverts upward.
            if z_score > -self.exit_threshold:
                self.ExitPosition()

        elif self.current_position == -1:

            # Short spread exits when z-score mean-reverts downward.
            if z_score < self.exit_threshold:
                self.ExitPosition()

        self.last_trade_date = self.Time.date()

        # ------------------------------------------------------------
        # 6. SIGNAL PLOTS
        # ------------------------------------------------------------
        self.Plot("Pair Signal", "Z Score", z_score)
        self.Plot("Pair Signal", "Entry Upper", self.entry_threshold)
        self.Plot("Pair Signal", "Entry Lower", -self.entry_threshold)
        self.Plot("Pair Signal", "Exit Upper", self.exit_threshold)
        self.Plot("Pair Signal", "Exit Lower", -self.exit_threshold)

    def EnterLongSpread(self):

        # Long spread:
        # Buy AAPL, short MSFT.
        leg_weight = self.position_size / 2.0

        self.SetHoldings(self.symbol1, leg_weight)
        self.SetHoldings(self.symbol2, -leg_weight)

        self.current_position = 1

        self.Debug(
            "Entered long spread: long AAPL, short MSFT on "
            + str(self.Time.date())
        )

    def EnterShortSpread(self):

        # Short spread:
        # Short AAPL, buy MSFT.
        leg_weight = self.position_size / 2.0

        self.SetHoldings(self.symbol1, -leg_weight)
        self.SetHoldings(self.symbol2, leg_weight)

        self.current_position = -1

        self.Debug(
            "Entered short spread: short AAPL, long MSFT on "
            + str(self.Time.date())
        )

    def ExitPosition(self):

        self.Liquidate(self.symbol1)
        self.Liquidate(self.symbol2)

        self.current_position = 0

        self.Debug(
            "Exited spread position on "
            + str(self.Time.date())
        )

    def OnData(self, data):

        # ------------------------------------------------------------
        # 1. CHECK DATA FOR PLOTS
        # ------------------------------------------------------------
        if self.symbol1 not in data or data[self.symbol1] is None:
            return

        if self.symbol2 not in data or data[self.symbol2] is None:
            return

        aapl_price = self.Securities[self.symbol1].Price
        msft_price = self.Securities[self.symbol2].Price

        if aapl_price <= 0 or msft_price <= 0:
            return

        if self.initial_aapl_price is None:
            self.initial_aapl_price = aapl_price

        if self.initial_msft_price is None:
            self.initial_msft_price = msft_price

        # ------------------------------------------------------------
        # 2. PRIMARY BENCHMARK: CASH
        # ------------------------------------------------------------
        cash_benchmark_value = self.initial_cash

        self.Plot(
            "Strategy Equity",
            "Portfolio Value",
            self.Portfolio.TotalPortfolioValue
        )

        self.Plot(
            "Strategy Equity",
            "Cash Benchmark",
            cash_benchmark_value
        )

        # ------------------------------------------------------------
        # 3. SECONDARY MARKET REFERENCES
        # ------------------------------------------------------------
        aapl_buy_hold = (
            self.initial_cash
            * aapl_price
            / self.initial_aapl_price
        )

        msft_buy_hold = (
            self.initial_cash
            * msft_price
            / self.initial_msft_price
        )

        self.Plot(
            "Market References",
            "Buy Hold AAPL",
            aapl_buy_hold
        )

        self.Plot(
            "Market References",
            "Buy Hold MSFT",
            msft_buy_hold
        )

        # ------------------------------------------------------------
        # 4. POSITION STATE
        # ------------------------------------------------------------
        self.Plot(
            "Position State",
            "Current Position",
            self.current_position
        )

        gross_exposure = (
            abs(self.Portfolio[self.symbol1].HoldingsValue)
            + abs(self.Portfolio[self.symbol2].HoldingsValue)
        )

        if self.Portfolio.TotalPortfolioValue > 0:

            gross_exposure_weight = (
                gross_exposure
                / self.Portfolio.TotalPortfolioValue
            )

            self.Plot(
                "Position State",
                "Gross Exposure",
                gross_exposure_weight
            )