#!/usr/bin/env python3
# entrypoint.py  (替换你现有的 main 脚本)
"""
主入口脚本：训练 / 生成 / 评估 / benchmark 等。

已新增：
- 支持 mode == "mc_benchmark" 调用 lmr.glue_benchmark 中的 MC benchmark runner
 （请确保 lmr/glue_benchmark.py 中存在 run_benchmark_with_mc 函数）
"""

import os
import re
import shutil
import csv
import sys
import json
from pathlib import Path

import torch
import torch.distributed as dist
import math
from dotenv import load_dotenv

import hydra
from omegaconf import OmegaConf
from safetensors.torch import load_file
import pdb
# pdb.set_trace()
# Adjust imports to match your project structure
from lmr.config import initialize_config
from lmr.tokenizer import Tokenizer
from lmr.models import get_model 
from lmr.data import get_dataset_splits
from lmr.checkpointing import Checkpointing
from lmr.utils.seed import set_seed
from lmr.training import Bert_Trainer, Trainer
from lmr.generation import Generator
from lmr.benchmark import Benchmark

# 新增 MC benchmark 支持导入
import lmr.glue_benchmark
from lmr.glue_benchmark import run_glue_benchmark
# from lmr.benchmark_with_mc import run_benchmark_with_mc

from lmr.ddp import unwrap_model

DATASET_DIR = Path("datasets")
CHECKPOINT_DIR = Path("/work/jf381/checkpoints")
BENCHMARK_DIR = Path("output")

# =============================================================================
# Helper Functions: Robust Loading
# =============================================================================

def load_weight_data(path_obj, device="cpu"):
    """支持文件夹(sharded safetensors), 单文件(.safetensors), 或 (.pt)"""
    model_state = {}
    
    if path_obj.is_dir():
        index_file = path_obj / "model.safetensors.index.json"
        if index_file.exists():
            print(f"🔹 Detected sharded safetensors folder: {path_obj.name}")
            with open(index_file, 'r') as f:
                index_data = json.load(f)
            weight_map = index_data.get("weight_map", {})
            shards = set(weight_map.values())
            for shard_name in shards:
                shard_path = path_obj / shard_name
                model_state.update(load_file(str(shard_path), device=str(device)))
            return model_state
        else:
            possible = list(path_obj.glob("*.safetensors")) + list(path_obj.glob("*.pt"))
            if not possible:
                return None
            path_obj = possible[0]

    if path_obj.suffix == ".safetensors":
        print(f"🔹 Loading single safetensors: {path_obj.name}")
        return load_file(str(path_obj), device=str(device))
    else:
        print(f"🔹 Loading pickle (.pt): {path_obj.name}")
        ckpt = torch.load(path_obj, map_location=device)
        return ckpt.get("model", ckpt.get("state_dict", ckpt))

def load_state_dict_robust(model, checkpoint_path, strict=False):
    """
    强力加载：自动处理前缀(module., model.)并匹配分片权重。
    """
    path_obj = Path(checkpoint_path)
    if not path_obj.exists():
        print(f"❌ Path not found: {checkpoint_path}")
        return False

    try:
        model_state = load_weight_data(path_obj)
        if model_state is None:
            print("❌ No weight data found in path.")
            return False

        ckpt_keys_map = {}
        for k in model_state.keys():
            clean_k = k.replace("module.", "").replace("_orig_mod.", "").replace("model.", "")
            ckpt_keys_map[clean_k] = k

        target_model = unwrap_model(model)
        target_state = target_model.state_dict()
        
        filtered_state = {}
        matched_count = 0

        for k_target, v_target in target_state.items():
            k_target_clean = k_target.replace("module.", "").replace("_orig_mod.", "").replace("model.", "")
            if k_target_clean in ckpt_keys_map:
                real_ckpt_key = ckpt_keys_map[k_target_clean]
                v_ckpt = model_state[real_ckpt_key]
                try:
                    if v_ckpt.shape == v_target.shape:
                        filtered_state[k_target] = v_ckpt
                        matched_count += 1
                except Exception:
                    # some safetensors keys may be non-tensor mapping values -> skip
                    continue

        msg = target_model.load_state_dict(filtered_state, strict=strict)
        print(f"✅ Loaded {matched_count} parameters. Status: {msg}")
        return True
    except Exception as e:
        print(f"❌ Failed to load checkpoint: {e}")
        import traceback
        traceback.print_exc()
        return False

def setup_model_and_tokenizer(config):
    print(f"🔧 Initializing Tokenizer: {config.tokenizer_base}")
    tokenizer = Tokenizer(config.tokenizer_base)
    print(f"🔧 Initializing Model: {config.model}")
    model = get_model(config.model, tokenizer.vocab_size, tokenizer=tokenizer)
    return tokenizer, model

# =============================================================================
# Mode Logics
# =============================================================================

def run_generation_task(config, model, tokenizer, output_dir, ckpt_name_tag=""):
    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
    model.to(device)
    model.eval()
    
    generator = Generator(config, model, tokenizer, device=device, output_dir=output_dir)
    results = generator.generate()
    
    metrics = {}
    if results and "verification" in results:
        metrics = {
            'acc': results['verification']['accuracy'],
            'correct': results['verification']['correct'],
            'total': results['verification']['total']
        }
        if results.get('token_accuracies'):
            import numpy as np
            metrics['token_acc'] = float(np.mean(results['token_accuracies']))

        if ckpt_name_tag:
            for fname in ["gsm8k_metrics.txt", "gsm8k_generations.txt"]:
                src = output_dir / fname
                if src.exists():
                    dst = output_dir / f"{src.stem}_{ckpt_name_tag}{src.suffix}"
                    shutil.move(src, dst)
    return metrics

def generate(config):
    """Mode: generate (单次运行，支持加载最近的文件夹或.pt权重)"""
    tokenizer, model = setup_model_and_tokenizer(config)

    import pdb
    # pdb.set_trace()
    ckpt_dir = CHECKPOINT_DIR / config.checkpoint_name
    
    load_mode = getattr(config.benchmark, "checkpoint_mode", "recent")
    print(f"🔍 Looking for [{load_mode}] weights in {ckpt_dir}...")

    checkpointing = Checkpointing(model, CHECKPOINT_DIR / config.checkpoint_name)
    checkpointing.load_model_states("recent")
    # all_items = list(ckpt_dir.iterdir()) if ckpt_dir.exists() else []
    # checkpoints = [f for f in all_items if ("optim" not in f.name and "sched" not in f.name and f.name != "metrics_summary.csv")]

    # def sort_key(f):
    #     nums = re.findall(r'\d+', f.name)
    #     return int(nums[-1]) if nums else 0
    # checkpoints.sort(key=sort_key)

    # target_ckpt = None
    # if load_mode == "best":
    #     best_candidates = [f for f in checkpoints if "best" in f.name.lower()]
    #     target_ckpt = best_candidates[0] if best_candidates else (checkpoints[-1] if checkpoints else None)
    # else:
    #     target_ckpt = checkpoints[-1] if checkpoints else None

    # if target_ckpt:
    #     print(f"🚀 Found target: {target_ckpt.name}")
    #     success = load_state_dict_robust(model, target_ckpt)
    #     if not success:
    #         print("⚠️ Load failed, check file integrity.")
    # else:
    #     print(f"❌ No checkpoints found in {ckpt_dir}")

    run_generation_task(config, model, tokenizer, output_dir=ckpt_dir)

def generate_all(config):
    """Mode: generate_all (扫描目录并运行所有 checkpoitns)"""
    tokenizer, model = setup_model_and_tokenizer(config)
    ckpt_dir = CHECKPOINT_DIR / config.checkpoint_name
    
    all_items = list(ckpt_dir.iterdir()) if ckpt_dir.exists() else []
    checkpoints = [f for f in all_items if ("optim" not in f.name and "sched" not in f.name and f.name != "metrics_summary.csv")]
    
    def sort_key(f):
        nums = re.findall(r'\d+', f.name)
        return int(nums[-1]) if nums else 0
    checkpoints.sort(key=sort_key)
    
    print(f"\n🔎 Found {len(checkpoints)} checkpoints.")
    summary_path = ckpt_dir / "metrics_summary.csv"
    
    with open(summary_path, mode='w', newline='') as f:
        writer = csv.writer(f)
        writer.writerow(["checkpoint", "accuracy", "token_accuracy", "correct", "total"])

    for ckpt_path in checkpoints:
        print(f"\n{'-'*40}\nProcessing: {ckpt_path.name}\n{'-'*40}")
        if load_state_dict_robust(model, ckpt_path):
            metrics = run_generation_task(config, model, tokenizer, ckpt_dir, ckpt_name_tag=ckpt_path.name)
            if metrics:
                with open(summary_path, mode='a', newline='') as f:
                    writer = csv.writer(f)
                    writer.writerow([
                        ckpt_path.name,
                        f"{metrics.get('acc', 0):.4f}",
                        f"{metrics.get('token_acc', 0):.4f}",
                        metrics.get('correct', 0),
                        metrics.get('total', 0)
                    ])

def train_model(config):
    tokenizer, model = setup_model_and_tokenizer(config)
    tokenized_dataset_dir = DATASET_DIR / config.tokenizer_base
    splits = get_dataset_splits(config.dataset, 1024, tokenized_dataset_dir)
    checkpointing = Checkpointing(model, CHECKPOINT_DIR / config.checkpoint_name)

    if "bert" in str(config.model).lower():
        trainer = Bert_Trainer(config.training, model, tokenizer, splits, checkpointing, None)
    else:
        trainer = Trainer(config.training, model, tokenizer, splits, checkpointing)
    # trainer.save_as_hf(save_dir='/work/jf381/checkpoints/test', push_config_fixes=True)
    trainer.train()

# =============================================================================
# Main Entry Point
# =============================================================================

@hydra.main(config_path="config", config_name="config", version_base="1.3")
def main(config):
    load_dotenv()
    set_seed(config)
    initialize_config(config)
    print("=== Config ===")
    print(OmegaConf.to_yaml(config))
    mode = config.mode

    if mode == "train":
        train_model(config)
    elif mode == "generate":
        generate(config)
    elif mode == "generate_all":
        generate_all(config)
    elif mode == "benchmark":
        tokenizer, model = setup_model_and_tokenizer(config)
        checkpointing = Checkpointing(model, CHECKPOINT_DIR / config.checkpoint_name)
        benchmarking = Benchmark(config.benchmark, model, tokenizer, checkpointing, BENCHMARK_DIR / config.checkpoint_name)
        benchmarking.run_benchmarks()
    elif mode == "glue_benchmark":
        # 原有 GLUE benchmark (finetune + eval pipeline)
        tokenizer, model = setup_model_and_tokenizer(config)
        checkpointing = Checkpointing(model, CHECKPOINT_DIR / config.checkpoint_name)
        run_glue_benchmark(config.benchmark, tokenizer, model, checkpointing, out_dir=BENCHMARK_DIR / config.checkpoint_name)
    elif mode == "mc_benchmark":
        # 新增: 多选题 benchmark（evaluation-only runner）
        tokenizer, model = setup_model_and_tokenizer(config)
        checkpointing = Checkpointing(model, CHECKPOINT_DIR / config.checkpoint_name)
        # 调用 lmr.glue_benchmark.run_benchmark_with_mc
        try:
            mc_config = config.get("mc_benchmark", config.benchmark if hasattr(config, "benchmark") else {})
        except Exception:
            mc_config = config.benchmark if hasattr(config, "benchmark") else {}
        # run_benchmark_with_mc 返回一个 list/dict 的结果对象
        res = run_benchmark_with_mc(mc_config, tokenizer, model, checkpointing, out_dir=BENCHMARK_DIR / config.checkpoint_name)
        # 把结果保存为 summary csv/json
        summary_path = BENCHMARK_DIR / config.checkpoint_name / "mc_benchmark_summary.json"
        Path(summary_path).parent.mkdir(parents=True, exist_ok=True)
        with open(summary_path, "w", encoding="utf-8") as f:
            json.dump(res, f, indent=2, ensure_ascii=False)
        print(f"[MC-BENCH] summary written to {summary_path}")
    else:
        print(f"❌ Unknown mode: {mode}")

if __name__ == "__main__":
    main()