#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
历史异动回测脚本

使用训练好的模型，对历史数据进行异动检测，生成异动记录

使用方法：
    # 回测指定日期范围
    python backtest.py --start 2024-01-01 --end 2024-12-01

    # 回测单天
    python backtest.py --start 2024-11-01 --end 2024-11-01

    # 只生成结果，不写入数据库
    python backtest.py --start 2024-01-01 --dry-run
"""

import os
import sys
import argparse
import json
from datetime import datetime, timedelta
from pathlib import Path
from typing import Dict, List, Tuple, Optional
from collections import defaultdict

import numpy as np
import pandas as pd
import torch
from tqdm import tqdm
from sqlalchemy import create_engine, text

# 添加父目录到路径
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from model import TransformerAutoencoder


# ==================== 配置 ====================

MYSQL_ENGINE = create_engine(
    "mysql+pymysql://root:Zzl5588161!@192.168.1.5:3306/stock",
    echo=False
)

# 特征列表（与训练一致）
FEATURES = [
    'alpha',
    'alpha_delta',
    'amt_ratio',
    'amt_delta',
    'rank_pct',
    'limit_up_ratio',
]

# 回测配置
BACKTEST_CONFIG = {
    'seq_len': 30,              # 序列长度
    'threshold_key': 'p95',     # 使用的阈值
    'min_alpha_abs': 0.5,       # 最小 Alpha 绝对值（过滤微小波动）
    'cooldown_minutes': 8,      # 同一概念冷却时间
    'max_alerts_per_minute': 15,  # 每分钟最多异动数
    'clip_value': 10.0,         # 极端值截断
}


# ==================== 模型加载 ====================

class AnomalyDetector:
    """异动检测器"""

    def __init__(self, checkpoint_dir: str = 'ml/checkpoints', device: str = 'auto'):
        self.checkpoint_dir = Path(checkpoint_dir)

        # 设备
        if device == 'auto':
            self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
        else:
            self.device = torch.device(device)

        # 加载配置
        self._load_config()

        # 加载模型
        self._load_model()

        # 加载阈值
        self._load_thresholds()

        print(f"AnomalyDetector 初始化完成")
        print(f"  设备: {self.device}")
        print(f"  阈值 ({BACKTEST_CONFIG['threshold_key']}): {self.threshold:.6f}")

    def _load_config(self):
        config_path = self.checkpoint_dir / 'config.json'
        with open(config_path, 'r') as f:
            self.config = json.load(f)

    def _load_model(self):
        model_path = self.checkpoint_dir / 'best_model.pt'
        checkpoint = torch.load(model_path, map_location=self.device)

        model_config = self.config['model'].copy()
        model_config['use_instance_norm'] = self.config.get('use_instance_norm', True)

        self.model = TransformerAutoencoder(**model_config)
        self.model.load_state_dict(checkpoint['model_state_dict'])
        self.model.to(self.device)
        self.model.eval()

    def _load_thresholds(self):
        thresholds_path = self.checkpoint_dir / 'thresholds.json'
        with open(thresholds_path, 'r') as f:
            thresholds = json.load(f)

        self.threshold = thresholds[BACKTEST_CONFIG['threshold_key']]

    @torch.no_grad()
    def compute_anomaly_scores(self, sequences: np.ndarray) -> np.ndarray:
        """
        计算异动分数

        Args:
            sequences: (n_sequences, seq_len, n_features)
        Returns:
            scores: (n_sequences,) 每个序列最后时刻的异动分数
        """
        # 截断极端值
        sequences = np.clip(sequences, -BACKTEST_CONFIG['clip_value'], BACKTEST_CONFIG['clip_value'])

        # 转为 tensor
        x = torch.FloatTensor(sequences).to(self.device)

        # 计算重构误差
        errors = self.model.compute_reconstruction_error(x, reduction='none')

        # 取最后一个时刻的误差
        scores = errors[:, -1].cpu().numpy()

        return scores

    def is_anomaly(self, score: float) -> bool:
        """判断是否异动"""
        return score > self.threshold


# ==================== 数据加载 ====================

def load_daily_features(data_dir: str, date: str) -> Optional[pd.DataFrame]:
    """加载单天的特征数据"""
    file_path = Path(data_dir) / f"features_{date}.parquet"

    if not file_path.exists():
        return None

    df = pd.read_parquet(file_path)
    return df


def get_available_dates(data_dir: str, start_date: str, end_date: str) -> List[str]:
    """获取可用的日期列表"""
    data_path = Path(data_dir)
    all_files = sorted(data_path.glob("features_*.parquet"))

    dates = []
    for f in all_files:
        date = f.stem.replace('features_', '')
        if start_date <= date <= end_date:
            dates.append(date)

    return dates


# ==================== 回测逻辑 ====================

def backtest_single_day(
    detector: AnomalyDetector,
    df: pd.DataFrame,
    date: str,
    seq_len: int = 30
) -> List[Dict]:
    """
    回测单天数据

    Args:
        detector: 异动检测器
        df: 当天的特征数据
        date: 日期
        seq_len: 序列长度

    Returns:
        alerts: 异动列表
    """
    alerts = []

    # 按概念分组
    grouped = df.groupby('concept_id', sort=False)

    # 冷却记录 {concept_id: last_alert_timestamp}
    cooldown = {}

    # 获取所有时间点
    all_timestamps = sorted(df['timestamp'].unique())

    if len(all_timestamps) < seq_len:
        return alerts

    # 对每个时间点进行检测（从第 seq_len 个开始）
    for t_idx in range(seq_len - 1, len(all_timestamps)):
        current_time = all_timestamps[t_idx]
        window_start_time = all_timestamps[t_idx - seq_len + 1]

        minute_alerts = []

        # 收集该时刻所有概念的序列
        concept_sequences = []
        concept_infos = []

        for concept_id, concept_df in grouped:
            # 获取该概念在时间窗口内的数据
            mask = (concept_df['timestamp'] >= window_start_time) & (concept_df['timestamp'] <= current_time)
            window_df = concept_df[mask].sort_values('timestamp')

            if len(window_df) < seq_len:
                continue

            # 取最后 seq_len 个点
            window_df = window_df.tail(seq_len)

            # 提取特征
            features = window_df[FEATURES].values

            # 处理缺失值
            features = np.nan_to_num(features, nan=0.0, posinf=0.0, neginf=0.0)

            # 获取当前时刻的信息
            current_row = window_df.iloc[-1]

            concept_sequences.append(features)
            concept_infos.append({
                'concept_id': concept_id,
                'timestamp': current_time,
                'alpha': current_row.get('alpha', 0),
                'alpha_delta': current_row.get('alpha_delta', 0),
                'amt_ratio': current_row.get('amt_ratio', 1),
                'limit_up_ratio': current_row.get('limit_up_ratio', 0),
                'limit_down_ratio': current_row.get('limit_down_ratio', 0),
                'rank_pct': current_row.get('rank_pct', 0.5),
                'stock_count': current_row.get('stock_count', 0),
                'total_amt': current_row.get('total_amt', 0),
            })

        if not concept_sequences:
            continue

        # 批量计算异动分数
        sequences_array = np.array(concept_sequences)
        scores = detector.compute_anomaly_scores(sequences_array)

        # 检测异动
        for i, (info, score) in enumerate(zip(concept_infos, scores)):
            concept_id = info['concept_id']
            alpha = info['alpha']

            # 过滤小波动
            if abs(alpha) < BACKTEST_CONFIG['min_alpha_abs']:
                continue

            # 检查冷却
            if concept_id in cooldown:
                last_alert = cooldown[concept_id]
                if isinstance(current_time, datetime):
                    time_diff = (current_time - last_alert).total_seconds() / 60
                else:
                    # timestamp 是字符串或其他格式
                    time_diff = BACKTEST_CONFIG['cooldown_minutes'] + 1  # 跳过冷却检查

                if time_diff < BACKTEST_CONFIG['cooldown_minutes']:
                    continue

            # 判断是否异动
            if not detector.is_anomaly(score):
                continue

            # 记录异动
            alert_type = 'surge_up' if alpha > 0 else 'surge_down'

            alert = {
                'concept_id': concept_id,
                'alert_time': current_time,
                'trade_date': date,
                'alert_type': alert_type,
                'anomaly_score': float(score),
                'threshold': detector.threshold,
                **info
            }

            minute_alerts.append(alert)
            cooldown[concept_id] = current_time

        # 按分数排序，限制数量
        minute_alerts.sort(key=lambda x: x['anomaly_score'], reverse=True)
        alerts.extend(minute_alerts[:BACKTEST_CONFIG['max_alerts_per_minute']])

    return alerts


# ==================== 数据库写入 ====================

def save_alerts_to_mysql(alerts: List[Dict], dry_run: bool = False) -> int:
    """保存异动到 MySQL"""
    if not alerts:
        return 0

    if dry_run:
        print(f"  [Dry Run] 将写入 {len(alerts)} 条异动")
        return len(alerts)

    saved = 0
    with MYSQL_ENGINE.begin() as conn:
        for alert in alerts:
            try:
                # 检查是否已存在
                check_sql = text("""
                    SELECT id FROM concept_minute_alert
                    WHERE concept_id = :concept_id
                    AND alert_time = :alert_time
                    AND trade_date = :trade_date
                """)
                exists = conn.execute(check_sql, {
                    'concept_id': alert['concept_id'],
                    'alert_time': alert['alert_time'],
                    'trade_date': alert['trade_date'],
                }).fetchone()

                if exists:
                    continue

                # 插入新记录
                insert_sql = text("""
                    INSERT INTO concept_minute_alert
                    (concept_id, concept_name, alert_time, alert_type, trade_date,
                     change_pct, zscore, importance_score, stock_count, extra_info)
                    VALUES
                    (:concept_id, :concept_name, :alert_time, :alert_type, :trade_date,
                     :change_pct, :zscore, :importance_score, :stock_count, :extra_info)
                """)

                conn.execute(insert_sql, {
                    'concept_id': alert['concept_id'],
                    'concept_name': alert.get('concept_name', ''),
                    'alert_time': alert['alert_time'],
                    'alert_type': alert['alert_type'],
                    'trade_date': alert['trade_date'],
                    'change_pct': alert.get('alpha', 0),
                    'zscore': alert['anomaly_score'],
                    'importance_score': alert['anomaly_score'],
                    'stock_count': alert.get('stock_count', 0),
                    'extra_info': json.dumps({
                        'detection_method': 'ml_autoencoder',
                        'threshold': alert['threshold'],
                        'alpha': alert.get('alpha', 0),
                        'amt_ratio': alert.get('amt_ratio', 1),
                    }, ensure_ascii=False)
                })

                saved += 1

            except Exception as e:
                print(f"  保存失败: {alert['concept_id']} - {e}")

    return saved


def export_alerts_to_csv(alerts: List[Dict], output_path: str):
    """导出异动到 CSV"""
    if not alerts:
        return

    df = pd.DataFrame(alerts)
    df.to_csv(output_path, index=False, encoding='utf-8-sig')
    print(f"已导出到: {output_path}")


# ==================== 主函数 ====================

def main():
    parser = argparse.ArgumentParser(description='历史异动回测')
    parser.add_argument('--data_dir', type=str, default='ml/data',
                        help='特征数据目录')
    parser.add_argument('--checkpoint_dir', type=str, default='ml/checkpoints',
                        help='模型检查点目录')
    parser.add_argument('--start', type=str, required=True,
                        help='开始日期 (YYYY-MM-DD)')
    parser.add_argument('--end', type=str, required=True,
                        help='结束日期 (YYYY-MM-DD)')
    parser.add_argument('--dry-run', action='store_true',
                        help='只计算，不写入数据库')
    parser.add_argument('--export-csv', type=str, default=None,
                        help='导出 CSV 文件路径')
    parser.add_argument('--device', type=str, default='auto',
                        help='设备 (auto/cuda/cpu)')

    args = parser.parse_args()

    print("=" * 60)
    print("历史异动回测")
    print("=" * 60)
    print(f"日期范围: {args.start} ~ {args.end}")
    print(f"数据目录: {args.data_dir}")
    print(f"模型目录: {args.checkpoint_dir}")
    print(f"Dry Run: {args.dry_run}")
    print("=" * 60)

    # 初始化检测器
    detector = AnomalyDetector(args.checkpoint_dir, args.device)

    # 获取可用日期
    dates = get_available_dates(args.data_dir, args.start, args.end)

    if not dates:
        print(f"未找到 {args.start} ~ {args.end} 范围内的数据")
        return

    print(f"\n找到 {len(dates)} 天的数据")

    # 回测
    all_alerts = []
    total_saved = 0

    for date in tqdm(dates, desc="回测进度"):
        # 加载数据
        df = load_daily_features(args.data_dir, date)

        if df is None or df.empty:
            continue

        # 回测单天
        alerts = backtest_single_day(
            detector, df, date,
            seq_len=BACKTEST_CONFIG['seq_len']
        )

        if alerts:
            all_alerts.extend(alerts)

            # 写入数据库
            saved = save_alerts_to_mysql(alerts, dry_run=args.dry_run)
            total_saved += saved

            if not args.dry_run:
                tqdm.write(f"  {date}: 检测到 {len(alerts)} 个异动，保存 {saved} 条")

    # 导出 CSV
    if args.export_csv and all_alerts:
        export_alerts_to_csv(all_alerts, args.export_csv)

    # 汇总
    print("\n" + "=" * 60)
    print("回测完成！")
    print("=" * 60)
    print(f"总计检测到: {len(all_alerts)} 个异动")
    print(f"保存到数据库: {total_saved} 条")

    # 统计
    if all_alerts:
        df_alerts = pd.DataFrame(all_alerts)
        print(f"\n异动类型分布:")
        print(df_alerts['alert_type'].value_counts())

        print(f"\n异动分数统计:")
        print(f"  Mean: {df_alerts['anomaly_score'].mean():.4f}")
        print(f"  Max:  {df_alerts['anomaly_score'].max():.4f}")
        print(f"  Min:  {df_alerts['anomaly_score'].min():.4f}")

    print("=" * 60)


if __name__ == "__main__":
    main()