#region imports
from AlgorithmImports import *
#endregion

import numpy as np
from collections import deque


"""
PAIR TRADING STRATEGY: EEM VS EFA

This strategy trades a relative-value pair between EEM and EFA.

EEM represents emerging-market equities.
EFA represents developed international equities.

The model calculates the ratio:

    EEM price / EFA 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:
   EEM is expensive relative to EFA.
   The strategy enters a short spread:
       short EEM
       long EFA

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

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. Therefore, the
primary benchmark is cash. EEM and EFA buy-and-hold curves are plotted only as
market references.
"""


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("EEM", Resolution.Daily).Symbol
        self.symbol2 = self.AddEquity("EFA", Resolution.Daily).Symbol

        # Cash benchmark is appropriate for a relative-value 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 about 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 EEM, short EFA
        # -1 = short spread: short EEM, long EFA
        self.current_position = 0

        # ------------------------------------------------------------
        # 5. REFERENCE PRICES
        # ------------------------------------------------------------
        self.initial_eem_price = None
        self.initial_efa_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))
                + " EEM prices and "
                + str(len(self.prices2))
                + " EFA prices"
            )

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

    def TradeLogic(self):

        # ------------------------------------------------------------
        # 1. PREVENT MULTIPLE DECISIONS ON THE 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 EEM, short EFA.
        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 EEM, short EFA on "
            + str(self.Time.date())
        )

    def EnterShortSpread(self):

        # Short spread:
        # Short EEM, buy EFA.
        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 EEM, long EFA 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

        eem_price = self.Securities[self.symbol1].Price
        efa_price = self.Securities[self.symbol2].Price

        if eem_price <= 0 or efa_price <= 0:
            return

        if self.initial_eem_price is None:
            self.initial_eem_price = eem_price

        if self.initial_efa_price is None:
            self.initial_efa_price = efa_price

        # ------------------------------------------------------------
        # 2. PRIMARY BENCHMARK: CASH
        # ------------------------------------------------------------
        self.Plot(
            "Strategy Equity",
            "Portfolio Value",
            self.Portfolio.TotalPortfolioValue
        )

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

        # ------------------------------------------------------------
        # 3. SECONDARY MARKET REFERENCES
        # ------------------------------------------------------------
        eem_buy_hold = (
            self.initial_cash
            * eem_price
            / self.initial_eem_price
        )

        efa_buy_hold = (
            self.initial_cash
            * efa_price
            / self.initial_efa_price
        )

        self.Plot(
            "Market References",
            "Buy Hold EEM",
            eem_buy_hold
        )

        self.Plot(
            "Market References",
            "Buy Hold EFA",
            efa_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
            )