#!/usr/bin/env python3
"""
Waymo E2E Evaluation Script for dVLM-AD (Fast_dLLM / AR) models.

功能：
  - 加载已训练的 Fast_dLLM 或标准 AR 模型
  - 对 dvlm-ad_waymo_e2e_val_cleaned.json 中的样本逐条推理
  - 输出完整 predictions.json（保留原始 CoT 文本、解析后的 JSON、轨迹对比）
  - 计算自驾评估指标：ADE / FDE / valid_ratio
  - 支持多 GPU 并行（--num_gpus N）

用法：
  python kewei_scripts/eval_dvlmad_waymo.py \
      --model_path output_models/my_checkpoint \
      --eval_json dvlmad_data/waymo_eval/dvlm-ad_waymo_e2e_val_cleaned.json \
      --image_root dvlmad_data \
      --output_dir eval_outputs/waymo_val \
      --max_new_tokens 1024
"""

import argparse
import json
import math
import os
import re
import sys
import time
import types
import traceback
from datetime import datetime
from pathlib import Path
from typing import Any, Dict, List, Optional, Tuple

import numpy as np
import torch
import torch.multiprocessing as mp
from PIL import Image
from tqdm import tqdm
from transformers import AutoModelForCausalLM, AutoProcessor, AutoTokenizer

# ──────────────────────────────────────────────────────────────────
# 路径设置：让 generation_functions 可以被导入
# ──────────────────────────────────────────────────────────────────
_script_dir = Path(__file__).resolve().parent
_project_root = _script_dir.parent
_models_dir = _project_root / "chengyue_script" / "dvlm" / "models"
_src_dir = _project_root / "src"
for _p in [str(_project_root), str(_src_dir), str(_models_dir)]:
    if _p not in sys.path:
        sys.path.insert(0, _p)

# ──────────────────────────────────────────────────────────────────
# 注册自定义模型（fast_dllm_qwen2_5_vl）到 HuggingFace Auto 类
# 必须在模块顶层执行，spawn 多进程时每个子进程都会重新 import，
# 从而保证子进程也完成注册。
# ──────────────────────────────────────────────────────────────────
try:
    from transformers import AutoConfig
    from lmflow.models.qwen2_5_vl_mdm.configuration import Fast_dLLM_Qwen2_5_VLConfig
    from lmflow.models.qwen2_5_vl_mdm.modeling import Fast_dLLM_Qwen2_5_VLForConditionalGeneration
    AutoConfig.register("fast_dllm_qwen2_5_vl", Fast_dLLM_Qwen2_5_VLConfig, exist_ok=True)
    AutoModelForCausalLM.register(
        Fast_dLLM_Qwen2_5_VLConfig,
        Fast_dLLM_Qwen2_5_VLForConditionalGeneration,
        exist_ok=True,
    )
except Exception as _e:
    import warnings
    warnings.warn(f"自定义模型注册失败，加载 checkpoint 时可能报错: {_e}")


# ══════════════════════════════════════════════════════════════════
#  工具函数
# ══════════════════════════════════════════════════════════════════

def resolve_image_path(image_field: str, image_root: str) -> str:
    """
    将 JSON 中的相对路径映射到本地绝对路径。

    JSON 中图片路径形如 "dataset/waymo_eval/images/val/xxx/yyy.jpg"
    本地实际路径形如 "{image_root}/waymo_eval/images/val/xxx/yyy.jpg"
    """
    # 去掉 "dataset/" 前缀
    rel = image_field
    for prefix in ("dataset/", "dataset\\"):
        if rel.startswith(prefix):
            rel = rel[len(prefix):]
            break
    return os.path.join(image_root, rel)


def _extract_float_pairs(s: str) -> Optional[List[List[float]]]:
    """
    从形如 '[[+00.10,-01.20], [+02.30,+00.00]]' 的字符串中用正则提取浮点数对。

    模型输出的轨迹格式包含 +/- 前缀和前导零（如 +00.10），这是非法 JSON，
    因此无法直接用 json.loads 解析，必须用正则提取原始数值。
    """
    pairs = re.findall(r'\[\s*([+-]?\d+\.?\d*)\s*,\s*([+-]?\d+\.?\d*)\s*\]', s)
    if not pairs:
        return None
    return [[float(a), float(b)] for a, b in pairs]


def parse_trajectory(response: str) -> Optional[List[List[float]]]:
    """
    从模型输出文本中解析轨迹 waypoints。

    模型输出应为 JSON 字符串，其中 "trajectory" 字段是
    [[x1,y1],[x2,y2],...] 格式的列表（或含 +/- 前缀的字符串形式）。

    注意：模型输出的数值格式为 [[+00.10,-01.20],...] —— 带 + 前缀和前导零，
    这不是合法 JSON，因此 json.loads 会失败，需要用正则提取。
    """
    if not response:
        return None

    # 尝试整体解析为 JSON
    try:
        obj = json.loads(response)
        traj = obj.get("trajectory")
        if traj is not None:
            if isinstance(traj, str):
                # 优先用正则提取（兼容 +00.10 等非标准 JSON 数值格式）
                result = _extract_float_pairs(traj)
                if result:
                    return result
                # 降级：标准 JSON 解析
                traj = json.loads(traj)
            if isinstance(traj, list) and len(traj) > 0:
                return [[float(p[0]), float(p[1])] for p in traj]
    except Exception:
        pass

    # 降级：用正则从原始文本中定位 trajectory 字段并提取数值对
    # 先找到 trajectory 字段的值字符串，再提取数值对
    patterns = [
        r'"trajectory"\s*:\s*"(\[\[.*?\]\])"',   # 值是 JSON 字符串形式
        r'"trajectory"\s*:\s*(\[\[.*?\]\])',       # 值是裸 JSON 数组形式
    ]
    for pat in patterns:
        m = re.search(pat, response, re.DOTALL)
        if m:
            result = _extract_float_pairs(m.group(1))
            if result:
                return result
            try:
                traj = json.loads(m.group(1))
                return [[float(p[0]), float(p[1])] for p in traj]
            except Exception:
                pass

    return None


def parse_full_output(response: str) -> Optional[Dict]:
    """尝试将模型输出解析为完整 JSON 对象。"""
    if not response:
        return None
    try:
        return json.loads(response)
    except Exception:
        pass
    # 从文本中抽取第一个 {...} 块
    m = re.search(r'\{.*\}', response, re.DOTALL)
    if m:
        try:
            return json.loads(m.group(0))
        except Exception:
            pass
    return None


def compute_trajectory_metrics(
    pred_traj: Optional[List[List[float]]],
    gt_traj: List[List[float]],
) -> Dict[str, Any]:
    """计算 ADE / FDE。"""
    if pred_traj is None or len(pred_traj) == 0:
        return {"ade": None, "fde": None, "valid": False, "num_points": 0}

    pred = np.array(pred_traj, dtype=np.float64)
    gt   = np.array(gt_traj,   dtype=np.float64)

    n = min(len(pred), len(gt))
    if n == 0:
        return {"ade": None, "fde": None, "valid": False, "num_points": 0}

    pred, gt = pred[:n], gt[:n]
    dists = np.linalg.norm(pred - gt, axis=1)
    return {
        "ade":        float(dists.mean()),
        "fde":        float(dists[-1]),
        "valid":      True,
        "num_points": n,
    }


# ══════════════════════════════════════════════════════════════════
#  模型加载 & 推理
# ══════════════════════════════════════════════════════════════════

class WaymoEvaluator:
    """加载模型并提供单样本推理接口。"""

    def __init__(
        self,
        model_path: str,
        base_model: str = "Qwen/Qwen2.5-VL-3B-Instruct",
        device: str = "cuda",
        torch_dtype: torch.dtype = torch.bfloat16,
        min_pixels: Optional[int] = None,
        max_pixels: Optional[int] = None,
    ):
        self.device = device
        self.is_mdm = False
        self.block_size = None
        self.mask_id = None

        print(f"[{device}] 加载 processor: {base_model}", flush=True)
        proc_kwargs = {}
        if min_pixels is not None:
            proc_kwargs["min_pixels"] = min_pixels
        if max_pixels is not None:
            proc_kwargs["max_pixels"] = max_pixels

        self.processor = AutoProcessor.from_pretrained(
            base_model, trust_remote_code=True, **proc_kwargs
        )

        print(f"[{device}] 加载模型: {model_path}", flush=True)
        try:
            # 优先用 AutoModelForCausalLM（适配自定义 Fast_dLLM 模型）
            self.model = AutoModelForCausalLM.from_pretrained(
                model_path,
                torch_dtype=torch_dtype,
                device_map=device,
                trust_remote_code=True,
            ).eval()
        except ValueError as e:
            # 标准 Qwen2.5-VL / Qwen2-VL 需要用 AutoModelForVision2Seq 加载
            if "Qwen2_5_VLConfig" in str(e) or "Qwen2VLConfig" in str(e) or "Unrecognized configuration class" in str(e):
                print(f"[{device}] AutoModelForCausalLM 不支持该配置，改用 AutoModelForVision2Seq...", flush=True)
                from transformers import AutoModelForVision2Seq
                self.model = AutoModelForVision2Seq.from_pretrained(
                    model_path,
                    torch_dtype=torch_dtype,
                    device_map=device,
                    trust_remote_code=True,
                ).eval()
            else:
                raise

        # 检测是否为 MDM 模型
        cfg = self.model.config
        if hasattr(cfg, "bd_size") and cfg.bd_size is not None:
            self.is_mdm = True
            self.block_size = cfg.bd_size
            tokenizer = AutoTokenizer.from_pretrained(model_path)
            mask_ids = tokenizer.encode("|<MASK>|", add_special_tokens=False)
            self.mask_id = mask_ids[0] if mask_ids else 151665
            print(f"[{device}] 检测到 MDM 模型，bd_size={self.block_size}，mask_id={self.mask_id}", flush=True)

            from generation_functions import Fast_dLLM_QwenForCausalLM
            self.model.mdm_sample = types.MethodType(
                Fast_dLLM_QwenForCausalLM.mdm_sample_cache, self.model
            )
        else:
            print(f"[{device}] 检测到标准 Qwen2.5-VL / AR 模型", flush=True)

    @torch.no_grad()
    def generate(
        self,
        image: Image.Image,
        question: str,
        max_new_tokens: int = 1024,
    ) -> str:
        """对单张图 + 问题进行推理，返回原始文本输出。"""
        messages = [
            {
                "role": "user",
                "content": [
                    {"type": "image", "image": image},
                    {"type": "text",  "text":  question},
                ],
            }
        ]

        text = self.processor.apply_chat_template(
            messages, tokenize=False, add_generation_prompt=True
        )
        inputs = self.processor(
            text=[text],
            images=[image],
            padding=True,
            return_tensors="pt",
        ).to(self.device)

        if self.is_mdm:
            # MDM 并行解码
            mdm_kwargs = dict(
                input_ids=inputs.input_ids,
                tokenizer=self.processor.tokenizer,
                block_size=self.block_size,
                max_tokens=max_new_tokens,
                mask_id=self.mask_id,
                return_stats=False,
            )
            if hasattr(inputs, "pixel_values") and inputs.pixel_values is not None:
                mdm_kwargs["pixel_values"] = inputs.pixel_values
            if hasattr(inputs, "image_grid_thw") and inputs.image_grid_thw is not None:
                mdm_kwargs["image_grid_thw"] = inputs.image_grid_thw

            output_ids = self.model.mdm_sample(**mdm_kwargs)
        else:
            # 标准自回归
            output_ids = self.model.generate(
                **inputs,
                max_new_tokens=max_new_tokens,
                do_sample=False,
            )

        # 去掉 prompt 部分
        generated_ids = output_ids[:, inputs.input_ids.shape[1]:]
        response = self.processor.batch_decode(
            generated_ids, skip_special_tokens=True
        )[0]
        return response.strip()


# ══════════════════════════════════════════════════════════════════
#  单进程评估逻辑
# ══════════════════════════════════════════════════════════════════

def evaluate_samples(
    evaluator: WaymoEvaluator,
    samples: List[Dict],
    image_root: str,
    max_new_tokens: int,
    rank: int = 0,
) -> List[Dict]:
    """
    对一批样本进行推理，返回每条样本的完整结果。

    返回值每条包含：
        sample_id, image, navigation_command,
        question, gt_answer_template,
        model_output_raw,   ← 原始模型输出（供 CoT 检查）
        model_output_parsed, ← 解析后的 JSON（若有效）
        pred_trajectory,    ← 解析出的预测轨迹
        gt_trajectory,      ← ground-truth 轨迹
        metrics             ← {ade, fde, valid, num_points}
        error               ← 若推理出错则记录
    """
    results = []

    for item in tqdm(samples, desc=f"GPU{rank} 推理", position=rank, leave=True):
        sample_id = item.get("sample_id", "")
        image_field = item.get("image", "")
        gt_traj = item.get("future waypoints", [])
        # GT 轨迹在 JSON 中以字符串形式存储，需要解析为列表
        if isinstance(gt_traj, str):
            try:
                gt_traj = json.loads(gt_traj)
            except Exception:
                gt_traj = []
        nav_cmd = item.get("navigation_command", "")
        conversations = item.get("conversations", [])

        # 取 human 端的问题
        question = ""
        gt_answer_template = ""
        for msg in conversations:
            if msg.get("from") == "human":
                question = msg.get("value", "")
            elif msg.get("from") == "assistant":
                gt_answer_template = msg.get("value", "")

        # 构建图片完整路径
        image_path = resolve_image_path(image_field, image_root)

        entry = {
            "sample_id":          sample_id,
            "image":              image_field,
            "image_full_path":    image_path,
            "navigation_command": nav_cmd,
            "question":           question,
            "gt_answer_template": gt_answer_template,
            "gt_trajectory":      gt_traj,
            "model_output_raw":   None,
            "model_output_parsed": None,
            "pred_trajectory":    None,
            "metrics":            {"ade": None, "fde": None, "valid": False},
            "error":              None,
        }

        if not question:
            entry["error"] = "missing question in conversations"
            results.append(entry)
            continue

        if not os.path.isfile(image_path):
            entry["error"] = f"image not found: {image_path}"
            results.append(entry)
            continue

        try:
            image = Image.open(image_path).convert("RGB")
            raw_output = evaluator.generate(image, question, max_new_tokens=max_new_tokens)
            entry["model_output_raw"] = raw_output
            entry["model_output_parsed"] = parse_full_output(raw_output)
            entry["pred_trajectory"] = parse_trajectory(raw_output)
            entry["metrics"] = compute_trajectory_metrics(entry["pred_trajectory"], gt_traj)
        except Exception as e:
            entry["error"] = traceback.format_exc()
            print(f"\n[GPU{rank}] 错误 sample_id={sample_id}: {e}", flush=True)

        results.append(entry)

    return results


# ══════════════════════════════════════════════════════════════════
#  多 GPU 并行
# ══════════════════════════════════════════════════════════════════

def worker_fn(
    rank: int,
    world_size: int,
    model_path: str,
    base_model: str,
    all_samples: List[Dict],
    image_root: str,
    max_new_tokens: int,
    min_pixels: Optional[int],
    max_pixels: Optional[int],
    return_dict: dict,
):
    """多进程 worker：负责自己分片的样本。

    无论成功还是崩溃，都保证向 return_dict[rank] 写入结果，
    避免主进程收集时出现 KeyError。
    """
    torch.cuda.set_device(rank)
    device = f"cuda:{rank}"
    shard = all_samples[rank::world_size]

    try:
        evaluator = WaymoEvaluator(
            model_path=model_path,
            base_model=base_model,
            device=device,
            min_pixels=min_pixels,
            max_pixels=max_pixels,
        )
        results = evaluate_samples(evaluator, shard, image_root, max_new_tokens, rank)
    except Exception as e:
        err_msg = traceback.format_exc()
        print(f"\n[GPU{rank}] Worker 崩溃: {e}\n{err_msg}", flush=True)
        # 把该 shard 的所有样本标记为失败，保留基本字段
        results = []
        for item in shard:
            results.append({
                "sample_id":          item.get("sample_id", ""),
                "image":              item.get("image", ""),
                "navigation_command": item.get("navigation_command", ""),
                "question":           "",
                "gt_answer_template": "",
                "gt_trajectory":      item.get("future waypoints", []),
                "model_output_raw":   None,
                "model_output_parsed": None,
                "pred_trajectory":    None,
                "metrics":            {"ade": None, "fde": None, "valid": False},
                "error":              f"Worker {rank} crashed: {e}",
            })

    return_dict[rank] = results


# ══════════════════════════════════════════════════════════════════
#  聚合指标
# ══════════════════════════════════════════════════════════════════

def aggregate_metrics(predictions: List[Dict]) -> Dict[str, Any]:
    """从全部预测结果计算汇总指标。"""
    total    = len(predictions)
    failed   = sum(1 for p in predictions if p.get("error"))
    evaluated = total - failed

    valid_preds = [p for p in predictions if p["metrics"].get("valid")]
    n_valid = len(valid_preds)

    agg = {
        "total_samples":          total,
        "evaluated":              evaluated,
        "failed":                 failed,
        "valid_trajectory":       n_valid,
        "valid_trajectory_ratio": n_valid / evaluated if evaluated > 0 else 0.0,
    }

    if n_valid > 0:
        ades = [p["metrics"]["ade"] for p in valid_preds]
        fdes = [p["metrics"]["fde"] for p in valid_preds]
        agg["ade_mean"]   = float(np.mean(ades))
        agg["ade_std"]    = float(np.std(ades))
        agg["ade_median"] = float(np.median(ades))
        agg["fde_mean"]   = float(np.mean(fdes))
        agg["fde_std"]    = float(np.std(fdes))
        agg["fde_median"] = float(np.median(fdes))
    else:
        agg["ade_mean"] = agg["fde_mean"] = None

    return agg


# ══════════════════════════════════════════════════════════════════
#  入口
# ══════════════════════════════════════════════════════════════════

def main():
    parser = argparse.ArgumentParser(
        description="dVLM-AD Waymo E2E Evaluation (Fast_dLLM / AR)",
        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
    )
    parser.add_argument("--model_path",     required=True,
                        help="已训练模型的 checkpoint 路径")
    parser.add_argument("--eval_json",
                        default="dvlmad_data/waymo_eval/dvlm-ad_waymo_e2e_val_cleaned.json",
                        help="评估数据集 JSON 路径")
    parser.add_argument("--image_root",
                        default="dvlmad_data",
                        help="图像根目录（JSON 中 dataset/ 前缀对应此目录）")
    parser.add_argument("--output_dir",
                        default="eval_outputs/waymo_val",
                        help="输出目录，存放 predictions.json 和 metrics.json")
    parser.add_argument("--base_model",
                        default="Qwen/Qwen2.5-VL-3B-Instruct",
                        help="processor 的基础模型名称")
    parser.add_argument("--max_new_tokens", type=int, default=1024,
                        help="最大生成 token 数")
    parser.add_argument("--max_samples",    type=int, default=-1,
                        help="限制评估样本数（-1 表示全部）")
    parser.add_argument("--num_gpus",       type=int, default=1,
                        help="并行 GPU 数量")
    parser.add_argument("--min_pixels",     type=int, default=200704,
                        help="图像最小像素数（训练时与此保持一致）")
    parser.add_argument("--max_pixels",     type=int, default=200704,
                        help="图像最大像素数（训练时与此保持一致）")
    args = parser.parse_args()

    # ── 路径处理 ──
    project_root = str(_project_root)
    eval_json  = args.eval_json  if os.path.isabs(args.eval_json)  else os.path.join(project_root, args.eval_json)
    image_root = args.image_root if os.path.isabs(args.image_root) else os.path.join(project_root, args.image_root)
    output_dir = args.output_dir if os.path.isabs(args.output_dir) else os.path.join(project_root, args.output_dir)
    os.makedirs(output_dir, exist_ok=True)

    # ── 加载数据 ──
    print(f"加载评估数据: {eval_json}", flush=True)
    with open(eval_json, "r") as f:
        all_samples = json.load(f)

    if args.max_samples > 0:
        all_samples = all_samples[:args.max_samples]
    print(f"共 {len(all_samples)} 条样本", flush=True)

    # ── 推理 ──
    t0 = time.time()
    if args.num_gpus > 1:
        mp.set_start_method("spawn", force=True)
        manager    = mp.Manager()
        return_dict = manager.dict()
        processes  = []
        for rank in range(args.num_gpus):
            p = mp.Process(
                target=worker_fn,
                args=(
                    rank, args.num_gpus, args.model_path, args.base_model,
                    all_samples, image_root, args.max_new_tokens,
                    args.min_pixels, args.max_pixels, return_dict,
                ),
            )
            p.start()
            processes.append(p)
        for p in processes:
            p.join()

        # 检查各 worker 退出状态
        for rank, p in enumerate(processes):
            if p.exitcode != 0:
                print(f"⚠️  GPU{rank} worker 异常退出 (exitcode={p.exitcode})", flush=True)

        # 按原始顺序合并（每个 GPU 取自 all_samples[rank::world_size]）
        # 若某个 worker 未写入 return_dict，用空列表兜底
        shards = []
        for r in range(args.num_gpus):
            if r in return_dict:
                shards.append(return_dict[r])
            else:
                print(f"⚠️  GPU{r} 结果缺失（worker 可能崩溃且未写入），该 shard 跳过", flush=True)
                # 用占位结果填充对应位置，保持顺序
                shard_size = len(all_samples[r::args.num_gpus])
                shards.append([{
                    "sample_id": all_samples[r + i * args.num_gpus].get("sample_id", ""),
                    "error": f"GPU{r} worker crashed without writing results",
                    "metrics": {"ade": None, "fde": None, "valid": False},
                    "pred_trajectory": None,
                    "gt_trajectory": all_samples[r + i * args.num_gpus].get("future waypoints", []),
                } for i in range(shard_size) if r + i * args.num_gpus < len(all_samples)])

        predictions = [None] * len(all_samples)
        for rank, shard in enumerate(shards):
            for local_i, result in enumerate(shard):
                global_i = rank + local_i * args.num_gpus
                if global_i < len(predictions):
                    predictions[global_i] = result
        predictions = [p for p in predictions if p is not None]
    else:
        evaluator = WaymoEvaluator(
            model_path=args.model_path,
            base_model=args.base_model,
            device="cuda",
            min_pixels=args.min_pixels,
            max_pixels=args.max_pixels,
        )
        predictions = evaluate_samples(
            evaluator, all_samples, image_root, args.max_new_tokens, rank=0
        )

    elapsed = time.time() - t0
    print(f"\n推理完成，耗时 {elapsed:.1f}s（{elapsed/len(predictions):.2f}s/样本）", flush=True)

    # ── 计算汇总指标 ──
    agg = aggregate_metrics(predictions)

    # ── 保存结果 ──
    output = {
        "metadata": {
            "model_path":     args.model_path,
            "eval_json":      eval_json,
            "image_root":     image_root,
            "timestamp":      datetime.now().isoformat(),
            "max_new_tokens": args.max_new_tokens,
            "num_gpus":       args.num_gpus,
            "elapsed_s":      round(elapsed, 2),
        },
        "aggregate_metrics": agg,
        "predictions":       predictions,
    }

    predictions_path = os.path.join(output_dir, "predictions.json")
    metrics_path     = os.path.join(output_dir, "metrics.json")

    with open(predictions_path, "w", encoding="utf-8") as f:
        json.dump(output, f, ensure_ascii=False, indent=2)

    with open(metrics_path, "w", encoding="utf-8") as f:
        json.dump({"metadata": output["metadata"], "aggregate_metrics": agg}, f, ensure_ascii=False, indent=2)

    # ── 打印摘要 ──
    print("\n" + "=" * 60)
    print("评估结果摘要")
    print("=" * 60)
    print(f"  总样本数     : {agg['total_samples']}")
    print(f"  推理成功     : {agg['evaluated']}")
    print(f"  推理失败     : {agg['failed']}")
    print(f"  轨迹解析成功 : {agg['valid_trajectory']}  ({agg['valid_trajectory_ratio']*100:.1f}%)")
    if agg.get("ade_mean") is not None:
        print(f"  ADE (mean)   : {agg['ade_mean']:.4f} m")
        print(f"  ADE (median) : {agg['ade_median']:.4f} m")
        print(f"  FDE (mean)   : {agg['fde_mean']:.4f} m")
        print(f"  FDE (median) : {agg['fde_median']:.4f} m")
    else:
        print("  ADE / FDE    : 无有效轨迹预测")
    print("=" * 60)
    print(f"完整结果已保存到:\n  {predictions_path}\n  {metrics_path}")


if __name__ == "__main__":
    main()
