#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
生成2024-2025年市场状态识别完整图表
"""

import sys
sys.path.insert(0, '/root/.openclaw/workspace/market-regime-identifier')

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import matplotlib.dates as mdates
from cyb50_market_classifier import fetch_cyb50_data, calculate_features, define_market_regime
import pickle
import warnings
warnings.filterwarnings('ignore')

# 设置中文字体
plt.rcParams['font.sans-serif'] = ['DejaVu Sans']
plt.rcParams['axes.unicode_minus'] = False

print("="*70)
print("生成2024年至今市场状态识别图表")
print("="*70)

# 获取数据
df = fetch_cyb50_data('2024-01-01', '2026-03-06')
if df is None:
    exit(1)

print(f"\n数据范围: {df.index[0].date()} ~ {df.index[-1].date()}")

# 计算特征和标签
features = calculate_features(df)
labels = define_market_regime(df, lookback=10)

# 训练模型
valid_idx = ~np.isnan(labels)
X = features[valid_idx]
y = labels[valid_idx]

from sklearn.ensemble import RandomForestClassifier
clf = RandomForestClassifier(
    n_estimators=100,
    max_depth=10,
    min_samples_split=20,
    min_samples_leaf=10,
    random_state=42,
    class_weight='balanced'
)
clf.fit(X, y)

# 预测所有数据
states = clf.predict(X)
probs = clf.predict_proba(X)

# 对齐数据
df_aligned = df.iloc[-len(states):].copy()
df_aligned['state'] = states
df_aligned['state_prob'] = [p[s] for s, p in zip(states, probs)]
df_aligned['prob_ranging'] = probs[:, 0]  # 震荡概率
df_aligned['prob_trend'] = probs[:, 1]    # 趋势概率
df_aligned['prob_reversal'] = probs[:, 2] # 反转概率

# 生成图表
fig, axes = plt.subplots(3, 1, figsize=(16, 12))

state_names = ['Ranging', 'Trend', 'Reversal']
colors = ['#2196F3', '#4CAF50', '#FF5722']  # 蓝、绿、橙

# 图1: 价格走势 + 状态标记
ax1 = axes[0]
for i, (name, color) in enumerate(zip(state_names, colors)):
    mask = df_aligned['state'] == i
    if mask.any():
        ax1.scatter(df_aligned.index[mask], df_aligned['close'][mask], 
                   c=color, label=name, alpha=0.7, s=30)

ax1.plot(df_aligned.index, df_aligned['close'], 'k-', alpha=0.3, linewidth=0.5)
ax1.set_ylabel('Price', fontsize=12)
ax1.set_title('CYB50 Market Regime Identification 2024-2025', fontsize=14, fontweight='bold')
ax1.legend(loc='upper left')
ax1.grid(True, alpha=0.3)

# 添加关键点位标注
for idx, row in df_aligned.iterrows():
    if idx.month == 1 and idx.day == 2:  # 年初
        ax1.annotate(f'{row["close"]:.0f}', 
                    xy=(idx, row['close']), 
                    xytext=(10, 10), textcoords='offset points',
                    fontsize=8, alpha=0.7)

# 图2: 状态概率时间序列
ax2 = axes[1]
ax2.fill_between(df_aligned.index, 0, df_aligned['prob_ranging'], 
                 alpha=0.5, label='Ranging', color=colors[0])
ax2.fill_between(df_aligned.index, df_aligned['prob_ranging'], 
                 df_aligned['prob_ranging'] + df_aligned['prob_trend'],
                 alpha=0.5, label='Trend', color=colors[1])
ax2.fill_between(df_aligned.index, 
                 df_aligned['prob_ranging'] + df_aligned['prob_trend'], 1,
                 alpha=0.5, label='Reversal', color=colors[2])

ax2.set_ylabel('Probability', fontsize=12)
ax2.set_title('State Probability Over Time', fontsize=12)
ax2.legend(loc='upper left')
ax2.grid(True, alpha=0.3)
ax2.set_ylim(0, 1)

# 图3: 状态分布统计
ax3 = axes[2]
state_counts = df_aligned['state'].value_counts().sort_index()
bars = ax3.bar(range(3), state_counts.values, color=colors, alpha=0.7)
ax3.set_xticks(range(3))
ax3.set_xticklabels(state_names)
ax3.set_ylabel('Days', fontsize=12)
ax3.set_title('State Distribution 2024-2025', fontsize=12)

# 添加数值标签
for i, (bar, count) in enumerate(zip(bars, state_counts.values)):
    pct = count / len(df_aligned) * 100
    ax3.text(bar.get_x() + bar.get_width()/2, bar.get_height() + 5,
             f'{count}d\n({pct:.1f}%)', 
             ha='center', va='bottom', fontsize=10)

plt.tight_layout()
plt.savefig('/root/.openclaw/workspace/market-regime-identifier/cyb50_regime_2024_2025.png', 
            dpi=150, bbox_inches='tight')
print("\n✓ 图表已保存: cyb50_regime_2024_2025.png")

# 生成详细报告
print("\n" + "="*70)
print("2024-2025年详细识别结果")
print("="*70)

# 按月份统计
print("\n【月度统计】")
print(f"{'月份':<10} {'总天数':<8} {'震荡':<8} {'趋势':<8} {'反转':<8} {'主要状态':<10}")
print("-"*70)

for year in [2024, 2025]:
    for month in range(1, 13):
        mask = (df_aligned.index.year == year) & (df_aligned.index.month == month)
        if not mask.any():
            continue
        
        month_data = df_aligned[mask]
        total = len(month_data)
        ranging = (month_data['state'] == 0).sum()
        trend = (month_data['state'] == 1).sum()
        reversal = (month_data['state'] == 2).sum()
        
        main_state = state_names[month_data['state'].mode()[0]]
        
        print(f"{year}-{month:02d}    {total:<8} {ranging:<8} {trend:<8} {reversal:<8} {main_state:<10}")

# 关键点位
print("\n【关键点位标注】")
print(f"{'日期':<12} {'收盘价':<10} {'状态':<10} {'置信度':<10} {'说明':<20}")
print("-"*70)

# 每月第一个交易日
for year in [2024, 2025]:
    for month in range(1, 13):
        mask = (df_aligned.index.year == year) & (df_aligned.index.month == month)
        if not mask.any():
            continue
        month_data = df_aligned[mask]
        first_day = month_data.iloc[0]
        
        date_str = month_data.index[0].strftime('%Y-%m-%d')
        price = first_day['close']
        state = state_names[int(first_day['state'])]
        prob = first_day['state_prob']
        
        # 简单说明
        if first_day['state'] == 0:
            desc = 'Consolidation'
        elif first_day['state'] == 1:
            if month_data['close'].iloc[-1] > price:
                desc = 'Uptrend'
            else:
                desc = 'Downtrend'
        else:
            desc = 'Reversal'
        
        print(f"{date_str:<12} {price:<10.2f} {state:<10} {prob:<10.2%} {desc:<20}")

print("\n" + "="*70)
print("✓ 报告生成完成!")
print("="*70)