"""
Qwen3.5-0.8B 本地模型微调示例

这个脚本用于微调本地已下载的 Qwen3.5-0.8B 模型。

使用方法:
    python examples/qwen3.5_0.8b_local_finetune.py

前提条件:
    - 已在本地下载 Qwen3.5-0.8B 模型
    - 模型路径配置在 local_model_path 变量中
"""

import os
import sys

if hasattr(os, 'environ'):
    os.environ.setdefault("ASCEND_LAUNCH_BLOCKING", "1")

sys.path.insert(0, os.path.join(os.path.dirname(__file__), ".."))

from finetunex.models import QwenConfig, load_qwen_model
from finetunex.data import load_dataset, format_dataset, InstructionDataset
from finetunex.trainer import FineTuneTrainer
from finetunex.utils import setup_environment, get_gpu_info, setup_logger


def main():
    # 设置环境和日志
    setup_environment(seed=42)
    logger = setup_logger("Qwen3.5-0.8B_Local_FineTuning")
    
    logger.info("=" * 60)
    logger.info("Qwen3.5-0.8B 本地模型微调")
    logger.info("=" * 60)
    
    # ==================== 配置区域 ====================
    
    # 本地模型路径（请根据实际情况修改）
    # 可以是绝对路径或相对于项目根目录的路径
    local_model_path = "./Qwen3.5-0.8B"
    
    # 如果模型在其他位置，请修改这里
    # 例如：
    # local_model_path = "/path/to/your/models/Qwen3.5-0.8B"
    # local_model_path = "D:\\AI_Models\\Qwen3.5-0.8B"
    
    # 检查模型路径是否存在
    if not os.path.exists(local_model_path):
        logger.error(f"模型路径不存在：{local_model_path}")
        logger.error("请修改脚本中的 local_model_path 变量为正确的模型路径")
        sys.exit(1)
    
    logger.info(f"使用本地模型：{local_model_path}")
    
    # 数据集路径
    dataset_path = os.path.join(os.path.dirname(__file__), "..", "data", "sample_dataset.json")
    
    # 输出目录
    output_dir = "./outputs/qwen3.5-0.8b-finetuned"
    
    # ==================== 模型配置 ====================
    
    config = QwenConfig(
        model_name=local_model_path,  # 使用本地模型路径
        lora_r=16,
        lora_alpha=32,
        lora_dropout=0.05,
        target_modules=[
            "q_proj",
            "k_proj",
            "v_proj",
            "o_proj",
            "gate_proj",
            "up_proj",
            "down_proj",
        ],
        per_device_train_batch_size=1,
        gradient_accumulation_steps=4,
        learning_rate=2e-4,
        num_train_epochs=3,
        max_seq_length=512,
        output_dir=output_dir,
        use_4bit=False,  # 华为升腾 NPU 不支持 4bit 量化
        trust_remote_code=True,
    )
    
    logger.info(f"模型配置：{config.model_name}")
    logger.info(f"LoRA 配置：r={config.lora_r}, alpha={config.lora_alpha}")
    logger.info(f"训练配置：epochs={config.num_train_epochs}, lr={config.learning_rate}")
    logger.info(f"输出目录：{output_dir}")
    
    # ==================== GPU 信息 ====================
    
    gpu_info = get_gpu_info()
    if gpu_info["available"]:
        logger.info(f"GPU 可用：{gpu_info['device_count']} 个设备")
        for i, dev in enumerate(gpu_info["devices"]):
            logger.info(f"  GPU {i}: {dev['name']} ({dev['max_memory']:.2f} GB)")
    else:
        logger.warning("GPU 不可用，将使用 CPU 训练（不推荐）")
    
    # ==================== 加载数据 ====================
    
    logger.info("\n加载数据集...")
    dataset = load_dataset(dataset_path, format="json")
    
    # 格式化数据集
    formatted_dataset = format_dataset(
        dataset,
        instruction_column="instruction",
        input_column="input",
        output_column="output",
    )
    
    logger.info(f"数据集大小：{len(formatted_dataset)} 样本")
    
    # ==================== 加载模型 ====================
    
    logger.info("\n加载本地模型...")
    model, tokenizer, peft_config = load_qwen_model(config)
    
    # ==================== 创建训练数据集 ====================
    
    logger.info("\n创建训练数据集...")
    train_dataset = InstructionDataset(
        formatted_dataset,
        tokenizer,
        max_length=config.max_seq_length,
    )
    
    # ==================== 创建训练器 ====================
    
    logger.info("\n创建训练器...")
    trainer = FineTuneTrainer(
        model=model,
        tokenizer=tokenizer,
        config=config,
        train_dataset=train_dataset,
    )
    
    # ==================== 设置训练参数 ====================
    
    trainer.setup_training(
        output_dir=config.output_dir,
        num_train_epochs=config.num_train_epochs,
        per_device_train_batch_size=config.per_device_train_batch_size,
        gradient_accumulation_steps=config.gradient_accumulation_steps,
        learning_rate=config.learning_rate,
        warmup_ratio=0.03,
        weight_decay=0.01,
        logging_steps=10,
        save_steps=50,
        bf16=True,
        fp16=False,
    )
    
    # ==================== 开始训练 ====================
    
    logger.info("\n" + "=" * 60)
    logger.info("开始训练...")
    logger.info("=" * 60)
    
    trainer.train()
    
    # ==================== 保存模型 ====================
    
    logger.info("\n保存模型...")
    trainer.save_model()
    
    logger.info("=" * 60)
    logger.info("训练完成！")
    logger.info(f"模型已保存到：{config.output_dir}")
    logger.info("=" * 60)
    
    # ==================== 测试推理 ====================
    
    logger.info("\n测试推理...")
    test_prompts = [
        "请解释什么是机器学习",
        "写一首关于春天的诗",
    ]
    
    for test_prompt in test_prompts:
        logger.info(f"\n输入：{test_prompt}")
        
        inputs = tokenizer(test_prompt, return_tensors="pt")
        if torch.cuda.is_available():
            inputs = inputs.to("cuda")
        elif hasattr(torch, 'npu') and torch.npu.is_available():
            inputs = inputs.to("npu:0")
        
        with torch.no_grad():
            outputs = model.generate(
                **inputs,
                max_new_tokens=150,
                temperature=0.7,
                do_sample=True,
                top_p=0.9,
            )
        
        response = tokenizer.decode(outputs[0], skip_special_tokens=True)
        logger.info(f"输出：{response}")
    
    logger.info("\n" + "=" * 60)
    logger.info("所有任务完成！")
    logger.info("使用以下命令进行推理：")
    logger.info(f"  python scripts/inference.py --model_path {output_dir} --interactive")
    logger.info("=" * 60)


if __name__ == "__main__":
    import torch
    main()