#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
XiaLi 云端 GPU 训练脚本 v1.0
=============================
两阶段训练 + ModelScope 自动下载数据 + 断点续训

用法（云GPU上）:
  python3 train_cloud.py              # 默认训练
  python3 train_cloud.py --resume      # 从最新检查点续训
  python3 train_cloud.py --epochs 10   # 指定预训练轮数
"""
import sys, os, json, time, glob, re, math, random
from pathlib import Path

import torch
import torch.nn as nn
from torch.utils.data import Dataset, DataLoader
import numpy as np

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from llm_chat_improved import (
    CharTokenizer, LLMModel, LLMChat,
    ConversationDataset, tokenizer_instance
)

# ── 配置 ──────────────────────────────────────────────────
class CloudConfig:
    # 模型架构
    d_model: int = 320
    n_heads: int = 8
    n_layers: int = 6
    d_ff: int = 1536
    max_length: int = 256          # 对齐用
    dropout: float = 0.1
    use_memory: bool = False       # 关掉记忆/思考/MoE 加速
    use_thinking: bool = False
    use_moe: bool = False

    # 预训练
    pt_batch_size: int = 64        # 32G显存可跑64, 24G改32
    pt_grad_accum: int = 1         # 大batch不需要累积
    pt_epochs: int = 5
    pt_lr: float = 5e-4
    pt_min_lr: float = 5e-5
    pt_warmup_steps: int = 200
    pt_max_length: int = 256       # 长序列学更好
    pt_weight_decay: float = 0.1
    pt_max_norm: float = 1.0

    # 对齐
    align_batch_size: int = 32
    align_grad_accum: int = 1
    align_epochs: int = 20
    align_lr: float = 1e-4
    align_min_lr: float = 1e-5
    align_warmup_steps: int = 50
    align_weight_decay: float = 0.01
    align_max_norm: float = 1.0

    # 通用
    device: str = "cuda" if torch.cuda.is_available() else "cpu"
    num_workers: int = 4
    use_amp: bool = True
    seed: int = 42
    save_dir: str = "/root/xiaoli-llm-chat"   # AutoDL 默认路径
    log_every: int = 10
    save_interval: int = 1
    keep_checkpoints: int = 3


# ── 词表加载 ──────────────────────────────────────────────
def load_vocab(vocab_path: str) -> dict:
    special = {'<PAD>': 0, '<UNK>': 1, '<SOS>': 2, '<EOS>': 3,
               '<START>': 4, '<END>': 5, '<SEP>': 6}
    with open(vocab_path, 'r') as f:
        normal_vocab = json.load(f)
    full_vocab = {**special, **normal_vocab}
    # 同步到全局 tokenizer
    tokenizer_instance.char_to_id = full_vocab.copy()
    tokenizer_instance.id_to_char = {v: k for k, v in full_vocab.items()}
    for attr in ['_multi_trie', '_multi_keys', '_single_map']:
        if hasattr(tokenizer_instance, attr): delattr(tokenizer_instance, attr)
    return full_vocab


# ── 预训练数据集（流式加载，不爆内存）────────────────────
class PretrainDataset(Dataset):
    def __init__(self, data_dir: str, vocab: dict, max_length: int = 128):
        self.vocab = vocab
        self.max_length = max_length
        self.pad_id = vocab.get('<PAD>', 0)
        self.texts = []
        data_dir = Path(data_dir)
        for fpath in sorted(data_dir.glob("*")):
            if fpath.suffix in ('.txt', '.json', '.jsonl'):
                self._load_file(fpath)
        if not self.texts:
            raise ValueError(f"预训练数据目录为空: {data_dir}")
        total_chars = sum(len(t) for t in self.texts)
        print(f"  预训练数据: {len(self.texts):,} 条, {total_chars:,} 字符")

    def _load_file(self, path: Path):
        try:
            if path.suffix == '.jsonl':
                with open(path, 'r', encoding='utf-8', errors='ignore') as f:
                    for line in f:
                        try:
                            item = json.loads(line)
                            t = item.get('text', '').strip()
                            if len(t) >= 10: self.texts.append(t)
                        except: continue
        except: pass

    def __len__(self): return len(self.texts)

    def __getitem__(self, idx):
        text = self.texts[idx]
        tokens = tokenizer_instance.encode(text)
        if len(tokens) > self.max_length:
            start = random.randint(0, len(tokens) - self.max_length)
            tokens = tokens[start:start + self.max_length]
        else:
            tokens = tokens + [self.pad_id] * (self.max_length - len(tokens))
        input_ids = torch.tensor(tokens[:-1], dtype=torch.long)
        target_ids = torch.tensor(tokens[1:], dtype=torch.long)
        loss_mask = torch.ones_like(target_ids)
        return {'input_ids': input_ids, 'target_ids': target_ids,
                'loss_mask': loss_mask, 'weights': 1.0}


# ── 训练器 ──────────────────────────────────────────────
class CloudTrainer:
    def __init__(self, model: nn.Module, config: CloudConfig, stage: str = "pretrain"):
        self.model = model.to(config.device)
        self.device = config.device
        self.config = config
        self.stage = stage
        self.criterion = nn.CrossEntropyLoss(ignore_index=0, label_smoothing=0.1)

        if stage == "pretrain":
            base_lr = config.pt_lr
            wd = config.pt_weight_decay
        else:
            base_lr = config.align_lr
            wd = config.align_weight_decay

        # 分层学习率：embedding/fc_out 用半速
        param_groups = []
        for name, param in model.named_parameters():
            if not param.requires_grad: continue
            if 'embedding' in name or 'fc_out' in name:
                param_groups.append({'params': param, 'lr': base_lr * 0.5, 'weight_decay': wd})
            elif 'norm' in name or 'rmsnorm' in name:
                param_groups.append({'params': param, 'lr': base_lr, 'weight_decay': 0.0})
            else:
                param_groups.append({'params': param, 'lr': base_lr, 'weight_decay': wd})

        self.optimizer = torch.optim.AdamW(param_groups, lr=base_lr, betas=(0.9, 0.95), eps=1e-8)
        self.scaler = torch.cuda.amp.GradScaler(enabled=True)
        self.scheduler = None

        if torch.cuda.is_available():
            torch.backends.cuda.matmul.allow_tf32 = True
            torch.backends.cudnn.allow_tf32 = True
            torch.backends.cudnn.benchmark = True

    def build_scheduler(self, total_steps: int):
        if self.stage == "pretrain":
            warmup = self.config.pt_warmup_steps
            min_lr = self.config.pt_min_lr
        else:
            warmup = self.config.align_warmup_steps
            min_lr = self.config.align_min_lr
        base_lr = self.optimizer.param_groups[0]['lr']

        def lr_lambda(step):
            if step < warmup:
                return step / max(1, warmup)
            progress = (step - warmup) / max(1, total_steps - warmup)
            cosine = 0.5 * (1 + math.cos(math.pi * progress))
            return (1 - min_lr / base_lr) * cosine + min_lr / base_lr

        self.scheduler = torch.optim.lr_scheduler.LambdaLR(self.optimizer, lr_lambda)

    def train_epoch(self, dataloader, epoch, total_epochs, stage_name="pretrain"):
        self.model.train()
        total_loss = 0
        n_batches = len(dataloader)
        accum_steps = self.config.pt_grad_accum if stage_name == "pretrain" else self.config.align_grad_accum
        max_norm = self.config.pt_max_norm if stage_name == "pretrain" else self.config.align_max_norm

        self.optimizer.zero_grad()
        for batch_idx, batch in enumerate(dataloader):
            input_ids = batch['input_ids'].to(self.device, non_blocking=True)
            target_ids = batch['target_ids'].to(self.device, non_blocking=True)
            mask = self.model.generate_mask(input_ids.size(1)).to(self.device, non_blocking=True)

            with torch.amp.autocast('cuda', enabled=True):
                outputs = self.model(input_ids, mask)
                loss = self.criterion(outputs.reshape(-1, outputs.size(-1)), target_ids.reshape(-1))
                loss = loss / accum_steps

            self.scaler.scale(loss).backward()

            if (batch_idx + 1) % accum_steps == 0:
                self.scaler.unscale_(self.optimizer)
                torch.nn.utils.clip_grad_norm_(self.model.parameters(), max_norm)
                self.scaler.step(self.optimizer)
                self.scaler.update()
                self.optimizer.zero_grad()
                if self.scheduler: self.scheduler.step()

            total_loss += loss.item() * accum_steps

            if (batch_idx + 1) % max(1, n_batches // 10) == 0:
                lr = self.optimizer.param_groups[0]['lr']
                mem = f" | Mem: {torch.cuda.max_memory_allocated()/1024**3:.1f}GiB" if torch.cuda.is_available() else ""
                print(f"  [{stage_name}] Epoch {epoch}/{total_epochs} | "
                      f"Batch {batch_idx+1}/{n_batches} | Loss: {total_loss/(batch_idx+1):.4f} | LR: {lr:.2e}{mem}")

        return total_loss / n_batches

    def save_checkpoint(self, path: str, vocab: dict, epoch: int, loss: float):
        core = getattr(self.model, '_orig_mod', self.model)
        torch.save({
            'model_state_dict': core.state_dict(),
            'optimizer_state_dict': self.optimizer.state_dict(),
            'scheduler_state_dict': self.scheduler.state_dict() if self.scheduler else None,
            'vocab': vocab,
            'epoch': epoch,
            'loss': loss,
            'config': {k: getattr(self.config, k) for k in dir(self.config)
                       if not k.startswith('_') and not callable(getattr(self.config, k))},
        }, path)
        print(f"  ✅ 已保存: {path} (epoch {epoch}, loss {loss:.4f})")

    def load_checkpoint(self, path: str) -> tuple:
        cp = torch.load(path, map_location=self.device, weights_only=False)
        core = getattr(self.model, '_orig_mod', self.model)
        core.load_state_dict(cp['model_state_dict'])
        self.optimizer.load_state_dict(cp['optimizer_state_dict'])
        if cp.get('scheduler_state_dict') and self.scheduler:
            self.scheduler.load_state_dict(cp['scheduler_state_dict'])
        return cp.get('epoch', 0), cp.get('loss', float('inf')), cp.get('vocab', {})


# ── 数据下载 ──────────────────────────────────────────────
def download_data(save_dir: str):
    """从 ModelScope 下载预训练数据"""
    data_dir = os.path.join(save_dir, 'pretrain_data')
    if os.path.exists(data_dir) and list(Path(data_dir).glob("*.jsonl")):
        print("  预训练数据已存在，跳过下载")
        return data_dir

    os.makedirs(data_dir, exist_ok=True)
    print("  📥 从 ModelScope 下载预训练数据...")
    try:
        import subprocess
        subprocess.run([
            sys.executable, '-m', 'pip', 'install', 'modelscope', '-q'
        ], check=True, capture_output=True)
        subprocess.run([
            sys.executable, '-m', 'modelscope', 'download',
            '--dataset', 'gongjy/minimind_dataset',
            '--include', 'pretrain_t2t_mini*',
            '--local_dir', data_dir,
        ], check=True)
        print(f"  ✅ 下载完成: {data_dir}")
    except Exception as e:
        print(f"  ⚠️ 自动下载失败: {e}")
        print(f"  ⚠️ 请手动下载放到 {data_dir}")
        print(f"  ⚠️ 命令: modelscope download --dataset gongjy/minimind_dataset --include 'pretrain_t2t_mini*' --local_dir {data_dir}")
    return data_dir


# ── 主函数 ──────────────────────────────────────────────
def main():
    import argparse
    parser = argparse.ArgumentParser()
    parser.add_argument("--resume", action="store_true", help="从最新检查点续训")
    parser.add_argument("--epochs", type=int, default=5, help="预训练轮数")
    parser.add_argument("--align-epochs", type=int, default=20, help="对齐轮数")
    parser.add_argument("--save-dir", type=str, default="/root/xiaoli-llm-chat")
    parser.add_argument("--batch-size", type=int, default=0, help="覆盖batch大小")
    args = parser.parse_args()

    config = CloudConfig()
    config.pt_epochs = args.epochs
    config.align_epochs = args.align_epochs
    config.save_dir = args.save_dir

    if args.batch_size > 0:
        config.pt_batch_size = args.batch_size
        config.align_batch_size = args.batch_size // 2

    if not torch.cuda.is_available():
        print("⚠️ 未检测到 GPU！云端训练建议使用 AutoDL / 矩池云 / 恒源云")

    print("=" * 65)
    print(f"  XiaLi 云端训练")
    print(f"  GPU: {torch.cuda.get_device_name(0) if torch.cuda.is_available() else 'N/A'}")
    print(f"  显存: {torch.cuda.get_device_properties(0).total_mem / 1024**3:.1f} GiB" if torch.cuda.is_available() else "")
    print(f"  模型: {config.d_model}d, {config.n_layers}L, {config.n_heads}H")
    print(f"  预训练: {config.pt_epochs} epochs, batch={config.pt_batch_size}")
    print(f"  对齐: {config.align_epochs} epochs, batch={config.align_batch_size}")
    print("=" * 65)

    os.makedirs(config.save_dir, exist_ok=True)

    # 1. 加载词表
    vocab_path = os.path.join(config.save_dir, 'vocab.json')
    if not os.path.exists(vocab_path):
        print(f"❌ 未找到 {vocab_path}")
        print(f"   请从 Gitee 仓库获取: https://gitee.com/shuiyu1123/xiaoli-llm-chat")
        return
    vocab = load_vocab(vocab_path)
    print(f"词表: {len(vocab)} tokens, vocab_size={max(vocab.values())+1}")

    # 2. 创建模型
    model = LLMModel(
        vocab_size=max(vocab.values())+1,
        d_model=config.d_model,
        n_heads=config.n_heads,
        n_layers=config.n_layers,
        d_ff=config.d_ff,
        max_len=config.max_length,
        dropout=config.dropout,
        use_memory=False, use_thinking=False, use_moe=False,
    )
    params = sum(p.numel() for p in model.parameters())
    print(f"参数量: {params:,} ({params/1e6:.2f}M)")

    # 3. 续训检测
    start_epoch = 1
    if args.resume:
        checkpoints = sorted(glob.glob(os.path.join(config.save_dir, 'pretrain_epoch_*.pth')))
        if checkpoints:
            latest = checkpoints[-1]
            trainer_pt = CloudTrainer(model, config, stage="pretrain")
            loaded_epoch, loaded_loss, _ = trainer_pt.load_checkpoint(latest)
            start_epoch = loaded_epoch + 1
            print(f"  从 {latest} 续训 (epoch {loaded_epoch}, loss {loaded_loss:.4f})")

    # 4. 下载数据
    pt_dir = download_data(config.save_dir)

    # 5. 预训练
    print(f"\n{'='*65}")
    print(f"  阶段一：预训练 ({config.pt_epochs} epochs)")
    print(f"{'='*65}")

    pt_dataset = PretrainDataset(pt_dir, vocab, config.pt_max_length)
    pt_loader = DataLoader(
        pt_dataset, batch_size=config.pt_batch_size, shuffle=True,
        num_workers=config.num_workers,
        pin_memory=True,
    )

    trainer = CloudTrainer(model, config, stage="pretrain")
    total_steps = len(pt_loader) * config.pt_epochs
    trainer.build_scheduler(total_steps)

    for epoch in range(start_epoch, config.pt_epochs + 1):
        t0 = time.time()
        loss = trainer.train_epoch(pt_loader, epoch, config.pt_epochs, "pretrain")
        epoch_time = time.time() - t0
        lr = trainer.optimizer.param_groups[0]['lr']
        print(f"  → Epoch {epoch}/{config.pt_epochs} | Loss: {loss:.4f} | LR: {lr:.2e} | 耗时: {epoch_time:.0f}s")

        if epoch % config.save_interval == 0:
            ckpt_path = os.path.join(config.save_dir, f'pretrain_epoch_{epoch}.pth')
            trainer.save_checkpoint(ckpt_path, vocab, epoch, loss)
            # 清理旧检查点
            all_ckpts = sorted(glob.glob(os.path.join(config.save_dir, 'pretrain_epoch_*.pth')))
            for old in all_ckpts[:-config.keep_checkpoints]:
                os.remove(old)

    print(f"\n  ✅ 预训练完成")

    # 6. 对齐
    print(f"\n{'='*65}")
    print(f"  阶段二：角色对齐 ({config.align_epochs} epochs)")
    print(f"{'='*65}")

    merged_path = os.path.join(config.save_dir, 'xiaoli_merged_data_v2.json')
    if not os.path.exists(merged_path):
        print(f"  ⚠️ 未找到合并数据: {merged_path}")
        print(f"  ⚠️ 跳过对齐阶段")
    else:
        with open(merged_path, 'r') as f:
            raw_data = json.load(f)
        pairs = []
        for item in raw_data:
            q = item.get("question", "") or item.get("input", "") or item.get("prompt", "")
            a = item.get("answer", "") or item.get("output", "") or item.get("response", "")
            if q and a: pairs.append((q, a))

        random.seed(config.seed)
        random.shuffle(pairs)
        split = int(len(pairs) * 0.95)
        train_pairs = pairs[:split]
        print(f"  对齐数据: {len(pairs)} 条 (训练 {len(train_pairs)} + 验证 {len(pairs)-split})")

        train_ds = ConversationDataset(train_pairs, vocab, config.max_length)
        train_loader = DataLoader(
            train_ds, batch_size=config.align_batch_size, shuffle=True,
            num_workers=config.num_workers, pin_memory=True,
        )

        trainer2 = CloudTrainer(model, config, stage="align")
        total_align_steps = len(train_loader) * config.align_epochs
        trainer2.build_scheduler(total_align_steps)

        for epoch in range(1, config.align_epochs + 1):
            t0 = time.time()
            loss = trainer2.train_epoch(train_loader, epoch, config.align_epochs, "align")
            epoch_time = time.time() - t0
            print(f"  → Epoch {epoch}/{config.align_epochs} | Loss: {loss:.4f} | 耗时: {epoch_time:.0f}s")

            if epoch % 5 == 0:
                ckpt_path = os.path.join(config.save_dir, f'xiaoli_align_epoch_{epoch}.pth')
                trainer2.save_checkpoint(ckpt_path, vocab, epoch, loss)

        # 保存最终模型
        final_path = os.path.join(config.save_dir, 'xiaoli_final.pth')
        trainer2.save_checkpoint(final_path, vocab, config.align_epochs, loss)
        print(f"\n  ✅ 对齐完成！最终模型: {final_path}")

    print(f"\n{'='*65}")
    print(f"  🎉 训练完成！")
    print(f"  模型保存在: {config.save_dir}")
    print(f"  {'='*65}")


if __name__ == "__main__":
    main()