#!/usr/bin/env python3
"""
微观结构因子计算模块
- CVD (累积成交量差)
- OFI (订单流不平衡)
- Delta 发散/收敛
- 大单冲击
- VPIN 近似
- 资金费率极值
- OI 价格配合度 (若有数据)
"""
import pandas as pd
import numpy as np
from typing import Dict, List, Optional, Tuple
from dataclasses import dataclass
from enum import Enum


class MicroSignal(Enum):
    STRONG_BULLISH = "STRONG_BULLISH"
    BULLISH = "BULLISH"
    NEUTRAL = "NEUTRAL"
    BEARISH = "BEARISH"
    STRONG_BEARISH = "STRONG_BEARISH"


@dataclass
class MicroFactors:
    """单根 K 线的微观因子"""
    timestamp: pd.Timestamp
    
    # Taker 量基础
    delta: float = 0.0
    cvd: float = 0.0
    ofi: float = 0.0
    
    # 发散/收敛
    divergence_bull: bool = False
    divergence_bear: bool = False
    convergence_bull: bool = False
    convergence_bear: bool = False
    
    # 大单
    large_buy: bool = False
    large_sell: bool = False
    large_trade: bool = False
    
    # VPIN 近似 (滚动窗口买卖不平衡)
    vpin_proxy: float = 0.0
    
    # 资金费率
    funding_rate: float = 0.0
    fr_extreme_long: bool = False
    fr_extreme_short: bool = False
    fr_zscore: float = 0.0
    
    # OI 变化 (若有)
    oi_change: float = 0.0
    oi_price_align: bool = False
    
    # 综合评分
    bullish_score: float = 0.0
    bearish_score: float = 0.0
    net_score: float = 0.0  # -1 到 1
    
    def signal(self) -> MicroSignal:
        if self.net_score >= 0.6:
            return MicroSignal.STRONG_BULLISH
        elif self.net_score >= 0.2:
            return MicroSignal.BULLISH
        elif self.net_score <= -0.6:
            return MicroSignal.STRONG_BEARISH
        elif self.net_score <= -0.2:
            return MicroSignal.BEARISH
        return MicroSignal.NEUTRAL


def compute_microstructure_factors(df: pd.DataFrame, 
                                    window_div: int = 20,
                                    window_vpin: int = 50,
                                    window_fr: int = 200,
                                    large_trade_pct: float = 0.99) -> pd.DataFrame:
    """
    计算完整微观因子，原地修改并返回 DataFrame
    
    必需列：
    - taker_buy_v, taker_sell_v (或 taker_buy_volume, taker_sell_volume)
    - v (volume)
    - c (close), h (high), l (low)
    - last_funding_rate (可选)
    - open_interest (可选)
    """
    df = df.copy()
    
    # === 1. 统一列名 ===
    # Taker 买卖量
    buy_col = None
    sell_col = None
    for c in ['taker_buy_v', 'taker_buy_volume', 'taker_buy_quote']:
        if c in df.columns:
            buy_col = c
            break
    for c in ['taker_sell_v', 'taker_sell_volume', 'taker_sell_quote']:
        if c in df.columns:
            sell_col = c
            break
    
    if not buy_col or not sell_col:
        # 尝试从 taker_buy_ratio 推算
        if 'taker_buy_ratio' in df.columns and 'v' in df.columns:
            df['taker_buy_v_calc'] = df['v'] * df['taker_buy_ratio']
            df['taker_sell_v_calc'] = df['v'] - df['taker_buy_v_calc']
            buy_col = 'taker_buy_v_calc'
            sell_col = 'taker_sell_v_calc'
        else:
            raise ValueError("找不到 Taker 买卖量列")
    
    # === 2. Delta & CVD ===
    df['delta'] = df[buy_col] - df[sell_col]
    df['cvd'] = df['delta'].cumsum()
    
    # === 3. OFI (Order Flow Imbalance) ===
    df['ofi'] = df['delta'] / (df[buy_col] + df[sell_col] + 1e-9)
    
    # === 4. Delta 发散/收敛 (滚动窗口高低点对比) ===
    # 价格创新高/新低 vs CVD 创新高/新低
    df['price_hh'] = df['c'].rolling(window_div).max()
    df['price_ll'] = df['c'].rolling(window_div).min()
    df['cvd_hh'] = df['cvd'].rolling(window_div).max()
    df['cvd_ll'] = df['cvd'].rolling(window_div).min()
    
    # 价格位置
    df['at_price_hh'] = (df['c'] == df['price_hh'])
    df['at_price_ll'] = (df['c'] == df['price_ll'])
    df['at_cvd_hh'] = (df['cvd'] == df['cvd_hh'])
    df['at_cvd_ll'] = (df['cvd'] == df['cvd_ll'])
    
    # 熊市发散：价格创新高，CVD 未创新高
    df['divergence_bear'] = df['at_price_hh'] & ~df['at_cvd_hh']
    # 牛市发散：价格创新低，CVD 未创新低
    df['divergence_bull'] = df['at_price_ll'] & ~df['at_cvd_ll']
    # 牛市收敛：价格创新高，CVD 同步创新高
    df['convergence_bull'] = df['at_price_hh'] & df['at_cvd_hh']
    # 熊市收敛：价格创新低，CVD 同步创新低
    df['convergence_bear'] = df['at_price_ll'] & df['at_cvd_ll']
    
    # === 5. 大单冲击 ===
    vol_99 = df['v'].rolling(100).quantile(large_trade_pct)
    df['large_trade'] = df['v'] > vol_99
    df['large_buy'] = df['large_trade'] & (df['delta'] > 0)
    df['large_sell'] = df['large_trade'] & (df['delta'] < 0)
    
    # === 6. VPIN 近似 ===
    # 滚动窗口内买卖不平衡度
    buy_sum = df[buy_col].rolling(window_vpin).sum()
    sell_sum = df[sell_col].rolling(window_vpin).sum()
    df['vpin_proxy'] = (buy_sum - sell_sum) / (buy_sum + sell_sum + 1e-9)
    
    # === 7. 资金费率因子 ===
    if 'last_funding_rate' in df.columns:
        df['funding_rate'] = df['last_funding_rate']
        # Z-score
        fr_mean = df['funding_rate'].rolling(window_fr).mean()
        fr_std = df['funding_rate'].rolling(window_fr).std()
        df['fr_zscore'] = (df['funding_rate'] - fr_mean) / (fr_std + 1e-9)
        
        # 极值 (95/5 分位)
        fr_95 = df['funding_rate'].rolling(window_fr).quantile(0.95)
        fr_05 = df['funding_rate'].rolling(window_fr).quantile(0.05)
        df['fr_extreme_long'] = df['funding_rate'] > fr_95
        df['fr_extreme_short'] = df['funding_rate'] < fr_05
    else:
        df['funding_rate'] = 0.0
        df['fr_zscore'] = 0.0
        df['fr_extreme_long'] = False
        df['fr_extreme_short'] = False
    
    # === 8. OI 因子 (若有) ===
    if 'open_interest' in df.columns:
        df['oi_change'] = df['open_interest'].diff()
        price_change = df['c'].diff()
        df['oi_price_align'] = np.sign(df['oi_change']) == np.sign(price_change)
    else:
        df['oi_change'] = 0.0
        df['oi_price_align'] = False
    
    # === 9. 综合评分 ===
    df['bullish_score'] = 0.0
    df['bearish_score'] = 0.0
    
    # Delta 方向 (权重 0.2)
    df.loc[df['delta'] > 0, 'bullish_score'] += 0.2
    df.loc[df['delta'] < 0, 'bearish_score'] += 0.2
    
    # CVD 趋势 (权重 0.15)
    cvd_slope = df['cvd'].diff(5)
    df.loc[cvd_slope > 0, 'bullish_score'] += 0.15
    df.loc[cvd_slope < 0, 'bearish_score'] += 0.15
    
    # OFI (权重 0.15)
    df.loc[df['ofi'] > 0.1, 'bullish_score'] += 0.15
    df.loc[df['ofi'] < -0.1, 'bearish_score'] += 0.15
    df.loc[df['ofi'] > 0.3, 'bullish_score'] += 0.1  # 额外奖励
    df.loc[df['ofi'] < -0.3, 'bearish_score'] += 0.1
    
    # 发散/收敛 (权重 0.2) - 强信号
    df.loc[df['divergence_bull'], 'bullish_score'] += 0.3
    df.loc[df['divergence_bear'], 'bearish_score'] += 0.3
    df.loc[df['convergence_bull'], 'bullish_score'] += 0.1
    df.loc[df['convergence_bear'], 'bearish_score'] += 0.1
    
    # 大单 (权重 0.1)
    df.loc[df['large_buy'], 'bullish_score'] += 0.15
    df.loc[df['large_sell'], 'bearish_score'] += 0.15
    
    # VPIN (权重 0.1)
    df.loc[df['vpin_proxy'] > 0.2, 'bullish_score'] += 0.1
    df.loc[df['vpin_proxy'] < -0.2, 'bearish_score'] += 0.1
    df.loc[df['vpin_proxy'] > 0.4, 'bullish_score'] += 0.1
    df.loc[df['vpin_proxy'] < -0.4, 'bearish_score'] += 0.1
    
    # 资金费率 (权重 0.1) - 反向指标
    df.loc[df['fr_extreme_short'], 'bullish_score'] += 0.15  # 费率极负 = 空头拥挤 = 反转多
    df.loc[df['fr_extreme_long'], 'bearish_score'] += 0.15   # 费率极正 = 多头拥挤 = 反转空
    df.loc[(df['fr_zscore'] > 0) & (df['fr_zscore'] < 2), 'bullish_score'] += 0.05  # 适度正费率 = 多头健康
    df.loc[(df['fr_zscore'] < 0) & (df['fr_zscore'] > -2), 'bearish_score'] += 0.05
    
    # OI 配合 (权重 0.05)
    df.loc[df['oi_price_align'], 'bullish_score'] += 0.05
    df.loc[df['oi_price_align'], 'bearish_score'] += 0.05
    
    # 归一化净分数
    df['net_score'] = (df['bullish_score'] - df['bearish_score']).clip(-1, 1)
    
    # === 10. 信号分类 ===
    def score_to_signal(s):
        if s >= 0.6: return 'STRONG_BULLISH'
        elif s >= 0.2: return 'BULLISH'
        elif s <= -0.6: return 'STRONG_BEARISH'
        elif s <= -0.2: return 'BEARISH'
        return 'NEUTRAL'
    
    df['micro_signal'] = df['net_score'].apply(score_to_signal)
    
    return df


def get_latest_micro_factors(df: pd.DataFrame) -> Dict:
    """获取最新一根 K 线的微观因子摘要"""
    last = df.iloc[-1]
    
    return {
        'timestamp': str(last.name),
        'delta': float(last['delta']),
        'cvd': float(last['cvd']),
        'ofi': float(last['ofi']),
        'divergence_bull': bool(last['divergence_bull']),
        'divergence_bear': bool(last['divergence_bear']),
        'convergence_bull': bool(last['convergence_bull']),
        'convergence_bear': bool(last['convergence_bear']),
        'large_buy': bool(last['large_buy']),
        'large_sell': bool(last['large_sell']),
        'vpin_proxy': float(last['vpin_proxy']),
        'funding_rate': float(last['funding_rate']),
        'fr_zscore': float(last['fr_zscore']),
        'fr_extreme_long': bool(last['fr_extreme_long']),
        'fr_extreme_short': bool(last['fr_extreme_short']),
        'oi_change': float(last['oi_change']),
        'oi_price_align': bool(last['oi_price_align']),
        'bullish_score': float(last['bullish_score']),
        'bearish_score': float(last['bearish_score']),
        'net_score': float(last['net_score']),
        'micro_signal': str(last['micro_signal'])
    }


def get_micro_summary(df: pd.DataFrame, lookback: int = 10) -> Dict:
    """最近 N 根的微观因子统计"""
    recent = df.tail(lookback)
    
    return {
        'avg_delta': float(recent['delta'].mean()),
        'avg_ofi': float(recent['ofi'].mean()),
        'cvd_trend': 'UP' if recent['cvd'].iloc[-1] > recent['cvd'].iloc[0] else 'DOWN',
        'cvd_change': float(recent['cvd'].iloc[-1] - recent['cvd'].iloc[0]),
        'divergence_bull_count': int(recent['divergence_bull'].sum()),
        'divergence_bear_count': int(recent['divergence_bear'].sum()),
        'large_buy_count': int(recent['large_buy'].sum()),
        'large_sell_count': int(recent['large_sell'].sum()),
        'avg_vpin': float(recent['vpin_proxy'].mean()),
        'funding_rate': float(recent['funding_rate'].iloc[-1]),
        'fr_regime': 'EXTREME_LONG' if recent['fr_extreme_long'].iloc[-1] else 
                     'EXTREME_SHORT' if recent['fr_extreme_short'].iloc[-1] else 'NORMAL',
        'net_score': float(recent['net_score'].iloc[-1]),
        'signal': str(recent['micro_signal'].iloc[-1]),
        'signal_distribution': recent['micro_signal'].value_counts().to_dict()
    }


if __name__ == "__main__":
    print("测试微观因子计算...")
    df = pd.read_parquet('/root/quant_pipeline/data/btc_15m_with_fr_oi.parquet')
    print(f"原始数据: {df.shape}")
    
    # 计算因子
    df = compute_microstructure_factors(df)
    print(f"计算完成: {df.shape}")
    print(f"新增列: {[c for c in df.columns if c not in ['o','h','l','c','v','close_time','quote_volume','count','taker_buy_v','taker_buy_quote','ignore','taker_sell_v','taker_sell_quote','taker_buy_ratio','timestamp','last_funding_rate']]}")
    
    # 最新因子
    latest = get_latest_micro_factors(df)
    print(f"\n最新微观因子:")
    for k, v in latest.items():
        print(f"  {k}: {v}")
    
    # 最近 10 根摘要
    summary = get_micro_summary(df, 10)
    print(f"\n最近 10 根摘要:")
    for k, v in summary.items():
        print(f"  {k}: {v}")
    
    # 保存带因子的数据
    out_file = '/root/quant_pipeline/data/btc_15m_with_micro.parquet'
    df.to_parquet(out_file)
    print(f"\n已保存至: {out_file}")