"""
Daily swing mean-reversion engine based on Connors RSI (CRSI) pullbacks.

This implementation uses a standard RSI calculation via:
  * price.diff()
  * positive/negative clipping
  * Wilder-style smoothing via pandas .ewm(...)

No external TA libraries required.
"""

from __future__ import annotations

from dataclasses import dataclass

import numpy as np
import pandas as pd


def rsi_wilder(close: pd.Series, window: int) -> pd.Series:
    """
    Compute the standard RSI using Wilder-style exponentially smoothed averages.

    RSI = 100 - 100 / (1 + RS)
    RS = AvgGain / AvgLoss

    Parameters
    ----------
    close
        Price series (DatetimeIndex recommended).
    window
        Lookback period for smoothing (e.g. 2-period RSI).
    """
    if window <= 0:
        raise ValueError("window must be > 0")
    c = close.astype(np.float64)
    delta = c.diff()
    gain = delta.clip(lower=0.0)
    loss = (-delta).clip(lower=0.0)

    alpha = 1.0 / float(window)
    avg_gain = gain.ewm(alpha=alpha, adjust=False, min_periods=window).mean()
    avg_loss = loss.ewm(alpha=alpha, adjust=False, min_periods=window).mean()

    rs = avg_gain / avg_loss.replace(0.0, np.nan)
    rsi = 100.0 - (100.0 / (1.0 + rs))
    return rsi


@dataclass
class SwingEngine:
    """
    Connors RSI-like pullback strategy on daily bars.

    Entry (vectorized state machine):
      * close > SMA(trend_sma)
      * RSI(rsi_window) < rsi_threshold

    Exit:
      * close > SMA(exit_sma)

    The engine outputs:
      * `micro_position`: {-1,0,1} positions (0/1 for this long-only CRSI pullback)
    """

    trend_sma: int = 200
    rsi_window: int = 2
    rsi_threshold: float = 10.0
    exit_sma: int = 5

    def transform(self, daily: pd.DataFrame) -> pd.DataFrame:
        """
        Transform daily OHLCV into signal + position columns.

        Expected columns
        -----------------
        daily
            Must contain `close` column.
        """
        if daily.empty:
            raise ValueError("daily is empty")
        if "close" not in daily.columns:
            raise KeyError("daily must contain a 'close' column")

        df = daily.sort_index().copy()
        close = df["close"].astype(np.float64)

        sma_trend = df["close"].rolling(window=self.trend_sma, min_periods=self.trend_sma).mean()
        sma_exit = df["close"].rolling(window=self.exit_sma, min_periods=self.exit_sma).mean()
        rsi = rsi_wilder(close, self.rsi_window)

        df["sma_200"] = sma_trend
        df["sma_5"] = sma_exit
        df["rsi_2"] = rsi

        # Vectorized state machine.
        entry_signal = (close > df["sma_200"]) & (df["rsi_2"] < float(self.rsi_threshold))
        exit_signal = close > df["sma_5"]

        df["entry_signal"] = entry_signal
        df["exit_signal"] = exit_signal

        # target_position is a state variable: 1 on entry, 0 on exit, otherwise carry forward.
        target_position = np.where(entry_signal.to_numpy(), 1.0, np.where(exit_signal.to_numpy(), 0.0, np.nan))
        target_position = pd.Series(target_position, index=df.index, dtype=np.float64)
        target_position = target_position.ffill().fillna(0.0)

        # Shift by 1 bar to eliminate lookahead: signal at t -> execute at t+1.
        df["micro_position"] = target_position.shift(1).fillna(0.0).astype(np.int8)
        return df