import numpy as np
import talib
from typing import Optional, Any, List, Dict
from src.core_data import Bar, Order
from src.strategies.base_strategy import Strategy


class TVDZScoreStrategy(Strategy):
    """
    内嵌 TVD (Condat 算法) + Z-Score ATR 的趋势突破策略。
    无任何外部依赖（如 pytv），纯 NumPy 实现。
    """

    def __init__(
        self,
        context: Any,
        main_symbol: str,
        enable_log: bool,
        trade_volume: int,
        tvd_lam: float = 50.0,
        atr_window: int = 14,
        z_window: int = 100,
        vol_threshold: float = -0.5,
        entry_threshold_atr: float = 3.0,
        stop_atr_multiplier: float = 3.0,
        order_direction: Optional[List[str]] = None,
    ):
        super().__init__(context, main_symbol, enable_log)
        self.trade_volume = trade_volume
        self.order_direction = order_direction or ["BUY", "SELL"]
        self.tvd_lam = tvd_lam
        self.atr_window = atr_window
        self.z_window = z_window
        self.vol_threshold = vol_threshold
        self.entry_threshold_atr = entry_threshold_atr
        self.stop_atr_multiplier = stop_atr_multiplier

        self.position_meta: Dict[str, Any] = self.context.load_state()
        self.main_symbol = main_symbol
        self.order_id_counter = 0

        self.log(f"TVDZScoreStrategy Initialized | λ={tvd_lam}, VolThresh={vol_threshold}")

    @staticmethod
    def _tvd_condat(y, lam):
        """Condat's O(N) TVD algorithm."""
        n = y.size
        if n == 0:
            return y.copy()
        x = y.astype(np.float64)
        k = 0
        k0 = 0
        vmin = x[0] - lam
        vmax = x[0] + lam
        for i in range(1, n):
            if x[i] < vmin:
                while k < i:
                    x[k] = vmin
                    k += 1
                k0 = i
                vmin = x[i] - lam
                vmax = x[i] + lam
            elif x[i] > vmax:
                while k < i:
                    x[k] = vmax
                    k += 1
                k0 = i
                vmin = x[i] - lam
                vmax = x[i] + lam
            else:
                vmin = max(vmin, x[i] - lam)
                vmax = min(vmax, x[i] + lam)
                if vmin > vmax:
                    k = k0
                    s = np.sum(x[k0:i+1])
                    s /= (i - k0 + 1)
                    x[k0:i+1] = s
                    k = i + 1
                    k0 = k
                    if k0 < n:
                        vmin = x[k0] - lam
                        vmax = x[k0] + lam
        while k < n:
            x[k] = vmin
            k += 1
        return x

    def _compute_zscore_atr_last(self, high, low, close) -> float:
        n = len(close)
        min_req = self.atr_window + self.z_window - 1
        if n < min_req:
            return np.nan
        start = max(0, n - (self.z_window + self.atr_window))
        seg_h, seg_l, seg_c = high[start:], low[start:], close[start:]
        atr_full = talib.ATR(seg_h, seg_l, seg_c, timeperiod=self.atr_window)
        atr_valid = atr_full[self.atr_window - 1:]
        if len(atr_valid) < self.z_window:
            return np.nan
        window_atr = atr_valid[-self.z_window:]
        mu = np.mean(window_atr)
        sigma = np.std(window_atr)
        last_atr = window_atr[-1]
        return (last_atr - mu) / sigma if sigma > 1e-12 else 0.0

    def on_open_bar(self, open_price: float, symbol: str):
        self.symbol = symbol
        bar_history = self.get_bar_history()
        if len(bar_history) < max(100, self.atr_window + self.z_window):
            return

        closes = np.array([b.close for b in bar_history], dtype=np.float64)
        highs = np.array([b.high for b in bar_history], dtype=np.float64)
        lows = np.array([b.low for b in bar_history], dtype=np.float64)

        # === TVD 平滑 ===
        tvd_prices = self._tvd_condat(closes, self.tvd_lam)
        tvd_price = tvd_prices[-1]

        # === Z-Score ATR ===
        current_atr = talib.ATR(highs, lows, closes, timeperiod=self.atr_window)[-1]
        if current_atr <= 0:
            return

        deviation = closes[-1] - tvd_price
        deviation_in_atr = deviation / current_atr

        position_volume = self.get_current_positions().get(self.symbol, 0)
        if position_volume != 0:
            self.manage_open_position(position_volume, bar_history[-1], current_atr, tvd_price)
            return

        direction = None
        if "BUY" in self.order_direction and deviation_in_atr > self.entry_threshold_atr:
            direction = "BUY"
        elif "SELL" in self.order_direction and deviation_in_atr < -self.entry_threshold_atr:
            direction = "SELL"

        if direction:
            self.log(f"Signal Fired | Dir: {direction}, Dev: {deviation_in_atr:.2f} ATR")
            entry_price = closes[-1]
            stop_loss = (
                entry_price - self.stop_atr_multiplier * current_atr
                if direction == "BUY"
                else entry_price + self.stop_atr_multiplier * current_atr
            )
            meta = {"entry_price": entry_price, "stop_loss": stop_loss}
            self.send_market_order(direction, self.trade_volume, "OPEN", meta)
            self.save_state(self.position_meta)

    def manage_open_position(self, volume: int, current_bar: Bar, current_atr: float, tvd_price: float):
        meta = self.position_meta.get(self.symbol)
        if not meta:
            return
        stop_loss = meta["stop_loss"]
        if (volume > 0 and current_bar.low <= stop_loss) or (volume < 0 and current_bar.high >= stop_loss):
            self.log(f"Stop Loss Hit at {stop_loss:.4f}")
            self.close_position("CLOSE_LONG" if volume > 0 else "CLOSE_SHORT", abs(volume))

    def close_position(self, direction: str, volume: int):
        self.send_market_order(direction, volume, offset="CLOSE")
        if self.symbol in self.position_meta:
            del self.position_meta[self.symbol]
            self.save_state(self.position_meta)

    def send_market_order(self, direction: str, volume: int, offset: str, meta: Optional[Dict] = None):
        if offset == "OPEN" and meta:
            self.position_meta[self.symbol] = meta
        order_id = f"{self.symbol}_{direction}_MARKET_{self.order_id_counter}"
        self.order_id_counter += 1
        order = Order(id=order_id, symbol=self.symbol, direction=direction, volume=volume,
                      price_type="MARKET", submitted_time=self.get_current_time(), offset=offset)
        self.send_order(order)

    def on_rollover(self, old_symbol: str, new_symbol: str):
        super().on_rollover(old_symbol, new_symbol)
        self.position_meta = {}
        self.log("Rollover: Strategy state reset.")