# -*- coding: utf-8 -*-
"""
Kronos-Enhanced FT-Transformer - 训练引擎
==========================================
纯NumPy实现的全流程训练引擎
支持: 训练/验证/测试 | 早停 | 学习率调度 | 月度拟合度校验 | 自动迭代优化
"""

import numpy as np
import json
import time
import os
import sys
from datetime import datetime, timedelta


class Trainer:
    """
    训练引擎
    ========
    功能:
      1. 数值梯度下降训练FT-Transformer
      2. 四维奖励机制驱动的损失函数
      3. 早停机制（patience）
      4. 余弦退火学习率调度
      5. 月度拟合度校验
      6. 自动迭代优化闭环
    """

    def __init__(self, model, loss_fn, config, checkpoint_dir='./checkpoints', report_dir='./reports'):
        self.model = model
        self.loss_fn = loss_fn
        self.config = config
        self.checkpoint_dir = checkpoint_dir
        self.report_dir = report_dir

        self.batch_size = config.get('batch_size', 8)
        self.lr = config.get('learning_rate', 1e-4)
        self.max_epochs = config.get('max_epochs', 50)
        self.patience = config.get('early_stopping_patience', 10)
        self.grad_clip = config.get('gradient_clip_norm', 1.0)
        self.warmup_epochs = config.get('warmup_epochs', 3)
        self.scheduler_type = config.get('scheduler_type', 'cosine')
        self.weight_decay = config.get('weight_decay', 1e-5)

        self.best_val_loss = float('inf')
        self.best_epoch = 0
        self.patience_counter = 0
        self.history = {'train_loss': [], 'val_loss': [], 'reward_score': [],
                        'lr': [], 'epoch_time': []}
        self.grad_eps = 1e-5  # 数值梯度步长

    def _cosine_lr(self, epoch):
        """余弦退火学习率"""
        if epoch < self.warmup_epochs:
            return self.lr * (epoch + 1) / self.warmup_epochs
        progress = (epoch - self.warmup_epochs) / max(1, self.max_epochs - self.warmup_epochs)
        return self.lr * 0.5 * (1 + np.cos(np.pi * progress))

    def _compute_sga_gradient(self, params_dict, x_cont, x_cat, y, features=None):
        """
        随机梯度逼近 (Simultaneous Perturbation Stochastic Approximation)
        仅需2次前向传播即可估计所有参数的梯度方向，适合CPU环境
        """
        eps = self.grad_eps
        delta = {}

        # 为每个梯度目标生成随机扰动向量
        grad_targets = {
            'head_W': self.model.head_W,
            'head_b': self.model.head_b,
            'head_norm_gamma': self.model.head_norm_gamma,
            'head_norm_beta': self.model.head_norm_beta,
        }
        if self.model.blocks:
            last_block = self.model.blocks[-1]
            grad_targets['last_ffn_W2'] = last_block.ffn.W2
            grad_targets['last_ffn_b2'] = last_block.ffn.b2

        for name, param in grad_targets.items():
            delta[name] = np.random.choice([-1.0, 1.0], size=param.shape)

        # 前向：所有参数 +eps*delta
        base_pred = self.model.forward(x_cont, x_cat, training=False)[0]
        base_loss = self.loss_fn.compute(base_pred, y, features)

        # 加扰动
        saved = {}
        for name, param in grad_targets.items():
            saved[name] = param.copy()
            perturbed = param + eps * delta[name]

            if name == 'head_W':
                self.model.head_W = perturbed
            elif name == 'head_b':
                self.model.head_b = perturbed
            elif name == 'head_norm_gamma':
                self.model.head_norm_gamma = perturbed
            elif name == 'head_norm_beta':
                self.model.head_norm_beta = perturbed
            elif name == 'last_ffn_W2':
                last_block.ffn.W2 = perturbed
            elif name == 'last_ffn_b2':
                last_block.ffn.b2 = perturbed

        pert_pred = self.model.forward(x_cont, x_cat, training=False)[0]
        pert_loss = self.loss_fn.compute(pert_pred, y, features)

        # 恢复参数
        for name, param in grad_targets.items():
            orig = saved[name]
            if name == 'head_W':
                self.model.head_W = orig
            elif name == 'head_b':
                self.model.head_b = orig
            elif name == 'head_norm_gamma':
                self.model.head_norm_gamma = orig
            elif name == 'head_norm_beta':
                self.model.head_norm_beta = orig
            elif name == 'last_ffn_W2':
                last_block.ffn.W2 = orig
            elif name == 'last_ffn_b2':
                last_block.ffn.b2 = orig

        # 计算SGA梯度估计
        loss_diff = pert_loss - base_loss
        grads = {}
        for name, param in grad_targets.items():
            grads[name] = loss_diff * delta[name] / eps

        return grads, base_loss

    def _apply_gradients(self, grads, lr):
        """应用梯度更新（带裁剪）"""
        for name, grad in grads.items():
            # 梯度裁剪
            grad_norm = np.linalg.norm(grad)
            if grad_norm > self.grad_clip:
                grad = grad * self.grad_clip / grad_norm

            # SGD更新
            if name == 'head_W':
                self.model.head_W -= lr * grad
            elif name == 'head_b':
                self.model.head_b -= lr * grad
            elif name == 'head_norm_gamma':
                self.model.head_norm_gamma -= lr * grad
            elif name == 'head_norm_beta':
                self.model.head_norm_beta -= lr * grad
            elif name == 'last_ffn_W2' and self.model.blocks:
                self.model.blocks[-1].ffn.W2 -= lr * grad
            elif name == 'last_ffn_b2' and self.model.blocks:
                self.model.blocks[-1].ffn.b2 -= lr * grad

    def train(self, X_cont_train, X_cat_train, y_train,
              X_cont_val=None, X_cat_val=None, y_val=None, verbose=True):
        """
        执行训练
        """
        print("\n" + "=" * 70)
        print("  FT-Transformer 训练启动")
        print("=" * 70)
        print("  训练集: {} 样本 | 验证集: {} 样本".format(
            len(X_cont_train), len(X_cont_val) if X_cont_val is not None else 0))
        print("  批次大小: {} | 学习率: {} | 最大轮数: {} | 早停patience: {}".format(
            self.batch_size, self.lr, self.max_epochs, self.patience))
        print("  梯度方式: SGA随机梯度逼近（高效CPU优化）| 梯度步长: {}".format(self.grad_eps))
        print("=" * 70)

        train_start = datetime.now()
        self.best_val_loss = float('inf')

        for epoch in range(1, self.max_epochs + 1):
            epoch_start = time.time()
            current_lr = self._cosine_lr(epoch)

            # === 训练阶段 ===
            train_losses = []
            reward_scores = []
            indices = np.random.permutation(len(X_cont_train))

            n_batches = 0
            for start in range(0, len(indices), self.batch_size):
                batch_idx = indices[start:start + self.batch_size]
                x_c = X_cont_train[batch_idx]
                x_cat = X_cat_train[batch_idx]
                y_batch = y_train[batch_idx]

                # 计算梯度（每3个batch计算一次以加速）
                if n_batches % 2 == 0:
                    grads, batch_loss = self._compute_sga_gradient(
                        {}, x_c, x_cat, y_batch, x_c)
                    self._apply_gradients(grads, current_lr)

                preds, _ = self.model.forward(x_c, x_cat, training=True)
                r_score = self.loss_fn.reward.compute(preds, y_batch, x_c)[0]

                train_losses.append(batch_loss)
                reward_scores.append(r_score)
                n_batches += 1

            avg_train_loss = np.mean(train_losses)
            avg_reward = np.mean(reward_scores)

            # === 验证阶段 ===
            val_loss = float('inf')
            if X_cont_val is not None and len(X_cont_val) > 0:
                val_preds, _ = self.model.predict(X_cont_val, X_cat_val)
                val_loss = self.loss_fn.compute(val_preds, y_val, X_cont_val)

            epoch_time = time.time() - epoch_start

            # 记录历史
            self.history['train_loss'].append(float(avg_train_loss))
            self.history['val_loss'].append(float(val_loss))
            self.history['reward_score'].append(float(avg_reward))
            self.history['lr'].append(float(current_lr))
            self.history['epoch_time'].append(epoch_time)

            # 早停检查
            if val_loss < self.best_val_loss:
                self.best_val_loss = val_loss
                self.best_epoch = epoch
                self.patience_counter = 0
                # 保存最佳模型
                self._save_checkpoint(epoch, val_loss)
            else:
                self.patience_counter += 1

            # 打印日志
            if verbose and (epoch % 5 == 0 or epoch == 1):
                print("  Epoch {:3d}/{:3d} | LR: {:.6f} | Train Loss: {:.6f} | "
                      "Val Loss: {:.6f} | Reward: {:.4f} | Time: {:.1f}s | Best: Ep{}".format(
                          epoch, self.max_epochs, current_lr, avg_train_loss,
                      val_loss, avg_reward, epoch_time, self.best_epoch))

            if self.patience_counter >= self.patience:
                print("\n  [早停] 连续 {} 轮无改善，训练终止于 Epoch {}".format(
                    self.patience, epoch))
                break

        train_elapsed = (datetime.now() - train_start).total_seconds()
        print("\n" + "-" * 70)
        print("  训练完成 | 总耗时: {:.1f}s | 最佳Epoch: {} | 最佳Val Loss: {:.6f}".format(
            train_elapsed, self.best_epoch, self.best_val_loss))
        print("-" * 70)

        # 加载最佳模型
        self._load_best_checkpoint()

        return self.history

    def _save_checkpoint(self, epoch, val_loss):
        path = os.path.join(self.checkpoint_dir, 'best_model.pkl')
        self.model.save(path)
        meta = {'epoch': epoch, 'val_loss': val_loss, 'best_val_loss': self.best_val_loss,
                'history': self.history, 'timestamp': datetime.now().isoformat()}
        meta_path = os.path.join(self.checkpoint_dir, 'best_model_meta.json')
        with open(meta_path, 'w', encoding='utf-8') as f:
            json.dump(meta, f, indent=2, ensure_ascii=False, default=str)

    def _load_best_checkpoint(self):
        path = os.path.join(self.checkpoint_dir, 'best_model.pkl')
        if os.path.exists(path):
            self.model.load(path)
            print("  [INFO] 已加载最佳模型: Epoch {}".format(self.best_epoch))


class MonthlyValidator:
    """
    月度拟合度校验器
    ================
    每自然月计算预测与实际走势的拟合度
    拟合度 >= 80%: 预测成功，保留参数
    拟合度 < 80%: 触发自动修正优化
    """

    def __init__(self, config):
        self.success_threshold = config['fit_threshold']['success_threshold']
        self.warning_threshold = config['fit_threshold']['warning_threshold']
        self.auto_retrain = config['auto_retrain_on_fail']
        self.max_retrain = config['max_retrain_attempts']
        self.adjustments = config['retrain_adjustments']

    def validate(self, predictions, actuals, pred_days=60):
        """
        月度拟合度校验

        参数:
            predictions: (N, 3) 模型预测 [趋势, 变化率, 波动率]
            actuals: (N, 3) 真实值

        返回:
            result: dict, 校验结果
        """
        N = len(predictions)
        if N == 0:
            return {'status': 'no_data', 'fit_score': 0, 'details': {}}

        # 1. 趋势方向准确率
        trend_correct = np.mean(np.sign(predictions[:, 0]) == np.sign(actuals[:, 0]))

        # 2. 价格变化率拟合度 (R²)
        ss_res = np.sum((predictions[:, 1] - actuals[:, 1]) ** 2)
        ss_tot = np.sum((actuals[:, 1] - np.mean(actuals[:, 1])) ** 2)
        r2_score = 1 - ss_res / (ss_tot + 1e-10)

        # 3. 波动率拟合度
        vol_correct = np.mean(np.abs(predictions[:, 2] - actuals[:, 2]) < 0.1)

        # 4. 综合拟合度（加权）
        fit_score = (0.4 * trend_correct + 0.3 * max(0, r2_score) + 0.3 * vol_correct)

        # 5. 判定
        if fit_score >= self.success_threshold:
            status = 'success'
            action = 'retain'
        elif fit_score >= self.warning_threshold:
            status = 'warning'
            action = 'monitor'
        else:
            status = 'failure'
            action = 'retrain' if self.auto_retrain else 'manual_review'

        result = {
            'status': status,
            'fit_score': round(float(fit_score), 4),
            'threshold': self.success_threshold,
            'action': action,
            'samples': N,
            'pred_days': pred_days,
            'details': {
                'trend_accuracy': round(float(trend_correct), 4),
                'price_r2': round(float(r2_score), 4),
                'volatility_accuracy': round(float(vol_correct), 4),
                'avg_trend_error': round(float(np.mean(np.abs(predictions[:, 0] - actuals[:, 0]))), 4),
                'avg_price_error': round(float(np.mean(np.abs(predictions[:, 1] - actuals[:, 1]))), 4),
            },
            'reward_detail': None,
        }

        return result

    def generate_retrain_plan(self, validation_result, current_lr, current_epoch):
        """生成重训计划"""
        plan = {
            'trigger': validation_result['status'],
            'current_fit_score': validation_result['fit_score'],
            'target_fit_score': self.success_threshold,
            'adjustments': {},
            'priority_order': [],
        }

        # 根据具体薄弱环节排序优化优先级
        details = validation_result['details']
        weak_dims = []

        if details['trend_accuracy'] < 0.7:
            weak_dims.append(('trend', details['trend_accuracy'], '增加波浪理论权重'))
        if details['price_r2'] < 0.5:
            weak_dims.append(('price', details['price_r2'], '调整量价因子权重'))
        if details['volatility_accuracy'] < 0.6:
            weak_dims.append(('volatility', details['volatility_accuracy'], '调整风控参数'))

        weak_dims.sort(key=lambda x: x[1])

        plan['adjustments'] = {
            'learning_rate': current_lr * self.adjustments['learning_rate_decay'],
            'reward_weight_shift': self.adjustments['reward_weight_shift'],
            'max_epochs': 30,
            'weak_dimensions': weak_dims,
        }
        plan['priority_order'] = [d[0] for d in weak_dims]

        return plan


class FitReporter:
    """拟合度报表生成器"""

    @staticmethod
    def generate_report(training_history, validation_results, model_info,
                        output_dir='./reports', format='json'):
        """生成完整的训练+验证报告"""
        os.makedirs(output_dir, exist_ok=True)

        report = {
            'generated_at': datetime.now().isoformat(),
            'model_info': model_info,
            'training': {
                'total_epochs': len(training_history.get('train_loss', [])),
                'final_train_loss': training_history['train_loss'][-1] if training_history.get('train_loss') else None,
                'best_val_loss': min(training_history['val_loss']) if training_history.get('val_loss') else None,
                'final_reward_score': training_history['reward_score'][-1] if training_history.get('reward_score') else None,
                'loss_curve': training_history.get('train_loss', []),
                'val_loss_curve': training_history.get('val_loss', []),
                'reward_curve': training_history.get('reward_score', []),
                'lr_curve': training_history.get('lr', []),
            },
            'validation': validation_results,
            'summary': FitReporter._generate_summary(training_history, validation_results),
        }

        if format == 'json':
            path = os.path.join(output_dir, 'training_report.json')
            with open(path, 'w', encoding='utf-8') as f:
                json.dump(report, f, indent=2, ensure_ascii=False, default=str)
            print("[INFO] 报告已保存: {}".format(path))
            return path

        return report

    @staticmethod
    def _generate_summary(history, val_results):
        """生成摘要"""
        summary_lines = []
        summary_lines.append("=" * 60)
        summary_lines.append("  Kronos-Enhanced FT-Transformer 训练报告")
        summary_lines.append("=" * 60)

        if history.get('train_loss'):
            summary_lines.append("\n[训练概况]")
            summary_lines.append("  总轮数: {}".format(len(history['train_loss'])))
            summary_lines.append("  最终训练损失: {:.6f}".format(history['train_loss'][-1]))
            if history.get('val_loss'):
                summary_lines.append("  最佳验证损失: {:.6f}".format(min(history['val_loss'])))
            if history.get('reward_score'):
                summary_lines.append("  最终奖励分数: {:.4f}".format(history['reward_score'][-1]))

        if val_results:
            summary_lines.append("\n[月度校验]")
            for i, vr in enumerate(val_results):
                status_cn = {'success': '达标', 'warning': '警告', 'failure': '不达标'}.get(vr['status'], vr['status'])
                summary_lines.append("  校验#{}: 拟合度={:.4f} (阈值={}) - {}".format(
                    i + 1, vr['fit_score'], vr['threshold'], status_cn))
                if vr.get('details'):
                    d = vr['details']
                    summary_lines.append("    趋势准确率: {:.4f} | 价格R²: {:.4f} | 波动率准确率: {:.4f}".format(
                        d.get('trend_accuracy', 0), d.get('price_r2', 0), d.get('volatility_accuracy', 0)))

        summary_lines.append("\n" + "=" * 60)
        return '\n'.join(summary_lines)


if __name__ == '__main__':
    print("Trainer模块加载测试")
    print("  Trainer: 数值梯度训练引擎")
    print("  MonthlyValidator: 月度拟合度校验")
    print("  FitReporter: 结构化报表生成")
    print("[OK] 模块加载成功")