"""
价格结构分析：15分钟BTC的每一个转折点
不预设指标，从价格本身找规律
"""

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

df = pd.read_parquet("/root/quant_pipeline/data/btc_15m_binance.parquet")
df['time'] = pd.to_datetime(df['ts'])
df = df.set_index('time').sort_index()

o = df['open'].values.astype(float)
h = df['high'].values.astype(float)
l = df['low'].values.astype(float)
c = df['close'].values.astype(float)
v = df['volume'].values.astype(float)
N = len(df)

print(f"📊 数据: {N:,}根15分K线  {df.index[0]} ~ {df.index[-1]}")
print(f"   约{len(df)//96}个交易日\n")

# ═══ 1. 找转折点（swing highs / swing lows）═══
def find_swings(high, low, lookback=5):
    """
    找局部高点和低点
    一个点是swing high：它比左右各lookback根K线都高
    一个点是swing low：它比左右各lookback根K线都低
    """
    N = len(high)
    swing_high = np.zeros(N, dtype=bool)
    swing_low = np.zeros(N, dtype=bool)
    
    for i in range(lookback, N - lookback):
        # swing high
        if all(high[i] >= high[i-lookback:i]) and all(high[i] >= high[i+1:i+lookback+1]):
            swing_high[i] = True
        # swing low
        if all(low[i] <= low[i-lookback:i]) and all(low[i] <= low[i+1:i+lookback+1]):
            swing_low[i] = True
    
    return swing_high, swing_low

swing_high, swing_low = find_swings(h, l, lookback=5)

print(f"找到 {swing_high.sum()} 个局部高点, {swing_low.sum()} 个局部低点\n")

# ═══ 2. 量化每一段 "坡度" ═══
# 从低点→高点 = 上涨段，从高点→低点 = 下跌段

def analyze_swing_segments(c, swing_high, swing_low):
    """
    提取每个摆动段的结构特征
    返回每个完成段的特征
    """
    segments = []
    
    # 找到所有转折点的索引（按时间排序）
    turning_points = sorted(set(
        list(np.where(swing_high)[0]) + list(np.where(swing_low)[0])
    ))
    
    if len(turning_points) < 2:
        return segments
    
    for i in range(1, len(turning_points)):
        start = turning_points[i-1]
        end = turning_points[i]
        
        if end - start < 3:  # 太短的段跳过
            continue
        
        start_type = "high" if swing_high[start] else "low"
        end_type = "high" if swing_high[end] else "low"
        
        # 方向
        if start_type == "low" and end_type == "high":
            direction = "up"
        elif start_type == "high" and end_type == "low":
            direction = "down"
        else:
            continue  # 同向的转折点跳过
        
        # 价格变化
        if direction == "up":
            move_pct = (c[end] / c[start] - 1) * 100
            high_in_move = np.max(c[start:end+1])
            low_in_move = np.min(c[start:end+1])
        else:
            move_pct = (c[end] / c[start] - 1) * 100  # 负数
            high_in_move = np.max(c[start:end+1])
            low_in_move = np.min(c[start:end+1])
        
        bars = end - start
        hours = bars * 0.25
        
        # 坡度 = 每小时的变动百分比
        slope = move_pct / hours if hours > 0 else 0
        
        # 成交量
        vol_sum = np.sum(v[start:end+1])
        vol_per_bar = vol_sum / bars
        
        # 前一段的信息（如果存在）
        if i >= 2:
            prev_start = turning_points[i-2]
            prev_end = turning_points[i-1]
            prev_move = (c[prev_end] / c[prev_start] - 1) * 100
        else:
            prev_move = 0
        
        segments.append({
            "idx": i,
            "start_idx": start,
            "end_idx": end,
            "start_time": df.index[start],
            "end_time": df.index[end],
            "direction": direction,
            "move_pct": move_pct,
            "high_in_move": high_in_move,
            "low_in_move": low_in_move,
            "bars": bars,
            "hours": hours,
            "slope_pct_per_hour": slope,
            "vol_sum": vol_sum,
            "vol_per_bar": vol_per_bar,
            "prev_move": prev_move,
        })
    
    return segments

segments = analyze_swing_segments(c, swing_high, swing_low)
print(f"提取 {len(segments)} 个完整摆动段\n")

# 转为df
seg_df = pd.DataFrame(segments)

# ═══ 3. 统计：一段大涨幅/跌幅结束后的市场行为 ═══
print("="*80)
print("📊 1. 每段摆动的基本统计")
print("="*80)

up_segs = seg_df[seg_df['direction'] == 'up']
down_segs = seg_df[seg_df['direction'] == 'down']

print(f"\n🔴 上涨段: {len(up_segs)}段")
print(f"   平均涨幅: {up_segs['move_pct'].mean():+.2f}%")
print(f"   中位涨幅: {up_segs['move_pct'].median():+.2f}%")
print(f"   最大涨幅: {up_segs['move_pct'].max():+.2f}%")
print(f"   平均耗时: {up_segs['hours'].mean():.1f}小时")
print(f"   平均坡度: {up_segs['slope_pct_per_hour'].mean():+.2f}%/小时")

print(f"\n🔵 下跌段: {len(down_segs)}段")
print(f"   平均跌幅: {down_segs['move_pct'].mean():+.2f}%")
print(f"   中位跌幅: {down_segs['move_pct'].median():+.2f}%")
print(f"   最大跌幅: {down_segs['move_pct'].min():+.2f}%")
print(f"   平均耗时: {down_segs['hours'].mean():.1f}小时")
print(f"   平均坡度: {down_segs['slope_pct_per_hour'].mean():+.2f}%/小时")

# ═══ 4. 关键问题：大涨之后会怎样？大跌之后会怎样？ ═══
print("\n" + "="*80)
print("📊 2. 大涨/大跌后, 市场怎么走？")
print("="*80)

for threshold in [1.0, 2.0, 3.0, 5.0]:
    # 大涨段
    big_up = up_segs[up_segs['move_pct'] > threshold]
    if len(big_up) < 5:
        continue
    
    print(f"\n📈 大涨>+{threshold:.0f}%的段 ({len(big_up)}段):")
    
    # 看后续一段
    next_moves = []
    for _, seg in big_up.iterrows():
        next_idx = seg['idx'] + 1
        if next_idx < len(segments):
            next_seg = segments[next_idx]
            next_moves.append(next_seg['move_pct'])
    
    if next_moves:
        next_arr = np.array(next_moves)
        print(f"   下一段的平均收益: {np.mean(next_arr):+.2f}%")
        print(f"   下一段继续上涨的概率: {(next_arr > 0).mean()*100:.0f}%")
        print(f"   下一段下跌的概率: {(next_arr < 0).mean()*100:.0f}%")
        print(f"   下一段继续大涨(>+1%)概率: {(next_arr > 1).mean()*100:.0f}%")
        print(f"   下一段大跌(<-1%)概率: {(next_arr < -1).mean()*100:.0f}%")
    
    # 大跌段
    big_down = down_segs[down_segs['move_pct'] < -threshold]
    if len(big_down) < 5:
        continue
    
    print(f"\n📉 大跌<{-threshold:.0f}%的段 ({len(big_down)}段):")
    
    next_moves = []
    for _, seg in big_down.iterrows():
        next_idx = seg['idx'] + 1
        if next_idx < len(segments):
            next_seg = segments[next_idx]
            next_moves.append(next_seg['move_pct'])
    
    if next_moves:
        next_arr = np.array(next_moves)
        print(f"   下一段的平均收益: {np.mean(next_arr):+.2f}%")
        print(f"   下一段反弹的概率: {(next_arr > 0).mean()*100:.0f}%")
        print(f"   下一段继续跌的概率: {(next_arr < 0).mean()*100:.0f}%")
        print(f"   下一段大幅反弹(>+1%)概率: {(next_arr > 1).mean()*100:.0f}%")
        print(f"   下一段继续大跌(<-1%)概率: {(next_arr < -1).mean()*100:.0f}%")


# ═══ 5. 坡度分析：急涨/急跌后的行为 ═══
print("\n" + "="*80)
print("📊 3. 坡度分析 —— 急涨/急跌后怎么走？")
print("="*80)

# 将坡度分5组
for direction, segs in [("上涨", up_segs), ("下跌", down_segs)]:
    if len(segs) < 10:
        continue
    
    slopes = segs['slope_pct_per_hour'].values
    move_pcts = segs['move_pct'].values
    
    # 按坡度分5等份
    bins = np.percentile(slopes, [20, 40, 60, 80])
    
    print(f"\n{direction}段按坡度分组:")
    print(f"   {'坡度范围':<20} {'段数':>5} {'平均涨幅':>10} {'转向后下一段平均':>18}")
    print(f"   {'-'*55}")
    
    for b_idx in range(5):
        if b_idx == 0:
            mask = slopes <= bins[0]
            label = f"最慢(≤{bins[0]:+.2f})"
        elif b_idx == 4:
            mask = slopes > bins[3]
            label = f"最快(>{bins[3]:+.2f})"
        else:
            mask = (slopes > bins[b_idx-1]) & (slopes <= bins[b_idx])
            label = f"{bins[b_idx-1]:+.2f}~{bins[b_idx]:+.2f}"
        
        cnt = mask.sum()
        if cnt < 3:
            continue
        
        avg_move = np.mean(move_pcts[mask])
        
        # 看后续行为
        next_moves = []
        for idx in segs[mask].index:
            seg_row = segs.loc[idx]
            next_idx = seg_row['idx'] + 1
            if next_idx < len(segments):
                next_moves.append(segments[next_idx]['move_pct'])
        
        if next_moves:
            next_avg = np.mean(next_moves)
            next_wr = (np.array(next_moves) > 0).mean() * 100
            print(f"   {label:<20} {cnt:>5} {avg_move:>+9.2f}% {next_avg:>+10.2f}% (继续方向下一段胜率{next_wr:.0f}%)")
        else:
            print(f"   {label:<20} {cnt:>5} {avg_move:>+9.2f}%")


# ═══ 6. 看看是否有某种特定的"坡度+幅度"组合能预测 ═══
print("\n" + "="*80)
print("📊 4. 组合特征：坡度 + 幅度 + 前一段方向 → 下一段方向")
print("="*80)

# 将每一段和前后关联起来
from collections import defaultdict

pattern_stats = defaultdict(list)

for i in range(2, len(segments) - 1):
    prev = segments[i-1]
    curr = segments[i]
    next_seg = segments[i+1]
    
    # 前一段的方向和大小
    prev_dir = "U" if prev['move_pct'] > 0 else "D"
    prev_mag = "big" if abs(prev['move_pct']) > 2 else ("mid" if abs(prev['move_pct']) > 1 else "small")
    
    # 当前段的方向和大小
    curr_dir = "U" if curr['move_pct'] > 0 else "D"
    curr_mag = "big" if abs(curr['move_pct']) > 2 else ("mid" if abs(curr['move_pct']) > 1 else "small")
    
    # 坡度分类
    curr_slope = "steep" if abs(curr['slope_pct_per_hour']) > 0.5 else ("mid" if abs(curr['slope_pct_per_hour']) > 0.2 else "slow")
    
    # 下一段的方向
    next_dir = "U" if next_seg['move_pct'] > 0 else "D"
    next_move = next_seg['move_pct']
    
    # 模式
    pattern = f"{prev_dir}-{prev_mag} → {curr_dir}-{curr_mag}({curr_slope})"
    pattern_stats[pattern].append((next_dir, next_move))

# 统计每种模式的下文
print(f"\n{'模式':<40} {'出现次数':>6} {'续涨率':>8} {'续跌率':>8} {'平均下一段幅度':>12}")
print(f"{'-'*75}")

valid_patterns = [(p, v) for p, v in pattern_stats.items() if len(v) >= 5]
valid_patterns.sort(key=lambda x: len(x[1]), reverse=True)

for pattern, outcomes in valid_patterns[:20]:
    total = len(outcomes)
    up_rate = sum(1 for d, m in outcomes if d == "U") / total * 100
    down_rate = sum(1 for d, m in outcomes if d == "D") / total * 100
    avg_next = np.mean([m for d, m in outcomes])
    print(f"  {pattern:<40} {total:>6} {up_rate:>7.0f}% {down_rate:>7.0f}% {avg_next:>+11.3f}%")

# 特别关注：大跌之后的反弹
print("\n" + "="*80)
print("📊 5. 大跌后(<-2%)的第一段反弹特征")
print("="*80)

crash_indices = []
for i, seg in enumerate(segments):
    if seg['direction'] == 'down' and seg['move_pct'] < -2:
        if i + 1 < len(segments):
            crash_indices.append(i)

print(f"跌幅>2%后: {len(crash_indices)}次")

bounce_segs = []
for idx in crash_indices:
    bounce = segments[idx + 1]
    crash = segments[idx]
    bounce_segs.append({
        "crash_pct": crash['move_pct'],
        "crash_hours": crash['hours'],
        "crash_vol": crash['vol_per_bar'],
        "bounce_pct": bounce['move_pct'],
        "bounce_hours": bounce['hours'],
        "bounce_slope": bounce['slope_pct_per_hour'],
        "bounce_direction": bounce['direction'],
    })

bounce_df = pd.DataFrame(bounce_segs)
total_bounces = len(bounce_df)
up_bounces = (bounce_df['bounce_direction'] == 'up').sum()
print(f"后续第一段反弹上涨: {up_bounces}/{total_bounces} = {up_bounces/total_bounces*100:.0f}%")
print(f"反弹平均幅度: {bounce_df['bounce_pct'].mean():+.2f}%")
print(f"反弹中位幅度: {bounce_df['bounce_pct'].median():+.2f}%")
print(f"反弹平均耗时: {bounce_df['bounce_hours'].mean():.1f}小时")

# 大跌后直接反弹（反转）vs 继续下跌（延续）
reversal = bounce_df[bounce_df['bounce_direction'] == 'up']
continuation = bounce_df[bounce_df['bounce_direction'] == 'down']

print(f"\n  反转(反弹+): {len(reversal)}次 平均{bounce_df.loc[reversal.index, 'bounce_pct'].mean():+.2f}%")
print(f"  延续(继续-): {len(continuation)}次 平均{bounce_df.loc[continuation.index, 'bounce_pct'].mean():+.2f}%")

# 看大跌的不同特征是否影响反弹概率
print(f"\n  🔍 大跌特征vs反弹概率:")
# 按跌幅分
for threshold in [2, 3, 4, 5]:
    sub = bounce_df[bounce_df['crash_pct'] < -threshold]
    if len(sub) < 3:
        continue
    sub_up = (sub['bounce_direction'] == 'up').sum()
    print(f"    大跌>{threshold}% ({len(sub)}次): 反弹率{sub_up/len(sub)*100:.0f}%, 平均反弹{sub['bounce_pct'].mean():+.2f}%")

# 按速度分
median_hours = bounce_df['crash_hours'].median()
fast_crash = bounce_df[bounce_df['crash_hours'] <= median_hours]
slow_crash = bounce_df[bounce_df['crash_hours'] > median_hours]
print(f"    急跌(≤{median_hours:.0f}h) ({len(fast_crash)}次): 反弹率{(fast_crash['bounce_direction']=='up').mean()*100:.0f}%")
print(f"    慢跌(>{median_hours:.0f}h) ({len(slow_crash)}次): 反弹率{(slow_crash['bounce_direction']=='up').mean()*100:.0f}%")