# -*- coding: utf-8 -*-
"""
蒸鱼股票分析系统 - 增强版市场状态推导模块
基于QuantGPT market_regime.py适配
从技术指标信号推导: regime(趋势/高波动/震荡) + style(大盘/小盘/动量/反转) + risk_level + headline
与现有market_regime.py互补：现有模块基于ADX方向，本模块基于多因子信号
"""

import numpy as np
import pandas as pd
from datetime import datetime
import warnings
warnings.filterwarnings('ignore')


class FactorMarketRegime:
    """从因子信号推导市场状态"""

    def __init__(self):
        self.lookback = 60
        self.vol_window = 20
        self.trend_window = 20

    def detect(self, close, high=None, low=None, volume=None, returns=None):
        """
        从价格数据推导综合市场状态

        Parameters
        ----------
        close : array-like - 收盘价序列
        high : array-like, optional
        low : array-like, optional
        volume : array-like, optional
        returns : array-like, optional - 收益率序列（如不传则自动计算）

        Returns
        -------
        dict : {
            regime: str - 趋势市/高波动/震荡市,
            style: str - 大盘·动量驱动/大盘·反转驱动/小盘·均衡驱动,
            risk_level: str - 低/中/高,
            headline: str - 一句话总结,
            details: dict - 各维度详情,
            signals: list - 关键信号,
        }
        """
        close = np.asarray(close, dtype=float)
        if high is not None:
            high = np.asarray(high, dtype=float)
        if low is not None:
            low = np.asarray(low, dtype=float)
        if volume is not None:
            volume = np.asarray(volume, dtype=float)

        # 计算收益率
        if returns is None:
            returns = np.diff(close) / close[:-1]
            returns = np.concatenate([[0], returns])

        n = len(close)
        if n < 30:
            return self._default_result()

        # ─── 各维度计算 ───
        trend_score = self._calc_trend_score(close, returns)
        volatility_score = self._calc_volatility_score(returns)
        momentum_score = self._calc_momentum_score(returns)
        breadth_score = self._calc_breadth_score(close, high, low)
        trend_direction = self._calc_trend_direction(close, returns)
        trend_speed = self._calc_trend_speed(close, returns, high, low)

        # ─── 推导regime ───
        if trend_score > 0.6 and volatility_score < 0.6:
            regime = '趋势市'
        elif volatility_score > 0.7:
            regime = '高波动'
        else:
            regime = '震荡市'

        # ─── 推导style ───
        if momentum_score > 0.55:
            style_prefix = '动量驱动'
        elif momentum_score < 0.35:
            style_prefix = '反转驱动'
        else:
            style_prefix = '均衡驱动'

        if trend_score > 0.55:
            size = '大盘'
        elif trend_score < 0.35:
            size = '小盘'
        else:
            size = '中盘'

        style = f'{size}·{style_prefix}'

        # ─── 推导risk_level ───
        if volatility_score > 0.7 and breadth_score < 0.4:
            risk_level = '高'
        elif volatility_score < 0.4 and breadth_score > 0.6:
            risk_level = '低'
        else:
            risk_level = '中'

        # ─── headline ───
        headline = self._generate_headline(regime, style, risk_level,
                                            trend_score, volatility_score, momentum_score, close)

        # ─── 信号列表 ───
        signals = self._collect_signals(trend_score, volatility_score, momentum_score,
                                          breadth_score, close, returns)

        return {
            'regime': regime,
            'trend_direction': trend_direction,
            'trend_speed': trend_speed,
            'style': style,
            'risk_level': risk_level,
            'headline': headline,
            'scores': {
                'trend': round(trend_score, 4),
                'volatility': round(volatility_score, 4),
                'momentum': round(momentum_score, 4),
                'breadth': round(breadth_score, 4),
            },
            'signals': signals,
            'data_points': n,
            'timestamp': datetime.now().isoformat(),
        }

    def _calc_trend_score(self, close, returns):
        """趋势强度 (0-1)"""
        n = len(close)
        if n < 20:
            return 0.5

        # MA斜率
        ma20 = np.convolve(close, np.ones(20)/20, mode='valid')
        if len(ma20) >= 2:
            ma_slope = (ma20[-1] - ma20[-2]) / ma20[-2] * 100
            trend_from_slope = min(abs(ma_slope) / 0.5, 1.0)
        else:
            trend_from_slope = 0.5

        # 价格相对MA20位置
        if len(close) >= 20:
            price_vs_ma = (close[-1] - ma20[-1]) / ma20[-1]
            trend_from_position = min(abs(price_vs_ma) / 0.05, 1.0)
        else:
            trend_from_position = 0.5

        # ADX近似（用收益率方向一致性）
        recent_ret = returns[-20:] if len(returns) >= 20 else returns
        positive_days = np.sum(recent_ret > 0)
        direction_consistency = abs(positive_days / len(recent_ret) - 0.5) * 2

        # 综合
        return trend_from_slope * 0.4 + trend_from_position * 0.3 + direction_consistency * 0.3

    def _calc_trend_direction(self, close, returns):
        """趋势方向: UP / DOWN / FLAT"""
        n = len(close)
        if n < 20:
            return 'FLAT'
        ma20 = np.convolve(close, np.ones(20)/20, mode='valid')
        ma_slope = (ma20[-1] - ma20[-2]) / ma20[-2] * 100 if len(ma20) >= 2 else 0
        price_vs_ma = (close[-1] - ma20[-1]) / ma20[-1] if len(ma20) > 0 else 0
        ret_20 = (close[-1] - close[-20]) / close[-20] if n >= 20 else 0
        recent_ret = returns[-20:] if len(returns) >= 20 else returns
        positive_days = int(np.sum(recent_ret > 0))

        down_score = 0
        if ma_slope < -0.10: down_score += 1
        if price_vs_ma < -0.02: down_score += 1
        if ret_20 < -0.02: down_score += 1
        if positive_days <= 6: down_score += 1

        up_score = 0
        if ma_slope > 0.10: up_score += 1
        if price_vs_ma > 0.02: up_score += 1
        if ret_20 > 0.02: up_score += 1
        if positive_days >= 14: up_score += 1

        if down_score >= 3:
            return 'DOWN'
        elif up_score >= 3:
            return 'UP'
        else:
            return 'FLAT'

    def _calc_trend_speed(self, close, returns, high, low):
        """趋势速度: FAST / NORMAL / SLOW"""
        n = len(close)
        if n < 20:
            return 'NORMAL'
        ret_20 = (close[-1] - close[-20]) / close[-20] if n >= 20 else 0
        recent_ret = returns[-20:] if len(returns) >= 20 else returns

        # 连续同向天数
        direction = 1 if ret_20 >= 0 else -1
        consecutive = 0
        max_consecutive = 0
        for r in recent_ret:
            if (direction > 0 and r > 0) or (direction < 0 and r < 0):
                consecutive += 1
                max_consecutive = max(max_consecutive, consecutive)
            else:
                consecutive = 0

        # ATR%
        atr_pct = 0
        if high is not None and low is not None and n >= 20:
            hi = high[-20:]
            lo = low[-20:]
            cl = close[-20:]
            tr1 = hi - lo
            tr2 = np.abs(hi - np.roll(cl, 1))
            tr3 = np.abs(lo - np.roll(cl, 1))
            tr = np.maximum(np.maximum(tr1, tr2), tr3)
            atr = np.mean(tr)
            atr_pct = atr / close[-1] if close[-1] > 0 else 0

        # 速度评分
        speed_score = 0
        abs_ret20 = abs(ret_20)
        if abs_ret20 > 0.10: speed_score += 2
        elif abs_ret20 > 0.05: speed_score += 1
        if max_consecutive >= 5: speed_score += 2
        elif max_consecutive >= 3: speed_score += 1
        if atr_pct > 0.03: speed_score += 1

        if speed_score >= 3:
            return 'FAST'
        elif speed_score >= 1:
            return 'NORMAL'
        else:
            return 'SLOW'

    def _calc_volatility_score(self, returns):
        """波动率强度 (0-1, 越大越高波)"""
        n = len(returns)
        if n < 20:
            return 0.5

        recent_vol = np.std(returns[-20:])
        long_vol = np.std(returns[-60:]) if n >= 60 else recent_vol
        if long_vol > 0:
            vol_ratio = recent_vol / long_vol
        else:
            vol_ratio = 1

        # 最大回撤
        cum_ret = np.cumprod(1 + returns[-60:]) if n >= 60 else np.cumprod(1 + returns)
        cum_max = np.maximum.accumulate(cum_ret)
        drawdown = (cum_ret - cum_max) / cum_max
        max_dd = abs(min(drawdown))

        vol_score = min(vol_ratio / 2.0, 1) * 0.5 + min(max_dd / 0.15, 1) * 0.5
        return vol_score

    def _calc_momentum_score(self, returns):
        """动量强度 (0-1, >0.5为动量主导)"""
        n = len(returns)
        if n < 20:
            return 0.5

        # 短期vs长期动量
        ret_5d = np.sum(returns[-5:]) if n >= 5 else np.sum(returns)
        ret_20d = np.sum(returns[-20:]) if n >= 20 else np.sum(returns)

        if abs(ret_20d) > 0.001:
            momentum_ratio = ret_5d / ret_20d
        else:
            momentum_ratio = 0.5

        # 价格动能
        if len(returns) >= 20:
            positive_mom = np.sum(returns[-20:] > 0) / 20
        else:
            positive_mom = 0.5

        score = 0.5 + (momentum_ratio - 0.5) * 0.3 + (positive_mom - 0.5) * 0.4
        return max(0, min(1, score))

    def _calc_breadth_score(self, close, high, low):
        """市场广度 (0-1, 越大越健康)"""
        n = len(close)
        if n < 20 or high is None or low is None:
            return 0.5

        # 价格在N日高低范围内的位置
        hi = high[-20:]
        lo = low[-20:]
        price_range = hi - lo
        valid = price_range > 0
        if valid.sum() == 0:
            return 0.5

        positions = (close[-20:] - lo) / price_range[valid]
        avg_position = np.nanmean(positions) if len(positions) > 0 else 0.5

        # 高点创新高能力
        recent_high = np.max(high[-10:])
        period_high = np.max(high[-60:]) if n >= 60 else recent_high
        high_ratio = recent_high / period_high if period_high > 0 else 0.5

        return avg_position * 0.5 + high_ratio * 0.5

    def _generate_headline(self, regime, style, risk_level,
                           trend_s, vol_s, mom_s, close):
        """生成一句话总结"""
        n = len(close)
        ret_20 = (close[-1] - close[-20]) / close[-20] if n >= 20 else 0

        direction = '上涨' if ret_20 > 0.02 else ('下跌' if ret_20 < -0.02 else '横盘')
        strength = '强势' if abs(ret_20) > 0.08 else '温和'

        parts = [f'当前市场处于{regime}状态', f'{direction}{strength}']
        parts.append(f'风格偏好{style}')
        parts.append(f'风险等级{risk_level}')

        return '，'.join(parts) + '。'

    def _collect_signals(self, trend_s, vol_s, mom_s, breadth_s, close, returns):
        """收集关键信号"""
        signals = []

        if trend_s > 0.7:
            signals.append({'type': 'strong_trend', 'weight': 2, 'desc': '趋势强劲'})
        elif trend_s < 0.3:
            signals.append({'type': 'weak_trend', 'weight': 1, 'desc': '趋势不明'})

        if vol_s > 0.7:
            signals.append({'type': 'high_volatility', 'weight': 2, 'desc': '波动率偏高'})
        elif vol_s < 0.3:
            signals.append({'type': 'low_volatility', 'weight': 1, 'desc': '波动率极低'})

        if mom_s > 0.65:
            signals.append({'type': 'momentum_dom', 'weight': 2, 'desc': '动量效应主导'})
        elif mom_s < 0.35:
            signals.append({'type': 'reversal_dom', 'weight': 2, 'desc': '反转效应显现'})

        if breadth_s > 0.6:
            signals.append({'type': 'healthy_breadth', 'weight': 1, 'desc': '市场广度良好'})
        elif breadth_s < 0.35:
            signals.append({'type': 'narrow_breadth', 'weight': 1, 'desc': '市场广度收窄'})

        # 极端信号
        n = len(close)
        if n >= 20:
            ret_20 = (close[-1] - close[-20]) / close[-20]
            if ret_20 > 0.1:
                signals.append({'type': 'overbought_warning', 'weight': 3, 'desc': f'20日涨幅{ret_20:.1%}，注意过热'})
            elif ret_20 < -0.1:
                signals.append({'type': 'oversold_signal', 'weight': 3, 'desc': f'20日跌幅{ret_20:.1%}，关注反弹'})

        return signals

    def _default_result(self):
        return {
            'regime': '震荡市', 'trend_direction': 'FLAT', 'trend_speed': 'NORMAL',
            'style': '中盘·均衡驱动', 'risk_level': '中',
            'headline': '数据不足，默认震荡市', 'scores': {
                'trend': 0.5, 'volatility': 0.5, 'momentum': 0.5, 'breadth': 0.5,
            },
            'signals': [], 'data_points': 0,
            'timestamp': datetime.now().isoformat(),
        }


# 便捷函数
def detect_factor_regime(close, high=None, low=None, volume=None, returns=None):
    """推导市场状态（全局便捷入口）"""
    detector = FactorMarketRegime()
    return detector.detect(close, high, low, volume, returns)


def get_regime_advice(regime_result):
    """根据市场状态生成操作建议"""
    regime = regime_result.get('regime', '震荡市')
    risk = regime_result.get('risk_level', '中')
    style = regime_result.get('style', '')

    advices = []

    if regime == '趋势市':
        if '动量' in style:
            advices.append({'action': '顺势而为', 'detail': '趋势+动量共振，适合追涨策略'})
        else:
            advices.append({'action': '趋势跟踪', 'detail': '趋势明确但动量衰减，注意止盈'})
    elif regime == '高波动':
        advices.append({'action': '降低仓位', 'detail': '高波动环境下控制风险优先'})
        advices.append({'action': '避免追涨', 'detail': '波动放大时追涨杀跌风险极高'})
    else:
        advices.append({'action': '区间操作', 'detail': '震荡市适用高抛低吸策略'})

    if risk == '高':
        advices.append({'action': '防守优先', 'detail': '风险偏高，建议降低仓位至半仓以下'})
    elif risk == '低':
        advices.append({'action': '可适度进攻', 'detail': '风险较低，可适当增加仓位'})

    return advices