#!/usr/bin/env python3
"""
将 eval_dvlmad_waymo.py 输出的 predictions.json 转换为 Waymo E2E 挑战赛
submission binproto 格式。

用法：
    python kewei_scripts/make_waymo_submission.py \
        --predictions eval_outputs/my_run/checkpoint-N/predictions.json \
        --output_dir  eval_outputs/my_run/checkpoint-N/submission \
        --authors     "Your Name" \
        --affiliation "Your Institution" \
        --account     "your@email.com" \
        --method_name "my-method"
"""

import argparse
import json
import math
import os
import re
import sys
from pathlib import Path
from typing import List, Optional, Tuple

import numpy as np

# ── Waymo / TF 依赖 ─────────────────────────────────────────────────────────
try:
    import tensorflow as tf
    from waymo_open_dataset.protos import (
        end_to_end_driving_submission_pb2 as wod_e2ed_submission_pb2,
    )
except ImportError as e:
    print(f"[ERROR] 缺少 Waymo / TensorFlow 依赖: {e}")
    print("请安装: pip install waymo-open-dataset-tf-2-12-0 tensorflow")
    sys.exit(1)

# ── 轨迹解析（与 waymo_submission.py 保持一致）────────────────────────────

NUM_RE = r'[+-]?(?:\d+(?:\.\d+)?|\.\d+)(?:[eE][+-]?\d+)?'

_PAIR_RE = re.compile(
    rf'\\?\[\s*({NUM_RE})\s*\\?,\s*({NUM_RE})\s*\\?\]',
    flags=re.IGNORECASE,
)
_PAIR_RE_PLAIN = re.compile(
    rf'\[\s*({NUM_RE})\s*,\s*({NUM_RE})\s*\]',
    flags=re.IGNORECASE,
)


def extract_trajectory_from_text(blob: str) -> List[List[float]]:
    """从模型原始输出文本中提取轨迹坐标（与参考脚本逻辑一致）。"""
    blob = blob[blob.index("trajectory"):]
    inner = re.sub(r'<\|mdm_start\|>|<\|mdm_end\|>', '', blob)
    inner = (inner
             .replace(r'\"', '"')
             .replace(r'\[', '[')
             .replace(r'\]', ']')
             .replace(r'\,', ',')
             .replace(r'\+', '+'))

    pairs = _PAIR_RE.findall(inner)
    if not pairs:
        pairs = _PAIR_RE_PLAIN.findall(inner)
    if not pairs:
        raise ValueError("未能从文本中提取到任何 [x, y] 坐标对")
    return [[float(x), float(y)] for x, y in pairs]


# ── JMT 插值（与参考脚本完全一致）───────────────────────────────────────────

def _finite_diff_velocity(p: np.ndarray, t: np.ndarray) -> np.ndarray:
    p, t = np.asarray(p, float), np.asarray(t, float)
    n = len(p)
    v = np.zeros(n, float)
    if n >= 2:
        v[0]  = (p[1] - p[0]) / (t[1] - t[0])
        v[-1] = (p[-1] - p[-2]) / (t[-1] - t[-2])
    if n >= 3:
        v[1:-1] = (p[2:] - p[:-2]) / (t[2:] - t[:-2])
    return v


def _finite_diff_accel(p: np.ndarray, t: np.ndarray) -> np.ndarray:
    p, t = np.asarray(p, float), np.asarray(t, float)
    n = len(p)
    a = np.zeros(n, float)
    if n < 3:
        return a
    a[0]  = 2 * (((p[1]-p[0])/(t[1]-t[0])) - ((p[2]-p[1])/(t[2]-t[1]))) / ((t[1]-t[0]) + (t[2]-t[1]))
    a[-1] = 2 * (((p[-1]-p[-2])/(t[-1]-t[-2])) - ((p[-2]-p[-3])/(t[-2]-t[-3]))) / ((t[-1]-t[-2]) + (t[-2]-t[-3]))
    for i in range(1, n - 1):
        dt1 = t[i] - t[i-1]
        dt2 = t[i+1] - t[i]
        a[i] = 2 * (((p[i+1]-p[i])/dt2) - ((p[i]-p[i-1])/dt1)) / (dt1 + dt2)
    return a


def _jmt_coeffs(p0, v0, a0, p1, v1, a1, T) -> np.ndarray:
    A0, A1, A2 = p0, v0, a0 / 2.0
    T2, T3, T4, T5 = T**2, T**3, T**4, T**5
    M = np.array([
        [  T3,    T4,     T5],
        [3*T2,  4*T3,   5*T4],
        [6*T,  12*T2,  20*T3],
    ], float)
    b = np.array([
        p1 - (A0 + A1*T + A2*T2),
        v1 - (A1 + 2*A2*T),
        a1 - (2*A2),
    ], float)
    A3, A4, A5 = np.linalg.solve(M, b)
    return np.array([A0, A1, A2, A3, A4, A5], float)


def _eval_quintic(coeffs: np.ndarray, tau: float) -> float:
    a0, a1, a2, a3, a4, a5 = coeffs
    return (((a5*tau + a4)*tau + a3)*tau + a2)*tau**2 + a1*tau + a0


def jmt_interpolate_xy_with_start(
    p_start: Tuple[float, float],
    traj_1to5: List[List[float]],
    t_new: np.ndarray,
) -> np.ndarray:
    """
    从当前位置 p_start (t=0) 和 5 个 1Hz 路点 (t=1..5) 用 JMT 插值生成
    t_new 处的轨迹点（4Hz 时 t_new = linspace(0.25, 5.0, 20)）。

    返回 shape (len(t_new), 2)。
    """
    P = np.vstack([np.asarray(p_start, float)[None, :],
                   np.asarray(traj_1to5, float)])       # (6, 2)
    t = np.arange(0.0, 6.0)                              # t = 0,1,2,3,4,5

    vx = _finite_diff_velocity(P[:, 0], t)
    vy = _finite_diff_velocity(P[:, 1], t)
    ax = _finite_diff_accel(P[:, 0], t)
    ay = _finite_diff_accel(P[:, 1], t)

    coeffs_x, coeffs_y, seg_starts = [], [], []
    for i in range(len(t) - 1):
        T = t[i+1] - t[i]
        coeffs_x.append(_jmt_coeffs(P[i, 0], vx[i], ax[i], P[i+1, 0], vx[i+1], ax[i+1], T))
        coeffs_y.append(_jmt_coeffs(P[i, 1], vy[i], ay[i], P[i+1, 1], vy[i+1], ay[i+1], T))
        seg_starts.append(t[i])
    seg_starts = np.asarray(seg_starts)

    t_new = np.clip(np.asarray(t_new, float), 0.0, 5.0)
    X = np.empty_like(t_new)
    Y = np.empty_like(t_new)
    for k, tk in enumerate(t_new):
        i = min(np.searchsorted(seg_starts, tk, side='right') - 1, len(seg_starts) - 1)
        i = max(i, 0)
        tau = tk - seg_starts[i]
        X[k] = _eval_quintic(coeffs_x[i], tau)
        Y[k] = _eval_quintic(coeffs_y[i], tau)
    return np.stack([X, Y], axis=1)   # (N, 2)


# ── 主逻辑 ────────────────────────────────────────────────────────────────────

def load_and_convert(
    predictions_json: str,
    verbose: bool = True,
) -> Tuple[dict, int, int]:
    """
    读取 predictions.json，将每条预测转换为 (frame_name → 20-point 4Hz traj)。

    返回 (pred_map, n_ok, n_skip)。
    """
    with open(predictions_json, "r", encoding="utf-8") as f:
        obj = json.load(f)

    # 兼容两种格式：顶层 list，或 {metadata, predictions: [...]}
    if isinstance(obj, list):
        preds = obj
    elif isinstance(obj, dict):
        preds = obj.get("predictions", obj)
    else:
        raise ValueError(f"无法识别的 predictions.json 格式: {type(obj)}")

    t_4hz = np.linspace(0.25, 5.0, 20)   # 4Hz, 5s → 20 waypoints

    pred_map = {}
    n_ok = n_skip = 0

    for item in preds:
        sample_id = item.get("sample_id", "")
        if not sample_id:
            n_skip += 1
            continue

        # 优先使用已解析的 pred_trajectory
        traj = item.get("pred_trajectory")

        # 若无效则从原始文本中重新提取
        if not traj:
            raw = item.get("model_output_raw") or ""
            if not raw or "trajectory" not in raw:
                if verbose:
                    print(f"  [SKIP] {sample_id}: 无有效轨迹", flush=True)
                n_skip += 1
                continue
            try:
                traj = extract_trajectory_from_text(raw)
            except Exception as e:
                if verbose:
                    print(f"  [SKIP] {sample_id}: 文本提取轨迹失败 ({e})", flush=True)
                n_skip += 1
                continue

        # 取前 5 个路点（1Hz, t=1..5s）
        traj5 = traj[:5]
        if len(traj5) < 1:
            n_skip += 1
            continue

        # 不足 5 个路点时用最后一个点填充
        while len(traj5) < 5:
            traj5.append(traj5[-1])

        try:
            traj_4hz = jmt_interpolate_xy_with_start((0.0, 0.0), traj5, t_4hz)
        except np.linalg.LinAlgError as e:
            if verbose:
                print(f"  [SKIP] {sample_id}: JMT 插值失败 ({e})", flush=True)
            n_skip += 1
            continue

        pred_map[sample_id] = traj_4hz
        n_ok += 1

    return pred_map, n_ok, n_skip


def build_submission(
    pred_map: dict,
    output_dir: str,
    num_shards: int,
    authors: List[str],
    affiliation: str,
    account_name: str,
    method_name: str,
    description: str,
    num_model_parameters: str,
) -> List[str]:
    """
    将 pred_map 打包为 Waymo binproto submission 文件，返回写出的文件路径列表。
    """
    os.makedirs(output_dir, exist_ok=True)

    # 按 frame_name 构建 FrameTrajectoryPredictions
    frame_preds = []
    for frame_name, traj_4hz in pred_map.items():
        predicted_trajectory = wod_e2ed_submission_pb2.TrajectoryPrediction(
            pos_x=traj_4hz[:, 0].tolist(),
            pos_y=traj_4hz[:, 1].tolist(),
        )
        frame_preds.append(
            wod_e2ed_submission_pb2.FrameTrajectoryPredictions(
                frame_name=frame_name,
                trajectory=predicted_trajectory,
            )
        )

    n = len(frame_preds)
    per_shard = math.ceil(n / max(num_shards, 1))
    n_shards_actual = math.ceil(n / per_shard) if per_shard > 0 else 1

    out_files = []
    for i in range(n_shards_actual):
        start = i * per_shard
        end   = min(start + per_shard, n)
        shard = wod_e2ed_submission_pb2.E2EDChallengeSubmission(
            predictions=frame_preds[start:end]
        )
        shard.submission_type = (
            wod_e2ed_submission_pb2.E2EDChallengeSubmission.SubmissionType.E2ED_SUBMISSION
        )
        shard.authors[:] = authors
        shard.affiliation = affiliation
        shard.account_name = account_name
        shard.unique_method_name = method_name
        shard.description = description
        shard.uses_public_model_pretraining = True
        shard.num_model_parameters = num_model_parameters

        fname = os.path.join(
            output_dir,
            f"mysubmission.binproto-{i:05d}-of-{n_shards_actual:05d}",
        )
        with tf.io.gfile.GFile(fname, "wb") as fp:
            fp.write(shard.SerializeToString())
        out_files.append(fname)
        print(f"  写出分片 {i+1}/{n_shards_actual}: {fname}  ({end-start} frames)", flush=True)

    return out_files


def main():
    parser = argparse.ArgumentParser(
        description="从 predictions.json 生成 Waymo E2E 挑战赛 submission",
        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
    )
    parser.add_argument(
        "--predictions", required=True,
        help="eval_dvlmad_waymo.py 输出的 predictions.json 路径",
    )
    parser.add_argument(
        "--output_dir", default=None,
        help="submission 文件输出目录（默认与 predictions.json 同目录下的 submission/）",
    )
    parser.add_argument(
        "--num_shards", type=int, default=1,
        help="将 submission 切分为几个 binproto 分片",
    )
    parser.add_argument("--authors",     default="Anonymous",
                        help="提交者姓名（逗号分隔多人）")
    parser.add_argument("--affiliation", default="Anonymous",
                        help="所属机构")
    parser.add_argument("--account",     default="anonymous@example.com",
                        help="Waymo 账户邮箱")
    parser.add_argument("--method_name", default="dvlm",
                        help="方法唯一标识名（字母/数字/连字符）")
    parser.add_argument("--description", default="dVLM for E2E planning",
                        help="方法简短描述")
    parser.add_argument("--num_model_parameters", default="3B",
                        help="模型参数量描述（如 3B、7B）")
    parser.add_argument("--quiet", action="store_true",
                        help="不打印 per-sample 跳过信息")
    args = parser.parse_args()

    pred_path = os.path.abspath(args.predictions)
    if not os.path.isfile(pred_path):
        print(f"[ERROR] predictions 文件不存在: {pred_path}")
        sys.exit(1)

    output_dir = args.output_dir or os.path.join(os.path.dirname(pred_path), "submission")

    print("=" * 60)
    print("Waymo E2E Submission Builder")
    print("=" * 60)
    print(f"  输入  : {pred_path}")
    print(f"  输出  : {output_dir}")
    print(f"  分片数: {args.num_shards}")
    print(f"  方法名: {args.method_name}")
    print()

    # Step 1: 加载并转换预测
    print("Step 1: 加载 predictions.json …")
    pred_map, n_ok, n_skip = load_and_convert(pred_path, verbose=not args.quiet)
    print(f"  有效帧: {n_ok}  跳过: {n_skip}  合计: {n_ok + n_skip}")
    if n_ok == 0:
        print("[ERROR] 没有任何有效预测，终止。")
        sys.exit(1)

    # Step 2: 构建 submission protobuf
    print("\nStep 2: 构建 submission binproto …")
    authors = [a.strip() for a in args.authors.split(",")]
    out_files = build_submission(
        pred_map       = pred_map,
        output_dir     = output_dir,
        num_shards     = args.num_shards,
        authors        = authors,
        affiliation    = args.affiliation,
        account_name   = args.account,
        method_name    = args.method_name,
        description    = args.description,
        num_model_parameters = args.num_model_parameters,
    )

    print("\n" + "=" * 60)
    print(f"✅ 完成！共写出 {len(out_files)} 个分片文件")
    print(f"   输出目录: {output_dir}")
    print()
    print("打包上传命令（在项目根目录执行）：")
    rel_out = os.path.relpath(output_dir)
    print(f"  tar cvf submission.tar {rel_out}")
    print( "  gzip submission.tar")
    print("=" * 60)


if __name__ == "__main__":
    main()
