"""
intraday_regime.py - 日内市场状态判别 + 策略匹配

来源: DeepSeek chat share 26iikphv8h94feze9q
核心算法:
  1. 趋势效率 R² (线性回归) - 判趋势 vs 震荡
  2. ADF 平稳检验 - 验证均值回归
  3. 历史波动率 - 区分高/低波动
  4. 开盘缺口 - 识别方向偏好

决策树:
  R² > 0.75  → 趋势跟踪 (顺势)
  R² < 0.30 + ADF 平稳 + 低波动 → 网格交易
  R² < 0.30 + ADF 平稳 + 高波动 → 布林带回归
  其他 → NO_TRADE (暂停)

⚠️ 这是港美股日内做T 元策略, 跟 crypto-t-monitor / longbridge-t-monitor 都独立
"""
import numpy as np
import pandas as pd
from typing import Dict, List, Tuple, Optional
from dataclasses import dataclass, field
from enum import Enum
import warnings
warnings.filterwarnings('ignore')


class MarketRegime(Enum):
    STRONG_TREND_UP = "强趋势上涨"
    STRONG_TREND_DOWN = "强趋势下跌"
    HIGH_VOL_SHAKE = "高波动剧烈震荡"
    LOW_VOL_STABLE = "低波动平稳震荡"
    CHAOTIC = "混乱无序"
    UNKNOWN = "无法判断"


class StrategyType(Enum):
    TREND_FOLLOWING = "趋势跟踪做T"
    GRID_TRADING = "网格交易做T"
    BOLLINGER_REVERSION = "布林带回归做T"
    NO_TRADE = "暂停交易"


@dataclass
class MarketDiagnosis:
    regime: MarketRegime
    r_squared: float
    volatility: float
    adf_pvalue: float
    recommended_strategy: StrategyType
    strategy_params: Dict
    confidence: float
    reasoning: str


class IntradayStrategySelector:
    """
    根据 5min K 线自动判别市场状态 + 推荐日内做T 策略
    
    用法:
        selector = IntradayStrategySelector()
        diagnosis = selector.diagnose(df_5min, open_gap_pct=0.3)
        print(diagnosis.recommended_strategy)
    """

    def __init__(self,
                 trend_r2_threshold: float = 0.75,
                 chaos_r2_threshold: float = 0.30,
                 adf_significance: float = 0.05,
                 vol_lookback: int = 20,
                 high_vol_threshold: float = 0.30,
                 grid_count: int = 3,
                 bb_period: int = 20,
                 bb_std: float = 2.0,
                 ema_period: int = 5):
        self.trend_r2_threshold = trend_r2_threshold
        self.chaos_r2_threshold = chaos_r2_threshold
        self.adf_significance = adf_significance
        self.vol_lookback = vol_lookback
        self.high_vol_threshold = high_vol_threshold
        self.grid_count = grid_count
        self.bb_period = bb_period
        self.bb_std = bb_std
        self.ema_period = ema_period

    def diagnose(self, df: pd.DataFrame, open_gap_pct: float = 0.0) -> MarketDiagnosis:
        if len(df) < 10:
            raise ValueError(f"需要至少 10 根 K 线, 拿到 {len(df)}")

        for col in ['open', 'high', 'low', 'close']:
            if col not in df.columns:
                raise ValueError(f"df 缺少 '{col}' 列")

        prices = df['close'].values
        r_squared = self._calculate_r_squared(prices)
        volatility = self._calculate_historical_volatility(prices)
        adf_pvalue = self._adf_test(prices)
        slope = self._calculate_trend_slope(prices)

        regime, confidence, reasoning = self._classify_regime(
            r_squared, volatility, adf_pvalue, slope, open_gap_pct
        )

        strategy, params = self._match_strategy(regime, df, volatility, r_squared)

        return MarketDiagnosis(
            regime=regime,
            r_squared=round(r_squared, 4),
            volatility=round(volatility, 4),
            adf_pvalue=round(adf_pvalue, 4),
            recommended_strategy=strategy,
            strategy_params=params,
            confidence=round(confidence, 2),
            reasoning=reasoning,
        )

    def _calculate_r_squared(self, prices: np.ndarray) -> float:
        n = len(prices)
        if n < 2:
            return 0.0
        x = np.arange(1, n + 1)
        y = prices
        x_mean = np.mean(x)
        y_mean = np.mean(y)
        numerator = np.sum((x - x_mean) * (y - y_mean))
        denominator = np.sqrt(np.sum((x - x_mean) ** 2) * np.sum((y - y_mean) ** 2))
        if denominator == 0:
            return 0.0
        r = numerator / denominator
        return r ** 2

    def _calculate_historical_volatility(self, prices: np.ndarray) -> float:
        if len(prices) < 2:
            return 0.0
        log_returns = np.diff(np.log(prices))
        return float(np.std(log_returns) * np.sqrt(252))

    def _calculate_trend_slope(self, prices: np.ndarray) -> float:
        n = len(prices)
        if n < 2:
            return 0.0
        x = np.arange(1, n + 1)
        y = prices
        x_mean = np.mean(x)
        y_mean = np.mean(y)
        slope = np.sum((x - x_mean) * (y - y_mean)) / np.sum((x - x_mean) ** 2)
        return float(slope / y_mean) if y_mean != 0 else 0.0

    def _adf_test(self, prices: np.ndarray) -> float:
        """简化 ADF (用一阶差分自相关近似)
        生产建议用 statsmodels.tsa.stattools.adfuller
        """
        try:
            from statsmodels.tsa.stattools import adfuller
            result = adfuller(prices, autolag='AIC')
            return float(result[1])
        except ImportError:
            # Fallback: 用一阶差分自相关近似
            diffs = np.diff(prices)
            if len(diffs) < 10:
                return 1.0
            autocorr = float(np.corrcoef(diffs[:-1], diffs[1:])[0, 1])
            if abs(autocorr) < 0.1:
                return 0.01
            elif abs(autocorr) < 0.3:
                return 0.05
            elif abs(autocorr) < 0.5:
                return 0.15
            else:
                return 0.50

    def _classify_regime(self, r2: float, vol: float, adf_p: float,
                         slope: float, gap: float) -> Tuple[MarketRegime, float, str]:

        if r2 > self.trend_r2_threshold:
            regime = MarketRegime.STRONG_TREND_UP if slope > 0.002 else MarketRegime.STRONG_TREND_DOWN
            confidence = min(r2, 1.0)
            reasoning = (f"R²={r2:.3f}>0.75, 市场呈现强趋势状态。"
                         f"线性回归斜率={slope:.4f}, 方向明确。"
                         f"此类行情适合顺势做T,严禁逆势网格。")
            return regime, confidence, reasoning

        if r2 < self.chaos_r2_threshold:
            if adf_p < self.adf_significance:
                if vol > self.high_vol_threshold:
                    regime = MarketRegime.HIGH_VOL_SHAKE
                    confidence = 0.70
                    reasoning = (f"R²={r2:.3f}<0.30, ADF p={adf_p:.3f}<0.05, "
                                 f"但波动率={vol:.2%}偏高。市场为高波动震荡,"
                                 f"适宜宽间距的逆势策略,需严格止损。")
                else:
                    regime = MarketRegime.LOW_VOL_STABLE
                    confidence = 0.85
                    reasoning = (f"R²={r2:.3f}<0.30, ADF p={adf_p:.3f}<0.05, "
                                 f"波动率={vol:.2%}适中。经典震荡市,"
                                 f"是网格和布林带回归策略的理想环境。")
            else:
                regime = MarketRegime.CHAOTIC
                confidence = 0.40
                reasoning = (f"R²={r2:.3f}<0.30, 但 ADF p={adf_p:.3f}>0.05, "
                             f"价格不具均值回归特性,属混乱状态,建议观望。")
            return regime, confidence, reasoning

        regime = MarketRegime.UNKNOWN
        confidence = 0.30
        reasoning = (f"R²={r2:.3f} 处于过渡区间(0.30-0.75), "
                     f"市场方向不明。建议等待模式清晰后再交易。")
        return regime, confidence, reasoning

    def _match_strategy(self, regime: MarketRegime, df: pd.DataFrame,
                        vol: float, r2: float) -> Tuple[StrategyType, Dict]:
        current_price = float(df['close'].iloc[-1])
        params = {}

        if regime == MarketRegime.STRONG_TREND_UP:
            strategy = StrategyType.TREND_FOLLOWING
            ema = float(df['close'].ewm(span=self.ema_period).mean().iloc[-1])
            params = {
                "direction": "long_only",
                "entry_trigger": f"价格回踩 {ema:.2f} (EMA{self.ema_period}) 不破",
                "stop_loss": f"{ema * 0.995:.2f}",
                "take_profit": f"{current_price * 1.02:.2f}",
            }

        elif regime == MarketRegime.STRONG_TREND_DOWN:
            strategy = StrategyType.TREND_FOLLOWING
            ema = float(df['close'].ewm(span=self.ema_period).mean().iloc[-1])
            params = {
                "direction": "short_only",
                "entry_trigger": f"价格反弹至 {ema:.2f} (EMA{self.ema_period}) 受阻",
                "stop_loss": f"{ema * 1.005:.2f}",
                "take_profit": f"{current_price * 0.98:.2f}",
            }

        elif regime == MarketRegime.LOW_VOL_STABLE:
            strategy = StrategyType.GRID_TRADING
            avg_amplitude = float(((df['high'] - df['low']) / df['close']).mean())
            grid_spacing = max(avg_amplitude * 0.8, 0.005)
            params = {
                "grid_spacing": f"{grid_spacing:.3%}",
                "grid_levels": self.grid_count,
                "base_price": f"{current_price:.2f}",
                "reverse_at_boundary": True,
            }

        elif regime == MarketRegime.HIGH_VOL_SHAKE:
            strategy = StrategyType.BOLLINGER_REVERSION
            rolling_std = float(df['close'].rolling(self.bb_period).std().iloc[-1])
            ma = float(df['close'].rolling(self.bb_period).mean().iloc[-1])
            upper = ma + self.bb_std * rolling_std
            lower = ma - self.bb_std * rolling_std
            params = {
                "upper_band": f"{upper:.2f}",
                "lower_band": f"{lower:.2f}",
                "sell_at_upper": True,
                "buy_at_lower": True,
                "stop_if_break": True,
            }

        else:
            strategy = StrategyType.NO_TRADE
            params = {"reason": "市场状态不清晰,等待趋势或明确震荡信号"}

        return strategy, params


def diagnose_market(df: pd.DataFrame, open_gap_pct: float = 0.0) -> MarketDiagnosis:
    """便捷函数"""
    return IntradayStrategySelector().diagnose(df, open_gap_pct)


if __name__ == '__main__':
    np.random.seed(42)

    # 场景 1: 震荡市
    n = 25
    base = 10.0
    noise = np.random.randn(n) * 0.05
    close = base + noise
    high = close + np.abs(np.random.randn(n) * 0.03)
    low = close - np.abs(np.random.randn(n) * 0.03)
    df = pd.DataFrame({
        'open': close - 0.01,
        'high': high,
        'low': low,
        'close': close,
        'volume': np.random.randint(1000, 5000, n),
    })
    diag = diagnose_market(df, open_gap_pct=0.0)
    print(f"场景 1 (震荡市): {diag.regime.value} | R²={diag.r_squared} | 策略: {diag.recommended_strategy.value}")
    print(f"  推理: {diag.reasoning}\n")

    # 场景 2: 强趋势
    close2 = base + np.cumsum(np.random.randn(n) * 0.02) * 2  # 上升趋势
    high2 = close2 + 0.05
    low2 = close2 - 0.05
    df2 = pd.DataFrame({
        'open': close2 - 0.01,
        'high': high2,
        'low': low2,
        'close': close2,
        'volume': np.random.randint(1000, 5000, n),
    })
    diag2 = diagnose_market(df2, open_gap_pct=0.5)
    print(f"场景 2 (趋势市): {diag2.regime.value} | R²={diag2.r_squared} | 策略: {diag2.recommended_strategy.value}")
    print(f"  推理: {diag2.reasoning}")
    print(f"  参数: {diag2.strategy_params}")