#region imports
from AlgorithmImports import *
#endregion

import numpy as np
from collections import deque


"""
CORN / LIVE CATTLE LONG-SHORT COMMODITY PAIR TRADE
WITH MODERATE RISK MANAGEMENT AND TRUE RESTART AFTER COOL-OFF

This strategy trades a relative-value pair between continuous futures:

1. Corn futures
2. Live Cattle futures

The model calculates:

    Corn price / Live Cattle price

It then computes a rolling z-score of that ratio.

Trading logic:

- If z-score is high:
      Corn is expensive relative to Live Cattle
      short Corn, long Live Cattle

- If z-score is low:
      Corn is cheap relative to Live Cattle
      long Corn, short Live Cattle

- If z-score mean-reverts:
      exit both legs

Risk management:

1. Spread stop loss:
   Exit if the combined spread loses more than 12.5% of initial gross exposure.

2. Single-leg stop loss:
   Exit if the worst leg loses more than 10% of initial gross spread exposure.

3. Adverse z-score stop:
   Exit if the z-score moves materially further against the open trade.

4. Maximum holding period:
   Exit after 45 calendar days.

5. Cool-off:
   After a stop loss, wait 10 calendar days, then restart normally.

Important design change:

The model resets its risk state after the cool-off period. This prevents the
algorithm from remaining stuck in a permanent defensive state.

Benchmark:

Cash is the benchmark because this is a long/short relative-value strategy.
"""


class CommodityPairsTradingAlgorithm(QCAlgorithm):

    def Initialize(self):

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

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

        # ------------------------------------------------------------
        # 2. ADD CONTINUOUS FUTURES
        # ------------------------------------------------------------
        self.corn = self.AddFuture(
            Futures.Grains.Corn,
            Resolution.Daily,
            dataNormalizationMode=DataNormalizationMode.BackwardsRatio,
            dataMappingMode=DataMappingMode.OpenInterest,
            contractDepthOffset=0,
            extendedMarketHours=True
        )

        self.cattle = self.AddFuture(
            Futures.Meats.LiveCattle,
            Resolution.Daily,
            dataNormalizationMode=DataNormalizationMode.BackwardsRatio,
            dataMappingMode=DataMappingMode.OpenInterest,
            contractDepthOffset=0,
            extendedMarketHours=True
        )

        self.corn_symbol = self.corn.Symbol
        self.cattle_symbol = self.cattle.Symbol

        # ------------------------------------------------------------
        # 3. SIGNAL PARAMETERS
        # ------------------------------------------------------------
        self.window = 45
        self.entry_threshold = 2.0
        self.exit_threshold = 0.25

        # ------------------------------------------------------------
        # 4. MODERATE RISK PARAMETERS
        # ------------------------------------------------------------
        self.leg_weight = 0.05

        self.spread_stop_loss_pct = 0.125
        self.single_leg_stop_loss_pct = 0.10
        self.adverse_z_stop = 4.0

        self.max_holding_days = 45

        self.cooloff_days = 10
        self.cooloff_until = datetime.min.date()
        self.cooloff_was_active = False

        # ------------------------------------------------------------
        # 5. ROLLING DATA
        # ------------------------------------------------------------
        self.corn_prices = deque(maxlen=self.window)
        self.cattle_prices = deque(maxlen=self.window)

        # 0 = no position
        # 1 = long Corn / short Cattle
        # -1 = short Corn / long Cattle
        self.current_position = 0

        self.last_trade_date = None

        self.entry_gross_exposure = None
        self.entry_portfolio_value = None
        self.entry_date = None
        self.entry_z_score = None

        # Cash benchmark.
        self.SetBenchmark(lambda time: self.initial_cash)

        # Build rolling window live. This avoids continuous futures
        # history-indexing errors during Initialize.
        self.Schedule.On(
            self.DateRules.EveryDay(),
            self.TimeRules.At(10, 0),
            self.TradeLogic
        )

    def TradeLogic(self):

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

        # ------------------------------------------------------------
        # 2. HANDLE COOL-OFF RESTART
        # ------------------------------------------------------------
        if self.Time.date() <= self.cooloff_until:

            self.cooloff_was_active = True

            self.Debug(
                "Cool-off active on "
                + str(self.Time.date())
                + ". No new trades until "
                + str(self.cooloff_until)
            )

            self.last_trade_date = self.Time.date()
            return

        # Once cool-off has expired, explicitly reset risk state.
        if self.cooloff_was_active:

            self.ResetTradeRiskState()

            self.cooloff_was_active = False

            self.Debug(
                "Cool-off ended on "
                + str(self.Time.date())
                + ". Strategy restarted."
            )

        # ------------------------------------------------------------
        # 3. GET CONTINUOUS FUTURES PRICES
        # ------------------------------------------------------------
        if not self.Securities.ContainsKey(self.corn_symbol):
            return

        if not self.Securities.ContainsKey(self.cattle_symbol):
            return

        corn_price = self.Securities[self.corn_symbol].Price
        cattle_price = self.Securities[self.cattle_symbol].Price

        if corn_price <= 0 or cattle_price <= 0:
            return

        self.corn_prices.append(float(corn_price))
        self.cattle_prices.append(float(cattle_price))

        # Wait until rolling window is full.
        if len(self.corn_prices) < self.window:
            self.Debug(
                "Building rolling window: "
                + str(len(self.corn_prices))
                + " / "
                + str(self.window)
            )
            return

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

        z_score, current_ratio = self.CalculateZScore(
            corn_price,
            cattle_price
        )

        if z_score is None:
            return

        # ------------------------------------------------------------
        # 4. RISK CHECKS ONLY WHEN POSITION IS OPEN
        # ------------------------------------------------------------
        if self.current_position != 0:

            risk_exit_reason = self.CheckRiskExits(z_score)

            if risk_exit_reason is not None:

                self.ExitSpread(risk_exit_reason)

                self.cooloff_until = (
                    self.Time.date()
                    + timedelta(days=self.cooloff_days)
                )

                self.cooloff_was_active = True

                self.Debug(
                    "Risk exit: "
                    + risk_exit_reason
                    + " on "
                    + str(self.Time.date())
                    + ". Cooling off until "
                    + str(self.cooloff_until)
                )

                self.last_trade_date = self.Time.date()
                return

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

            if z_score > self.entry_threshold:

                self.EnterShortSpread(z_score)

            elif z_score < -self.entry_threshold:

                self.EnterLongSpread(z_score)

        elif self.current_position == 1:

            # Long spread exits when ratio mean-reverts upward.
            if z_score > -self.exit_threshold:

                self.ExitSpread("Mean reversion exit")

        elif self.current_position == -1:

            # Short spread exits when ratio mean-reverts downward.
            if z_score < self.exit_threshold:

                self.ExitSpread("Mean reversion exit")

        self.last_trade_date = self.Time.date()

        # ------------------------------------------------------------
        # 6. 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)
        self.Plot("Pair Signal", "Corn / Cattle Ratio", current_ratio)

    def CalculateZScore(self, corn_price, cattle_price):

        ratio_series = (
            np.array(self.corn_prices)
            / np.array(self.cattle_prices)
        )

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

        if ratio_std <= 0:
            return None, None

        current_ratio = corn_price / cattle_price
        z_score = (current_ratio - ratio_mean) / ratio_std

        return z_score, current_ratio

    def EnterLongSpread(self, z_score):

        # Long spread:
        # Long Corn, short Live Cattle.
        corn_contract = self.corn.Mapped
        cattle_contract = self.cattle.Mapped

        if corn_contract is None or cattle_contract is None:
            return

        self.SetHoldings(corn_contract, self.leg_weight)
        self.SetHoldings(cattle_contract, -self.leg_weight)

        self.current_position = 1

        self.entry_portfolio_value = self.Portfolio.TotalPortfolioValue
        self.entry_gross_exposure = None
        self.entry_date = self.Time.date()
        self.entry_z_score = z_score

        self.Debug(
            "Entered long spread: long Corn, short Live Cattle on "
            + str(self.Time.date())
            + " z="
            + str(round(z_score, 2))
        )

    def EnterShortSpread(self, z_score):

        # Short spread:
        # Short Corn, long Live Cattle.
        corn_contract = self.corn.Mapped
        cattle_contract = self.cattle.Mapped

        if corn_contract is None or cattle_contract is None:
            return

        self.SetHoldings(corn_contract, -self.leg_weight)
        self.SetHoldings(cattle_contract, self.leg_weight)

        self.current_position = -1

        self.entry_portfolio_value = self.Portfolio.TotalPortfolioValue
        self.entry_gross_exposure = None
        self.entry_date = self.Time.date()
        self.entry_z_score = z_score

        self.Debug(
            "Entered short spread: short Corn, long Live Cattle on "
            + str(self.Time.date())
            + " z="
            + str(round(z_score, 2))
        )

    def ExitSpread(self, reason):

        # Liquidate all open futures positions.
        # This is safer than liquidating only the currently mapped contracts
        # because futures contracts can roll.
        for holding in self.Portfolio.Values:

            if (
                holding.Invested
                and holding.Symbol.SecurityType == SecurityType.Future
            ):
                self.Liquidate(holding.Symbol)

        self.current_position = 0
        self.ResetTradeRiskState()

        self.Debug(
            "Exited spread on "
            + str(self.Time.date())
            + " | Reason: "
            + reason
        )

    def CheckRiskExits(self, z_score):

        total_unrealized_pnl = 0.0
        gross_exposure = 0.0
        worst_leg_pnl = 0.0

        for holding in self.Portfolio.Values:

            if (
                holding.Invested
                and holding.Symbol.SecurityType == SecurityType.Future
            ):
                total_unrealized_pnl += holding.UnrealizedProfit
                gross_exposure += abs(holding.HoldingsValue)

                if holding.UnrealizedProfit < worst_leg_pnl:
                    worst_leg_pnl = holding.UnrealizedProfit

        if gross_exposure <= 0:
            return None

        if self.entry_gross_exposure is None:
            self.entry_gross_exposure = gross_exposure

        if self.entry_gross_exposure <= 0:
            return None

        spread_return = total_unrealized_pnl / self.entry_gross_exposure
        worst_leg_return = worst_leg_pnl / self.entry_gross_exposure

        self.Plot("Risk Management", "Spread Return", spread_return)
        self.Plot("Risk Management", "Worst Leg Return", worst_leg_return)
        self.Plot("Risk Management", "Spread Stop", -self.spread_stop_loss_pct)
        self.Plot("Risk Management", "Leg Stop", -self.single_leg_stop_loss_pct)

        # ------------------------------------------------------------
        # 1. SPREAD STOP LOSS
        # ------------------------------------------------------------
        if spread_return <= -self.spread_stop_loss_pct:
            return "Spread stop loss"

        # ------------------------------------------------------------
        # 2. SINGLE LEG STOP LOSS
        # ------------------------------------------------------------
        if worst_leg_return <= -self.single_leg_stop_loss_pct:
            return "Single-leg stop loss"

        # ------------------------------------------------------------
        # 3. ADVERSE Z-SCORE STOP
        # ------------------------------------------------------------
        if self.current_position == 1:

            if z_score <= -self.adverse_z_stop:
                return "Adverse z-score stop"

        if self.current_position == -1:

            if z_score >= self.adverse_z_stop:
                return "Adverse z-score stop"

        # ------------------------------------------------------------
        # 4. MAX HOLDING PERIOD
        # ------------------------------------------------------------
        if self.entry_date is not None:

            days_held = (
                self.Time.date()
                - self.entry_date
            ).days

            self.Plot("Risk Management", "Days Held", days_held)

            if days_held >= self.max_holding_days:
                return "Max holding period"

        return None

    def ResetTradeRiskState(self):

        self.entry_gross_exposure = None
        self.entry_portfolio_value = None
        self.entry_date = None
        self.entry_z_score = None

    def OnData(self, data):

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

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

        corn_price = self.Securities[self.corn_symbol].Price
        cattle_price = self.Securities[self.cattle_symbol].Price

        if corn_price > 0:
            self.Plot(
                "Continuous Futures Prices",
                "Corn",
                corn_price
            )

        if cattle_price > 0:
            self.Plot(
                "Continuous Futures Prices",
                "Live Cattle",
                cattle_price
            )

        self.Plot(
            "Position State",
            "Current Position",
            self.current_position
        )

        gross_exposure = 0.0

        for holding in self.Portfolio.Values:

            if holding.Invested:
                gross_exposure += abs(holding.HoldingsValue)

        if self.Portfolio.TotalPortfolioValue > 0:

            self.Plot(
                "Position State",
                "Gross Exposure",
                gross_exposure / self.Portfolio.TotalPortfolioValue
            )

        cooloff_active = 0

        if self.Time.date() <= self.cooloff_until:
            cooloff_active = 1

        self.Plot(
            "Risk Management",
            "Cool-Off Active",
            cooloff_active
        )