#!/usr/bin/env python3
"""
RAG 知识库初始化脚本
添加加密货币交易的专业知识到向量数据库
"""

import sys
import os

# 添加项目根目录到 Python 路径
sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))

from langchain_core.documents import Document
from app.services.rag_service import get_rag_service
import logging

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# 专业知识库内容
KNOWLEDGE_BASE = [
    {
        "title": "技术指标 - MACD（指数平滑异同移动平均线）",
        "content": """
## MACD 指标详解

### 基本原理
MACD 由 Gerald Appel 在 1970 年代提出，是趋势跟踪和动量指标的结合。

### 计算方法
1. **DIF（快线）** = EMA(12) - EMA(26)
2. **DEA（慢线）** = DIF 的 9 日 EMA
3. **MACD 柱状图** = (DIF - DEA) × 2

### 交易信号
**金叉买入信号**：
- DIF 上穿 DEA，且都在零轴上方 → 强烈买入
- DIF 上穿 DEA，但在零轴下方 → 谨慎买入

**死叉卖出信号**：
- DIF 下穿 DEA，且都在零轴下方 → 强烈卖出
- DIF 下穿 DEA，但在零轴上方 → 谨慎卖出

**背离信号**：
- 价格创新高，MACD 未创新高 → 顶背离，看跌
- 价格创新低，MACD 未创新低 → 底背离，看涨

### 最佳参数
- **日线级别**：12, 26, 9（标准）
- **4小时级别**：12, 26, 9
- **1小时级别**：6, 13, 5（灵敏）
- **15分钟级别**：6, 13, 5

### 实战策略
```python
# Freqtrade MACD 金叉策略
def populate_entry_trend(self, dataframe: DataFrame, metadata: dict) -> DataFrame:
    dataframe.loc[
        (
            (dataframe['macd'] > dataframe['macdsignal']) &
            (dataframe['macd'].shift(1) <= dataframe['macdsignal'].shift(1)) &
            (dataframe['volume'] > 0)
        ),
        'enter_long'] = 1
    return dataframe
```

### 注意事项
- ⚠️ 震荡市场容易产生假信号，建议结合趋势指标（如 MA）
- ⚠️ 加密货币波动大，建议调整参数或结合 RSI
- ✅ 配合成交量使用，成交量放大的信号更可靠
""",
        "category": "technical_analysis"
    },
    {
        "title": "技术指标 - RSI（相对强弱指数）",
        "content": """
## RSI 指标详解

### 基本原理
RSI 由 Welles Wilder 在 1978 年提出，衡量价格变动的速度和幅度，判断超买超卖。

### 计算方法
RSI = 100 - [100 / (1 + RS)]
其中 RS = n 日平均上涨幅度 / n 日平均下跌幅度

### 标准参数
- **常用周期**：14（标准）、6（快速）、28（慢速）
- **超买线**：70
- **超卖线**：30

### 交易信号
**超买超卖信号**：
- RSI > 70 → 超买，可能回调
- RSI < 30 → 超卖，可能反弹
- RSI > 80 → 极度超买（加密货币）
- RSI < 20 → 极度超卖（加密货币）

**背离信号（重要！）**：
- **顶背离**：价格创新高，RSI 未创新高 → 强烈卖出信号
- **底背离**：价格创新低，RSI 未创新低 → 强烈买入信号

**中线突破**：
- RSI 上穿 50 → 多头占优
- RSI 下破 50 → 空头占优

### 加密货币特殊用法
由于加密货币波动性大：
- 超买线可调整到 80
- 超卖线可调整到 20
- 周期可缩短到 9-12

### 实战策略
```python
# Freqtrade RSI 超卖反弹策略
def populate_entry_trend(self, dataframe: DataFrame, metadata: dict) -> DataFrame:
    dataframe.loc[
        (
            (dataframe['rsi'] < 30) &
            (dataframe['rsi'] > dataframe['rsi'].shift(1)) &  # RSI 开始回升
            (dataframe['volume'] > dataframe['volume'].rolling(20).mean())  # 成交量放大
        ),
        'enter_long'] = 1
    return dataframe

def populate_exit_trend(self, dataframe: DataFrame, metadata: dict) -> DataFrame:
    dataframe.loc[
        (
            (dataframe['rsi'] > 70) |  # 超买
            (dataframe['close'] < dataframe['close'].shift(1) * 0.98)  # 或跌破 2%
        ),
        'exit_long'] = 1
    return dataframe
```

### 组合使用建议
- **RSI + MACD**：RSI 超卖 + MACD 金叉 = 强买入
- **RSI + MA**：RSI 超卖 + 站上 MA20 = 确认反弹
- **RSI + 布林带**：价格触及下轨 + RSI < 30 = 超卖反弹
""",
        "category": "technical_analysis"
    },
    {
        "title": "K线形态 - 头肩顶与头肩底",
        "content": """
## 头肩顶/底形态识别

### 头肩顶（看跌反转形态）

**形态特征**：
1. **左肩**：价格上涨到高点后回落
2. **头部**：价格再次上涨，突破左肩高点，然后回落
3. **右肩**：价格第三次上涨，但未能突破头部高点，形成右肩
4. **颈线**：连接左肩和头部之间的低点与头部和右肩之间的低点

**确认信号**：
- ✅ 价格有效跌破颈线（收盘价跌破 3%）
- ✅ 成交量：左肩>头部>右肩，跌破颈线时放量
- ✅ 反抽颈线未能收复，形成二次确认

**目标价位**：
下跌目标 = 颈线位 - 头部高点到颈线的垂直距离

**实战要点**：
- 头肩顶形成时间越长，反转力度越强（一般 3-6 个月）
- 颈线倾斜度：水平或微微上倾更可靠
- 假突破：跌破颈线后 3 日内反抽未能收复，才确认

### 头肩底（看涨反转形态）

**形态特征**（与头肩顶相反）：
1. 价格下跌形成左肩
2. 继续下跌形成头部（最低点）
3. 反弹后再次下跌形成右肩，但不破头部低点
4. 突破颈线确认

**确认信号**：
- ✅ 价格有效突破颈线（收盘价突破 3%）
- ✅ 成交量：突破颈线时明显放大
- ✅ 回踩颈线支撑有效

**目标价位**：
上涨目标 = 颈线位 + 头部低点到颈线的垂直距离

### 交易策略
```python
# 头肩顶策略（简化版）
def identify_head_shoulders_top(dataframe):
    # 识别头肩顶
    left_shoulder = dataframe['high'].rolling(20).max()
    head = dataframe['high'].rolling(40).max()
    right_shoulder = dataframe['high'].rolling(20).max()
    
    # 颈线
    neckline = (dataframe['low'].rolling(20).min() + dataframe['low'].shift(20).rolling(20).min()) / 2
    
    # 跌破颈线 = 卖出信号
    sell_signal = dataframe['close'] < neckline * 0.97
    
    return sell_signal
```

### 失败形态（陷阱）
- ❌ 右肩高于头部 → 形态失败
- ❌ 跌破颈线后立即反弹收复 → 假突破
- ❌ 成交量反常（左肩最大） → 可靠性降低
""",
        "category": "chart_patterns"
    },
    {
        "title": "风险管理 - 仓位管理与资金分配",
        "content": """
## 专业仓位管理策略

### 核心原则
1. **永远不要满仓**：留有余力应对突发情况
2. **单笔亏损不超过账户 2%**：保护本金
3. **相关性分散**：不要在相关品种上重复押注

### 凯利公式（Kelly Criterion）
最优仓位 = (胜率 × 盈亏比 - 败率) / 盈亏比

**示例**：
- 胜率 60%
- 盈亏比 2:1（平均盈利是平均亏损的 2 倍）
- 最优仓位 = (0.6 × 2 - 0.4) / 2 = 0.4 = 40%

**保守修正**：实际仓位 = 凯利仓位 × 0.5（减半使用更安全）

### 固定比例法
每笔交易风险固定为账户的 1-2%

**计算公式**：
仓位 = 可承受亏损金额 / 单笔止损幅度

**示例**：
- 账户 10,000 USDT
- 单笔风险 2% = 200 USDT
- 止损幅度 5%
- 仓位 = 200 / 0.05 = 4,000 USDT（40%）

### 分批建仓法（推荐）
**三分法**：
1. 首次建仓：30%
2. 确认趋势：加仓 30%
3. 突破关键位：再加 40%

**金字塔加仓**：
- 首次：50%
- 第一次加仓：30%
- 第二次加仓：20%

**反金字塔加仓（激进）**：
- 首次：20%
- 第一次加仓：30%
- 第二次加仓：50%

### 不同市场环境的仓位策略
| 市场状态 | 总仓位 | 单币种仓位 |
|---------|--------|-----------|
| 强势牛市 | 70-80% | 20-30% |
| 震荡市场 | 30-50% | 10-15% |
| 弱势熊市 | 10-20% | 5-10% |
| 单边暴跌 | 0-10% | 空仓为主 |

### 止损设置
**技术止损**：
- 跌破关键支撑位（如 MA20、前低）
- 跌幅达到 5-8%（加密货币）
- K线形态破坏（如破颈线）

**时间止损**：
- 持仓 3 天无利润 → 平仓
- 持仓 7 天未达预期 → 减仓

**盈利保护止损**：
- 盈利 10%时，止损移至成本价
- 盈利 20%时，止损移至+5%
- 盈利 30%时，止损移至+10%

### 实战代码
```python
# Freqtrade 动态仓位管理
def custom_stake_amount(self, pair: str, current_time, current_rate,
                        proposed_stake, min_stake, max_stake,
                        entry_tag, side: str, **kwargs) -> float:
    # 获取账户余额
    wallet_balance = self.wallets.get_total(self.config['stake_currency'])
    
    # 计算止损距离
    stop_loss_distance = 0.05  # 5%
    
    # 固定风险 2%
    risk_per_trade = wallet_balance * 0.02
    
    # 计算仓位
    stake_amount = risk_per_trade / stop_loss_distance
    
    # 限制最大仓位 30%
    max_position = wallet_balance * 0.3
    stake_amount = min(stake_amount, max_position)
    
    return stake_amount
```

### 心理建设
- ✅ 亏损是交易的一部分，接受小亏损
- ✅ 不要因为连续亏损而加大仓位（赌徒心态）
- ✅ 不要因为连续盈利而放松警惕（过度自信）
- ❌ 永远不要 ALL IN
""",
        "category": "risk_management"
    },
    {
        "title": "交易策略 - 双均线策略（Golden Cross）",
        "content": """
## 双均线交叉策略

### 策略原理
利用不同周期的移动平均线交叉产生买卖信号，是最经典的趋势跟踪策略。

### 常用参数组合
**短线交易**：
- MA5 / MA10
- MA10 / MA20

**中线交易**：
- MA20 / MA60
- MA30 / MA90

**长线交易**：
- MA50 / MA200（金叉/死叉）
- MA100 / MA300

### 信号规则
**金叉（买入）**：
- 短期均线向上穿越长期均线
- 两条均线都向上倾斜（趋势确认）
- 价格在均线上方

**死叉（卖出）**：
- 短期均线向下穿越长期均线
- 两条均线都向下倾斜
- 价格在均线下方

### 优化方法
1. **加入成交量过滤**：
   - 金叉时成交量放大 > 20 日均量
   - 避免无量上涨的假突破

2. **加入 MACD 确认**：
   - 金叉 + MACD 金叉 = 双重确认
   - 提高胜率，但会错过部分行情

3. **动态止损**：
   - 价格跌破短期均线 → 平仓 50%
   - 价格跌破长期均线 → 全部平仓

4. **分批入场**：
   - 金叉时买入 50%
   - 价格回踩均线不破时加仓 30%
   - 突破前高时再加仓 20%

### Freqtrade 完整策略代码
```python
from freqtrade.strategy import IStrategy
from pandas import DataFrame
import talib.abstract as ta

class GoldenCrossStrategy(IStrategy):
    # 策略参数
    minimal_roi = {
        "0": 0.20,   # 20% 止盈
        "60": 0.10,  # 持仓1小时后 10% 止盈
        "120": 0.05  # 持仓2小时后 5% 止盈
    }
    
    stoploss = -0.08  # 8% 止损
    
    timeframe = '1h'
    
    # 策略指标
    def populate_indicators(self, dataframe: DataFrame, metadata: dict) -> DataFrame:
        # 计算移动平均线
        dataframe['ma_short'] = ta.SMA(dataframe, timeperiod=20)
        dataframe['ma_long'] = ta.SMA(dataframe, timeperiod=60)
        
        # 成交量均线
        dataframe['volume_ma'] = dataframe['volume'].rolling(20).mean()
        
        # MACD
        macd = ta.MACD(dataframe)
        dataframe['macd'] = macd['macd']
        dataframe['macdsignal'] = macd['macdsignal']
        
        return dataframe
    
    # 买入信号
    def populate_entry_trend(self, dataframe: DataFrame, metadata: dict) -> DataFrame:
        dataframe.loc[
            (
                # 金叉
                (dataframe['ma_short'] > dataframe['ma_long']) &
                (dataframe['ma_short'].shift(1) <= dataframe['ma_long'].shift(1)) &
                
                # 两条均线都向上
                (dataframe['ma_short'] > dataframe['ma_short'].shift(1)) &
                (dataframe['ma_long'] > dataframe['ma_long'].shift(1)) &
                
                # 成交量放大
                (dataframe['volume'] > dataframe['volume_ma'] * 1.2) &
                
                # MACD 确认
                (dataframe['macd'] > dataframe['macdsignal']) &
                
                # 确保有成交量
                (dataframe['volume'] > 0)
            ),
            'enter_long'] = 1
        
        return dataframe
    
    # 卖出信号
    def populate_exit_trend(self, dataframe: DataFrame, metadata: dict) -> DataFrame:
        dataframe.loc[
            (
                # 死叉
                (dataframe['ma_short'] < dataframe['ma_long']) &
                (dataframe['ma_short'].shift(1) >= dataframe['ma_long'].shift(1)) &
                
                # 或者价格跌破短期均线
                (dataframe['close'] < dataframe['ma_short'] * 0.97)
            ) |
            (
                # 或 MACD 死叉
                (dataframe['macd'] < dataframe['macdsignal']) &
                (dataframe['macd'].shift(1) >= dataframe['macdsignal'].shift(1))
            ),
            'exit_long'] = 1
        
        return dataframe
```

### 回测数据参考
基于 BTC/USDT 1小时级别（2023 年数据）：
- **胜率**：45%
- **盈亏比**：2.5:1
- **年化收益**：35%
- **最大回撤**：-15%

### 适用市场
- ✅ 趋势明显的牛市或熊市
- ✅ 波段较大的品种（如主流币）
- ❌ 震荡市场（容易被来回打止损）
- ❌ 低波动品种（信号少）

### 注意事项
1. **避免频繁交易**：震荡市会产生大量假信号，建议结合趋势指标（如 ADX）
2. **时间周期选择**：加密货币建议 1h 或 4h，避免 15m（噪音太多）
3. **滑点和手续费**：频繁交易会吃掉大部分利润，建议只在明确趋势时交易
""",
        "category": "trading_strategies"
    },
    {
        "title": "链上数据分析 - 交易所净流入/流出",
        "content": """
## 交易所资金流动分析

### 核心原理
- **净流入（Inflow）**：资金流入交易所 → 抛压增加 → 看跌
- **净流出（Outflow）**：资金流出交易所 → 长期持有 → 看涨

### 数据来源
- CryptoQuant
- Glassnode
- IntoTheBlock
- Nansen

### 关键指标

**1. 交易所余额变化**
- 余额持续下降 → 抛压减轻，利好
- 余额持续上升 → 抛压增加，利空
- 大额转入（>100 BTC）→ 可能即将砸盘

**2. 交易所储备比例**
- 低储备比例（<10%）→ 流动性紧张，利好
- 高储备比例（>15%）→ 抛压大，利空

**3. 矿工流向**
- 矿工向交易所转币 → 准备抛售，利空
- 矿工囤币 → 看好后市，利好

### 实战信号
**强烈看涨**：
- 交易所净流出 >5000 BTC/天
- 交易所余额创新低
- 巨鲸地址增持
- 稳定币流入交易所（准备抄底）

**强烈看跌**：
- 交易所净流入 >8000 BTC/天
- 矿工大量转入交易所
- 巨鲸地址减持
- 稳定币流出交易所

### 配合价格使用
| 链上信号 | 价格走势 | 判断 |
|---------|---------|-----|
| 净流出 | 上涨 | 健康上涨，可持有 |
| 净流出 | 下跌 | 底部吸筹，潜在反转 |
| 净流入 | 上涨 | 获利盘出逃，警惕回调 |
| 净流入 | 下跌 | 恐慌性抛售，超跌反弹 |

### 注意事项
- ⚠️ 单日数据波动大，建议看 7 日均值
- ⚠️ 大型交易所（Binance、Coinbase）的数据更有参考价值
- ⚠️ 结合价格和成交量综合判断
""",
        "category": "onchain_analysis"
    }
]


async def init_knowledge_base():
    """初始化知识库"""
    try:
        logger.info("开始初始化 RAG 知识库...")
        
        # 获取 RAG 服务
        rag_service = get_rag_service()
        
        # 转换为 RAG 服务期望的格式
        documents = []
        for item in KNOWLEDGE_BASE:
            doc = {
                "content": item["content"],
                "title": item["title"],
                "category": item["category"],
                "source": "official_knowledge_base"
            }
            documents.append(doc)
        
        logger.info(f"准备导入 {len(documents)} 条知识...")
        
        # 批量导入到默认集合
        collection_name = "trading_knowledge"
        await rag_service.ingest_documents(
            documents=documents,
            collection_name=collection_name
        )
        
        logger.info(f"✅ 知识库初始化完成！共导入 {len(documents)} 条专业知识")
        logger.info(f"集合名称: {collection_name}")
        
        # 测试检索
        logger.info("\n测试检索功能...")
        results = await rag_service.search(
            query="MACD 指标怎么用？",
            collection_name=collection_name,
            top_k=2
        )
        
        logger.info(f"检索到 {len(results)} 条相关结果：")
        for i, doc in enumerate(results, 1):
            title = doc.get('title', '未命名') if isinstance(doc, dict) else getattr(doc, 'metadata', {}).get('title', '未命名')
            logger.info(f"{i}. {title}")
        
        return True
        
    except Exception as e:
        logger.error(f"❌ 初始化知识库失败: {e}")
        import traceback
        traceback.print_exc()
        return False


if __name__ == "__main__":
    import asyncio
    success = asyncio.run(init_knowledge_base())
    sys.exit(0 if success else 1)