# src/risk/stops/atr_stop.py
import pandas as pd
import numpy as np
from src.risk.stops.base import StopLoss
from src.core.trade import TradeType


class ATRStop(StopLoss):
    """
    Stop basado en ATR (Average True Range).
    """

    def __init__(self, atr_period: int = 14, multiplier: float = 2.0):
        if atr_period <= 0:
            raise ValueError("atr_period debe ser > 0")
        if multiplier <= 0:
            raise ValueError("multiplier debe ser > 0")

        self.atr_period = atr_period
        self.multiplier = multiplier

    def _calculate_atr(self, data: pd.DataFrame) -> pd.Series:
        """
        Calcula el ATR clásico (Wilder).
        """
        high = data["high"]
        low = data["low"]
        close = data["close"].shift(1)

        tr = pd.concat(
            [
                high - low,
                (high - close).abs(),
                (low - close).abs(),
            ],
            axis=1,
        ).max(axis=1)

        atr = tr.ewm(alpha=1/self.atr_period, adjust=False).mean()
        return atr

    def get_stop_price(
        self,
        *,
        data: pd.DataFrame,
        idx: int,
        entry_price: float,
        trade_type: TradeType
    ) -> float:

        if not {"high", "low", "close"}.issubset(data.columns):
            raise ValueError("ATRStop requiere columnas high, low, close")

        atr_series = self._calculate_atr(data)

        atr_value = atr_series.iloc[idx]

        if not np.isfinite(atr_value):
            raise ValueError("ATR no disponible en este índice")

        distance = atr_value * self.multiplier

        if trade_type == TradeType.LONG:
            return entry_price - distance

        elif trade_type == TradeType.SHORT:
            return entry_price + distance

        else:
            raise ValueError(f"TradeType no soportado: {trade_type}")
