#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
训练脚本 V2 - 基于 Z-Score 特征的 LSTM Autoencoder

改进点：
1. 使用 Z-Score 特征（相对于同时间片历史的偏离）
2. 短序列：10分钟（不需要30分钟预热）
3. 开盘即可检测：9:30 直接有特征

模型输入：
- 过去10分钟的 Z-Score 特征序列
- 特征：alpha_zscore, amt_zscore, rank_zscore, momentum_3m, momentum_5m, limit_up_ratio

模型学习：
- 学习 Z-Score 序列的"正常演化模式"
- 异动 = Z-Score 序列的异常演化（重构误差大）
"""

import os
import sys
import argparse
import json
from datetime import datetime
from pathlib import Path
from typing import List, Tuple, Dict

import numpy as np
import pandas as pd
import torch
import torch.nn as nn
from torch.utils.data import Dataset, DataLoader
from torch.optim import AdamW
from torch.optim.lr_scheduler import CosineAnnealingWarmRestarts
from tqdm import tqdm

from model import TransformerAutoencoder, AnomalyDetectionLoss, count_parameters

# 性能优化
torch.backends.cudnn.benchmark = True
torch.backends.cuda.matmul.allow_tf32 = True
torch.backends.cudnn.allow_tf32 = True

try:
    import matplotlib
    matplotlib.use('Agg')
    import matplotlib.pyplot as plt
    HAS_MATPLOTLIB = True
except ImportError:
    HAS_MATPLOTLIB = False


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

TRAIN_CONFIG = {
    # 数据配置（改进！）
    'seq_len': 10,                    # 10分钟序列（不是30分钟！）
    'stride': 2,                      # 步长2分钟

    # 时间切分
    'train_end_date': '2024-06-30',
    'val_end_date': '2024-09-30',

    # V2 特征（Z-Score 为主）
    'features': [
        'alpha_zscore',       # Alpha 的 Z-Score
        'amt_zscore',         # 成交额的 Z-Score
        'rank_zscore',        # 排名的 Z-Score
        'momentum_3m',        # 3分钟动量
        'momentum_5m',        # 5分钟动量
        'limit_up_ratio',     # 涨停占比
    ],

    # 训练配置
    'batch_size': 4096,
    'epochs': 100,
    'learning_rate': 3e-4,
    'weight_decay': 1e-5,
    'gradient_clip': 1.0,

    # 早停配置
    'patience': 15,
    'min_delta': 1e-6,

    # 模型配置（小型 LSTM）
    'model': {
        'n_features': 6,
        'hidden_dim': 32,
        'latent_dim': 4,
        'num_layers': 1,
        'dropout': 0.2,
        'bidirectional': True,
    },

    # 标准化配置
    'clip_value': 5.0,  # Z-Score 已经标准化，clip 5.0 足够

    # 阈值配置
    'threshold_percentiles': [90, 95, 99],
}


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

def load_data_by_date(data_dir: str, features: List[str]) -> Dict[str, pd.DataFrame]:
    """按日期加载 V2 数据"""
    data_path = Path(data_dir)
    parquet_files = sorted(data_path.glob("features_v2_*.parquet"))

    if not parquet_files:
        raise FileNotFoundError(f"未找到 V2 数据文件: {data_dir}")

    print(f"找到 {len(parquet_files)} 个 V2 数据文件")

    date_data = {}

    for pf in tqdm(parquet_files, desc="加载数据"):
        date = pf.stem.replace('features_v2_', '')

        df = pd.read_parquet(pf)

        required_cols = features + ['concept_id', 'timestamp']
        missing_cols = [c for c in required_cols if c not in df.columns]
        if missing_cols:
            print(f"警告: {date} 缺少列: {missing_cols}, 跳过")
            continue

        date_data[date] = df

    print(f"成功加载 {len(date_data)} 天的数据")
    return date_data


def split_data_by_date(
    date_data: Dict[str, pd.DataFrame],
    train_end: str,
    val_end: str
) -> Tuple[Dict[str, pd.DataFrame], Dict[str, pd.DataFrame], Dict[str, pd.DataFrame]]:
    """按日期划分数据集"""
    train_data = {}
    val_data = {}
    test_data = {}

    for date, df in date_data.items():
        if date <= train_end:
            train_data[date] = df
        elif date <= val_end:
            val_data[date] = df
        else:
            test_data[date] = df

    print(f"数据集划分:")
    print(f"  训练集: {len(train_data)} 天 (<= {train_end})")
    print(f"  验证集: {len(val_data)} 天 ({train_end} ~ {val_end})")
    print(f"  测试集: {len(test_data)} 天 (> {val_end})")

    return train_data, val_data, test_data


def build_sequences_by_concept(
    date_data: Dict[str, pd.DataFrame],
    features: List[str],
    seq_len: int,
    stride: int
) -> np.ndarray:
    """按概念分组构建序列"""
    all_dfs = []
    for date, df in sorted(date_data.items()):
        df = df.copy()
        df['date'] = date
        all_dfs.append(df)

    if not all_dfs:
        return np.array([])

    combined = pd.concat(all_dfs, ignore_index=True)
    combined = combined.sort_values(['concept_id', 'date', 'timestamp'])

    all_sequences = []
    grouped = combined.groupby('concept_id', sort=False)
    n_concepts = len(grouped)

    for concept_id, concept_df in tqdm(grouped, desc="构建序列", total=n_concepts, leave=False):
        feature_data = concept_df[features].values
        feature_data = np.nan_to_num(feature_data, nan=0.0, posinf=0.0, neginf=0.0)

        n_points = len(feature_data)
        for start in range(0, n_points - seq_len + 1, stride):
            seq = feature_data[start:start + seq_len]
            all_sequences.append(seq)

    if not all_sequences:
        return np.array([])

    sequences = np.array(all_sequences)
    print(f"  构建序列: {len(sequences):,} 条 (来自 {n_concepts} 个概念)")

    return sequences


# ==================== 数据集 ====================

class SequenceDataset(Dataset):
    def __init__(self, sequences: np.ndarray):
        self.sequences = torch.FloatTensor(sequences)

    def __len__(self) -> int:
        return len(self.sequences)

    def __getitem__(self, idx: int) -> torch.Tensor:
        return self.sequences[idx]


# ==================== 训练器 ====================

class EarlyStopping:
    def __init__(self, patience: int = 10, min_delta: float = 1e-6):
        self.patience = patience
        self.min_delta = min_delta
        self.counter = 0
        self.best_loss = float('inf')
        self.early_stop = False

    def __call__(self, val_loss: float) -> bool:
        if val_loss < self.best_loss - self.min_delta:
            self.best_loss = val_loss
            self.counter = 0
        else:
            self.counter += 1
            if self.counter >= self.patience:
                self.early_stop = True
        return self.early_stop


class Trainer:
    def __init__(
        self,
        model: nn.Module,
        train_loader: DataLoader,
        val_loader: DataLoader,
        config: Dict,
        device: torch.device,
        save_dir: str = 'ml/checkpoints_v2'
    ):
        self.model = model.to(device)
        self.train_loader = train_loader
        self.val_loader = val_loader
        self.config = config
        self.device = device
        self.save_dir = Path(save_dir)
        self.save_dir.mkdir(parents=True, exist_ok=True)

        self.optimizer = AdamW(
            model.parameters(),
            lr=config['learning_rate'],
            weight_decay=config['weight_decay']
        )

        self.scheduler = CosineAnnealingWarmRestarts(
            self.optimizer, T_0=10, T_mult=2, eta_min=1e-6
        )

        self.criterion = AnomalyDetectionLoss()

        self.early_stopping = EarlyStopping(
            patience=config['patience'],
            min_delta=config['min_delta']
        )

        self.use_amp = torch.cuda.is_available()
        self.scaler = torch.cuda.amp.GradScaler() if self.use_amp else None
        if self.use_amp:
            print("  ✓ 启用 AMP 混合精度训练")

        self.history = {'train_loss': [], 'val_loss': [], 'learning_rate': []}
        self.best_val_loss = float('inf')

    def train_epoch(self) -> float:
        self.model.train()
        total_loss = 0.0
        n_batches = 0

        pbar = tqdm(self.train_loader, desc="Training", leave=False)
        for batch in pbar:
            batch = batch.to(self.device, non_blocking=True)
            self.optimizer.zero_grad(set_to_none=True)

            if self.use_amp:
                with torch.cuda.amp.autocast():
                    output, latent = self.model(batch)
                    loss, _ = self.criterion(output, batch, latent)

                self.scaler.scale(loss).backward()
                self.scaler.unscale_(self.optimizer)
                torch.nn.utils.clip_grad_norm_(self.model.parameters(), self.config['gradient_clip'])
                self.scaler.step(self.optimizer)
                self.scaler.update()
            else:
                output, latent = self.model(batch)
                loss, _ = self.criterion(output, batch, latent)
                loss.backward()
                torch.nn.utils.clip_grad_norm_(self.model.parameters(), self.config['gradient_clip'])
                self.optimizer.step()

            total_loss += loss.item()
            n_batches += 1
            pbar.set_postfix({'loss': f"{loss.item():.4f}"})

        return total_loss / n_batches

    @torch.no_grad()
    def validate(self) -> float:
        self.model.eval()
        total_loss = 0.0
        n_batches = 0

        for batch in self.val_loader:
            batch = batch.to(self.device, non_blocking=True)

            if self.use_amp:
                with torch.cuda.amp.autocast():
                    output, latent = self.model(batch)
                    loss, _ = self.criterion(output, batch, latent)
            else:
                output, latent = self.model(batch)
                loss, _ = self.criterion(output, batch, latent)

            total_loss += loss.item()
            n_batches += 1

        return total_loss / n_batches

    def save_checkpoint(self, epoch: int, val_loss: float, is_best: bool = False):
        model_to_save = self.model.module if hasattr(self.model, 'module') else self.model

        checkpoint = {
            'epoch': epoch,
            'model_state_dict': model_to_save.state_dict(),
            'optimizer_state_dict': self.optimizer.state_dict(),
            'scheduler_state_dict': self.scheduler.state_dict(),
            'val_loss': val_loss,
            'config': self.config,
        }

        torch.save(checkpoint, self.save_dir / 'last_checkpoint.pt')

        if is_best:
            torch.save(checkpoint, self.save_dir / 'best_model.pt')
            print(f"  ✓ 保存最佳模型 (val_loss: {val_loss:.6f})")

    def train(self, epochs: int):
        print(f"\n开始训练 ({epochs} epochs)...")
        print(f"设备: {self.device}")
        print(f"模型参数量: {count_parameters(self.model):,}")

        for epoch in range(1, epochs + 1):
            print(f"\nEpoch {epoch}/{epochs}")

            train_loss = self.train_epoch()
            val_loss = self.validate()

            self.scheduler.step()
            current_lr = self.optimizer.param_groups[0]['lr']

            self.history['train_loss'].append(train_loss)
            self.history['val_loss'].append(val_loss)
            self.history['learning_rate'].append(current_lr)

            print(f"  Train Loss: {train_loss:.6f}")
            print(f"  Val Loss:   {val_loss:.6f}")
            print(f"  LR:         {current_lr:.2e}")

            is_best = val_loss < self.best_val_loss
            if is_best:
                self.best_val_loss = val_loss
            self.save_checkpoint(epoch, val_loss, is_best)

            if self.early_stopping(val_loss):
                print(f"\n早停触发！")
                break

        print(f"\n训练完成！最佳验证损失: {self.best_val_loss:.6f}")
        self.save_history()

        return self.history

    def save_history(self):
        history_path = self.save_dir / 'training_history.json'
        with open(history_path, 'w') as f:
            json.dump(self.history, f, indent=2)
        print(f"训练历史已保存: {history_path}")

        if HAS_MATPLOTLIB:
            self.plot_training_curves()

    def plot_training_curves(self):
        fig, axes = plt.subplots(1, 2, figsize=(14, 5))
        epochs = range(1, len(self.history['train_loss']) + 1)

        ax1 = axes[0]
        ax1.plot(epochs, self.history['train_loss'], 'b-', label='Train Loss', linewidth=2)
        ax1.plot(epochs, self.history['val_loss'], 'r-', label='Val Loss', linewidth=2)
        ax1.set_xlabel('Epoch')
        ax1.set_ylabel('Loss')
        ax1.set_title('Training & Validation Loss (V2)')
        ax1.legend()
        ax1.grid(True, alpha=0.3)

        best_epoch = np.argmin(self.history['val_loss']) + 1
        best_val_loss = min(self.history['val_loss'])
        ax1.axvline(x=best_epoch, color='g', linestyle='--', alpha=0.7)
        ax1.scatter([best_epoch], [best_val_loss], color='g', s=100, zorder=5)

        ax2 = axes[1]
        ax2.plot(epochs, self.history['learning_rate'], 'g-', linewidth=2)
        ax2.set_xlabel('Epoch')
        ax2.set_ylabel('Learning Rate')
        ax2.set_title('Learning Rate Schedule')
        ax2.set_yscale('log')
        ax2.grid(True, alpha=0.3)

        plt.tight_layout()
        plt.savefig(self.save_dir / 'training_curves.png', dpi=150, bbox_inches='tight')
        plt.close()
        print(f"训练曲线已保存")


# ==================== 阈值计算 ====================

@torch.no_grad()
def compute_thresholds(
    model: nn.Module,
    data_loader: DataLoader,
    device: torch.device,
    percentiles: List[float] = [90, 95, 99]
) -> Dict[str, float]:
    """在验证集上计算阈值"""
    model.eval()
    all_errors = []

    print("计算异动阈值...")
    for batch in tqdm(data_loader, desc="Computing thresholds"):
        batch = batch.to(device)
        errors = model.compute_reconstruction_error(batch, reduction='none')
        seq_errors = errors[:, -1]  # 最后一个时刻
        all_errors.append(seq_errors.cpu().numpy())

    all_errors = np.concatenate(all_errors)

    thresholds = {}
    for p in percentiles:
        threshold = np.percentile(all_errors, p)
        thresholds[f'p{p}'] = float(threshold)
        print(f"  P{p}: {threshold:.6f}")

    thresholds['mean'] = float(np.mean(all_errors))
    thresholds['std'] = float(np.std(all_errors))
    thresholds['median'] = float(np.median(all_errors))

    return thresholds


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

def main():
    parser = argparse.ArgumentParser(description='训练 V2 模型')
    parser.add_argument('--data_dir', type=str, default='ml/data_v2', help='V2 数据目录')
    parser.add_argument('--epochs', type=int, default=100)
    parser.add_argument('--batch_size', type=int, default=4096)
    parser.add_argument('--lr', type=float, default=3e-4)
    parser.add_argument('--device', type=str, default='auto')
    parser.add_argument('--save_dir', type=str, default='ml/checkpoints_v2')
    parser.add_argument('--train_end', type=str, default='2024-06-30')
    parser.add_argument('--val_end', type=str, default='2024-09-30')
    parser.add_argument('--seq_len', type=int, default=10, help='序列长度（分钟）')

    args = parser.parse_args()

    config = TRAIN_CONFIG.copy()
    config['batch_size'] = args.batch_size
    config['epochs'] = args.epochs
    config['learning_rate'] = args.lr
    config['train_end_date'] = args.train_end
    config['val_end_date'] = args.val_end
    config['seq_len'] = args.seq_len

    if args.device == 'auto':
        device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    else:
        device = torch.device(args.device)

    print("=" * 60)
    print("概念异动检测模型训练 V2（Z-Score 特征）")
    print("=" * 60)
    print(f"数据目录: {args.data_dir}")
    print(f"设备: {device}")
    print(f"序列长度: {config['seq_len']} 分钟")
    print(f"批次大小: {config['batch_size']}")
    print(f"特征: {config['features']}")
    print("=" * 60)

    # 1. 加载数据
    print("\n[1/6] 加载 V2 数据...")
    date_data = load_data_by_date(args.data_dir, config['features'])

    # 2. 划分数据集
    print("\n[2/6] 划分数据集...")
    train_data, val_data, test_data = split_data_by_date(
        date_data, config['train_end_date'], config['val_end_date']
    )

    # 3. 构建序列
    print("\n[3/6] 构建序列...")
    print("训练集:")
    train_sequences = build_sequences_by_concept(
        train_data, config['features'], config['seq_len'], config['stride']
    )
    print("验证集:")
    val_sequences = build_sequences_by_concept(
        val_data, config['features'], config['seq_len'], config['stride']
    )

    if len(train_sequences) == 0:
        print("错误: 训练集为空！")
        return

    # 4. 预处理
    print("\n[4/6] 数据预处理...")
    clip_value = config['clip_value']
    print(f"  Z-Score 特征已标准化，截断: ±{clip_value}")

    train_sequences = np.clip(train_sequences, -clip_value, clip_value)
    if len(val_sequences) > 0:
        val_sequences = np.clip(val_sequences, -clip_value, clip_value)

    # 保存配置
    save_dir = Path(args.save_dir)
    save_dir.mkdir(parents=True, exist_ok=True)

    with open(save_dir / 'config.json', 'w') as f:
        json.dump(config, f, indent=2)

    # 5. 创建数据加载器
    print("\n[5/6] 创建数据加载器...")
    train_dataset = SequenceDataset(train_sequences)
    val_dataset = SequenceDataset(val_sequences) if len(val_sequences) > 0 else None

    print(f"  训练序列: {len(train_dataset):,}")
    print(f"  验证序列: {len(val_dataset) if val_dataset else 0:,}")

    n_gpus = torch.cuda.device_count() if torch.cuda.is_available() else 1
    num_workers = min(32, 8 * n_gpus) if sys.platform != 'win32' else 0

    train_loader = DataLoader(
        train_dataset,
        batch_size=config['batch_size'],
        shuffle=True,
        num_workers=num_workers,
        pin_memory=True,
        prefetch_factor=4 if num_workers > 0 else None,
        persistent_workers=True if num_workers > 0 else False,
        drop_last=True
    )

    val_loader = DataLoader(
        val_dataset,
        batch_size=config['batch_size'] * 2,
        shuffle=False,
        num_workers=num_workers,
        pin_memory=True,
    ) if val_dataset else None

    # 6. 训练
    print("\n[6/6] 训练模型...")
    model = TransformerAutoencoder(**config['model'])

    if torch.cuda.device_count() > 1:
        print(f"  使用 {torch.cuda.device_count()} 张 GPU 并行训练")
        model = nn.DataParallel(model)

    if val_loader is None:
        print("警告: 验证集为空，使用训练集的 10% 作为验证")
        split_idx = int(len(train_dataset) * 0.9)
        train_subset = torch.utils.data.Subset(train_dataset, range(split_idx))
        val_subset = torch.utils.data.Subset(train_dataset, range(split_idx, len(train_dataset)))
        train_loader = DataLoader(train_subset, batch_size=config['batch_size'], shuffle=True, num_workers=num_workers, pin_memory=True)
        val_loader = DataLoader(val_subset, batch_size=config['batch_size'], shuffle=False, num_workers=num_workers, pin_memory=True)

    trainer = Trainer(
        model=model,
        train_loader=train_loader,
        val_loader=val_loader,
        config=config,
        device=device,
        save_dir=args.save_dir
    )

    trainer.train(config['epochs'])

    # 计算阈值
    print("\n[额外] 计算异动阈值...")
    best_checkpoint = torch.load(save_dir / 'best_model.pt', map_location=device)

    # 创建新的单 GPU 模型用于计算阈值（避免 DataParallel 问题）
    threshold_model = TransformerAutoencoder(**config['model'])
    threshold_model.load_state_dict(best_checkpoint['model_state_dict'])
    threshold_model.to(device)
    threshold_model.eval()

    thresholds = compute_thresholds(threshold_model, val_loader, device, config['threshold_percentiles'])

    with open(save_dir / 'thresholds.json', 'w') as f:
        json.dump(thresholds, f, indent=2)

    print("\n" + "=" * 60)
    print("训练完成！")
    print(f"模型保存位置: {args.save_dir}")
    print("=" * 60)


if __name__ == "__main__":
    main()