"""
Pairs Trading strategy (BTC-PERP / ETH-PERP).

Computes the Z-score of the BTC-ETH spread over a rolling window.
When the spread moves beyond a threshold, trades mean reversion.

- Z > +2: BTC expensive -> short BTC, long ETH
- Z < -2: BTC cheap -> long BTC, short ETH
"""
import numpy as np
from nautilus_trader.trading.strategy import Strategy
from nautilus_trader.config import StrategyConfig


class PairsTradingConfig(StrategyConfig, frozen=True):
    pair: tuple[str, str]
    z_entry: float = 2.0
    z_exit: float = 0.5
    lookback_hours: int = 24
    trade_size: float = 0.001
    hedge_ratio: float = 0.05


class PairsTrading(Strategy):
    """
    Statistical arbitrage on BTC/ETH spread.

    Assumes BTC and ETH are cointegrated — the spread between
    them tends to revert to a mean. Trades the deviations.
    """

    def __init__(self, config: PairsTradingConfig) -> None:
        super().__init__(config)
        self.config = config
        self.price_history: dict[str, list[float]] = {
            self.config.pair[0]: [],
            self.config.pair[1]: [],
        }
        self.position_open = False

    def on_start(self) -> None:
        for inst in self.config.pair:
            self.subscribe_bars(f"{inst}-1-MINUTE-LAST-INTERNAL")
        self.log.info(f"Pairs trading: {self.config.pair[0]} / {self.config.pair[1]}")

    def on_bar(self, bar) -> None:
        inst_id = str(bar.bar_type.instrument_id)
        if inst_id not in self.price_history:
            return

        self.price_history[inst_id].append(bar.close.as_double())

        a_hist = self.price_history[self.config.pair[0]]
        b_hist = self.price_history[self.config.pair[1]]
        if len(a_hist) < 100 or len(b_hist) < 100:
            return

        maxlen = self.config.lookback_hours * 60
        self.price_history[self.config.pair[0]] = a_hist[-maxlen:]
        self.price_history[self.config.pair[1]] = b_hist[-maxlen:]

        a = np.array(a_hist[-100:])
        b = np.array(b_hist[-100:])
        spread = a - self.config.hedge_ratio * b

        std = spread.std()
        z = (spread[-1] - spread.mean()) / std if std > 0 else 0

        self._signal(z)

    def _signal(self, z: float) -> None:
        btc_pos = self.portfolio.net_position(self.config.pair[0])

        if z > self.config.z_entry and btc_pos <= 0:
            self._trade("SELL", "BUY")
            self.position_open = True
        elif z < -self.config.z_entry and btc_pos >= 0:
            self._trade("BUY", "SELL")
            self.position_open = True
        elif abs(z) < self.config.z_exit and self.position_open:
            self.close_all_positions(self.config.pair[0])
            self.close_all_positions(self.config.pair[1])
            self.position_open = False

    def _trade(self, a_side: str, b_side: str) -> None:
        self.submit_order(self.order_factory.market(
            instrument_id=self.config.pair[0],
            order_side=a_side,
            quantity=self.config.trade_size,
        ))
        self.submit_order(self.order_factory.market(
            instrument_id=self.config.pair[1],
            order_side=b_side,
            quantity=self.config.trade_size / self.config.hedge_ratio,
        ))