#!/usr/bin/env python3
"""传感器标定数据预处理和检测入口。"""

from __future__ import annotations

import argparse
import csv
import json
import math
import sys
import time
from pathlib import Path
from statistics import mean
from typing import Any

try:
    import yaml
except ImportError as exc:
    raise SystemExit("缺少 PyYAML，请先在当前 Python 环境安装 yaml 模块。") from exc


def now_us() -> int:
    return int(time.time() * 1_000_000)


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="汇总传感器标定数据质量，并可选生成相机标定目标检测结果。")
    parser.add_argument("--dataset-root", required=True, help="本次会话数据集根目录。")
    parser.add_argument("--dataset-index", default="", help="可选 dataset_index.yaml；填写后把预处理产物写回 metadata。")
    parser.add_argument("--image-index", action="append", default=[], help="images.yaml，可重复。")
    parser.add_argument("--target-detections", action="append", default=[], help="已有 target_detections.json，可重复。")
    parser.add_argument("--imu-samples", action="append", default=[], help="imu_samples.csv，可重复。")
    parser.add_argument("--imu-segments", action="append", default=[], help="imu_segments.yaml，可重复。")
    parser.add_argument("--pointcloud-index", action="append", default=[], help="pointclouds.yaml，可重复。")
    parser.add_argument("--laser-scan-index", action="append", default=[], help="laser_scans.yaml，可重复。")
    parser.add_argument("--robot-poses", action="append", default=[], help="robot_poses.csv，可重复。")
    parser.add_argument("--output-report", default="", help="预处理质量报告；默认写到 dataset_root/sensor/preprocess_quality_report.json。")
    parser.add_argument("--generate-image-detections", action="store_true", help="使用 OpenCV 从 image-index 中生成相机检测结果。")
    parser.add_argument("--output-target-detections", default="", help="生成的相机检测结果 JSON；默认写到 dataset_root/sensor/preprocessed_target_detections.json。")
    parser.add_argument("--target-type", choices=["checkerboard", "charuco"], default="checkerboard", help="相机检测目标类型。")
    parser.add_argument("--checkerboard-cols", type=int, default=0, help="棋盘格内角点列数。")
    parser.add_argument("--checkerboard-rows", type=int, default=0, help="棋盘格内角点行数。")
    parser.add_argument("--charuco-squares-x", type=int, default=0, help="ChArUco 棋盘格列数。")
    parser.add_argument("--charuco-squares-y", type=int, default=0, help="ChArUco 棋盘格行数。")
    parser.add_argument("--charuco-square-size-m", type=float, default=0.0, help="ChArUco 方格边长。")
    parser.add_argument("--charuco-marker-size-m", type=float, default=0.0, help="ChArUco marker 边长。")
    parser.add_argument("--aruco-dictionary", default="DICT_4X4_50", help="OpenCV ArUco 字典名。")
    parser.add_argument("--max-image-missing-ratio", type=float, default=0.0, help="允许图像文件缺失比例。")
    parser.add_argument("--max-imu-gyro-rps", type=float, default=20.0, help="IMU 角速度饱和检查阈值。")
    parser.add_argument("--max-imu-accel-mps2", type=float, default=80.0, help="IMU 加速度饱和检查阈值。")
    parser.add_argument("--update-dataset-index", action=argparse.BooleanOptionalAction, default=True, help="是否写回 dataset_index metadata。")
    return parser.parse_args()


def load_yaml(path: Path) -> dict[str, Any]:
    with path.open("r", encoding="utf-8") as stream:
        data = yaml.safe_load(stream) or {}
    if not isinstance(data, dict):
        raise ValueError(f"{path} 顶层必须是 YAML map。")
    return data


def write_json(path: Path, data: dict[str, Any]) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_text(json.dumps(data, ensure_ascii=False, indent=2), encoding="utf-8")


def write_yaml(path: Path, data: dict[str, Any]) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_text(yaml.safe_dump(data, sort_keys=False, allow_unicode=True), encoding="utf-8")


def resolve_path(raw_path: str, dataset_root: Path) -> Path:
    path = Path(raw_path).expanduser()
    if path.is_absolute():
        return path
    return (dataset_root / path).resolve(strict=False)


def relative_to_root(path: Path, dataset_root: Path) -> str:
    try:
        return str(path.resolve(strict=False).relative_to(dataset_root.resolve(strict=False)))
    except ValueError:
        return str(path.resolve(strict=False))


def as_float(value: Any, default: float = 0.0) -> float:
    try:
        return float(value)
    except (TypeError, ValueError):
        return default


def as_int(value: Any, default: int = 0) -> int:
    try:
        return int(value)
    except (TypeError, ValueError):
        return default


def quality_summary(values: list[float]) -> dict[str, Any]:
    if not values:
        return {"count": 0, "average": 0.0, "minimum": 0.0, "maximum": 0.0}
    return {
        "count": len(values),
        "average": mean(values),
        "minimum": min(values),
        "maximum": max(values),
    }


def read_image_index(path: Path, dataset_root: Path) -> dict[str, Any]:
    data = load_yaml(path)
    frames = data.get("frames", [])
    if not isinstance(frames, list):
        raise ValueError(f"{path}.frames 必须是列表。")
    missing_count = 0
    stamps: list[int] = []
    quality_values: list[float] = []
    for frame in frames:
        if not isinstance(frame, dict):
            continue
        stamp_us = as_int(frame.get("stamp_us"))
        if stamp_us > 0:
            stamps.append(stamp_us)
        if "quality_score" in frame:
            quality_values.append(as_float(frame.get("quality_score")))
        image_file = str(frame.get("file", "") or "")
        if image_file and not resolve_path(image_file, dataset_root).exists():
            missing_count += 1
    return {
        "file": str(path),
        "sensor_id": data.get("sensor_id", ""),
        "sensor_type": data.get("sensor_type", ""),
        "frame_id": data.get("frame_id", ""),
        "target_board_id": data.get("target_board_id", ""),
        "frame_count": len(frames),
        "missing_file_count": missing_count,
        "missing_file_ratio": missing_count / max(len(frames), 1),
        "start_us": min(stamps) if stamps else 0,
        "end_us": max(stamps) if stamps else 0,
        "quality": quality_summary(quality_values),
    }


def read_detection_file(path: Path) -> dict[str, Any]:
    with path.open("r", encoding="utf-8") as stream:
        data = json.load(stream)
    if not isinstance(data, dict):
        raise ValueError(f"{path} 顶层必须是 JSON object。")
    detections = data.get("detections", [])
    if not isinstance(detections, list):
        raise ValueError(f"{path}.detections 必须是列表。")
    detected_count = 0
    quality_values: list[float] = []
    stamps: list[int] = []
    for detection in detections:
        if not isinstance(detection, dict):
            continue
        if bool(detection.get("target_detected", False)):
            detected_count += 1
        if "quality_score" in detection:
            quality_values.append(as_float(detection.get("quality_score")))
        stamp_us = as_int(detection.get("stamp_us"))
        if stamp_us > 0:
            stamps.append(stamp_us)
    return {
        "file": str(path),
        "sensor_id": data.get("sensor_id", ""),
        "target_type": data.get("target_type", ""),
        "detection_count": len(detections),
        "target_detected_count": detected_count,
        "target_detected_ratio": detected_count / max(len(detections), 1),
        "start_us": min(stamps) if stamps else 0,
        "end_us": max(stamps) if stamps else 0,
        "quality": quality_summary(quality_values),
    }


def read_csv_rows(path: Path) -> list[dict[str, str]]:
    with path.open("r", encoding="utf-8", newline="") as stream:
        reader = csv.DictReader(stream)
        return list(reader)


def read_imu_samples(path: Path, args: argparse.Namespace) -> dict[str, Any]:
    rows = read_csv_rows(path)
    stamps: list[int] = []
    accel_norms: list[float] = []
    gyro_norms: list[float] = []
    saturation_count = 0
    for row in rows:
        stamp_us = as_int(row.get("stamp_us"))
        if stamp_us > 0:
            stamps.append(stamp_us)
        ax = as_float(row.get("ax_mps2"))
        ay = as_float(row.get("ay_mps2"))
        az = as_float(row.get("az_mps2"))
        gx = as_float(row.get("gx_rps"))
        gy = as_float(row.get("gy_rps"))
        gz = as_float(row.get("gz_rps"))
        accel_norm = math.sqrt(ax * ax + ay * ay + az * az)
        gyro_norm = math.sqrt(gx * gx + gy * gy + gz * gz)
        accel_norms.append(accel_norm)
        gyro_norms.append(gyro_norm)
        if accel_norm > args.max_imu_accel_mps2 or gyro_norm > args.max_imu_gyro_rps:
            saturation_count += 1
    return {
        "file": str(path),
        "sample_count": len(rows),
        "start_us": min(stamps) if stamps else 0,
        "end_us": max(stamps) if stamps else 0,
        "accel_norm_mps2": quality_summary(accel_norms),
        "gyro_norm_rps": quality_summary(gyro_norms),
        "saturation_count": saturation_count,
        "saturation_ratio": saturation_count / max(len(rows), 1),
    }


def read_imu_segments(path: Path) -> dict[str, Any]:
    data = load_yaml(path)
    segments = data.get("segments", [])
    if not isinstance(segments, list):
        raise ValueError(f"{path}.segments 必须是列表。")
    static_count = 0
    motion_count = 0
    durations: list[float] = []
    for segment in segments:
        if not isinstance(segment, dict):
            continue
        segment_type = str(segment.get("segment_type", "") or "")
        if segment_type == "static":
            static_count += 1
        elif segment_type == "motion":
            motion_count += 1
        start_us = as_int(segment.get("start_us"))
        end_us = as_int(segment.get("end_us"))
        if end_us > start_us > 0:
            durations.append((end_us - start_us) / 1_000_000.0)
    return {
        "file": str(path),
        "sensor_id": data.get("sensor_id", ""),
        "segment_count": len(segments),
        "static_segment_count": static_count,
        "motion_segment_count": motion_count,
        "duration_sec": quality_summary(durations),
    }


def read_frame_index(path: Path, list_key: str) -> dict[str, Any]:
    data = load_yaml(path)
    frames = data.get(list_key, [])
    if not isinstance(frames, list):
        raise ValueError(f"{path}.{list_key} 必须是列表。")
    stamps: list[int] = []
    quality_values: list[float] = []
    for frame in frames:
        if not isinstance(frame, dict):
            continue
        stamp_us = as_int(frame.get("stamp_us"))
        if stamp_us > 0:
            stamps.append(stamp_us)
        if "quality_score" in frame:
            quality_values.append(as_float(frame.get("quality_score")))
    return {
        "file": str(path),
        "sensor_id": data.get("sensor_id", ""),
        "sensor_type": data.get("sensor_type", ""),
        "frame_id": data.get("frame_id", ""),
        "frame_count": len(frames),
        "start_us": min(stamps) if stamps else 0,
        "end_us": max(stamps) if stamps else 0,
        "quality": quality_summary(quality_values),
    }


def read_robot_poses(path: Path) -> dict[str, Any]:
    rows = read_csv_rows(path)
    stamps: list[int] = []
    quality_values: list[float] = []
    for row in rows:
        stamp_us = as_int(row.get("stamp_us"))
        if stamp_us > 0:
            stamps.append(stamp_us)
        if row.get("quality_score") not in (None, ""):
            quality_values.append(as_float(row.get("quality_score")))
    return {
        "file": str(path),
        "pose_count": len(rows),
        "start_us": min(stamps) if stamps else 0,
        "end_us": max(stamps) if stamps else 0,
        "quality": quality_summary(quality_values),
    }


def import_cv2():
    try:
        import cv2
    except ImportError as exc:
        raise RuntimeError("生成图像检测结果需要安装 OpenCV Python 模块 cv2。") from exc
    return cv2


def detect_checkerboard(cv2, image_path: Path, pattern_size: tuple[int, int]) -> tuple[bool, list[dict[str, Any]]]:
    image = cv2.imread(str(image_path))
    if image is None:
        return False, []
    gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
    found, corners = cv2.findChessboardCorners(gray, pattern_size)
    if not found:
        return False, []
    criteria = (
        cv2.TERM_CRITERIA_EPS + cv2.TERM_CRITERIA_MAX_ITER,
        30,
        0.001,
    )
    corners = cv2.cornerSubPix(gray, corners, (11, 11), (-1, -1), criteria)
    result: list[dict[str, Any]] = []
    for index, corner in enumerate(corners.reshape(-1, 2)):
        result.append({"id": index, "u_px": float(corner[0]), "v_px": float(corner[1])})
    return True, result


def detect_charuco(cv2, image_path: Path, args: argparse.Namespace) -> tuple[bool, list[dict[str, Any]]]:
    if not hasattr(cv2, "aruco"):
        raise RuntimeError("当前 OpenCV 未启用 aruco 模块，无法生成 ChArUco 检测结果。")
    if args.charuco_squares_x <= 0 or args.charuco_squares_y <= 0:
        raise ValueError("ChArUco 检测必须设置 --charuco-squares-x 和 --charuco-squares-y。")
    if args.charuco_square_size_m <= 0.0 or args.charuco_marker_size_m <= 0.0:
        raise ValueError("ChArUco 检测必须设置 --charuco-square-size-m 和 --charuco-marker-size-m。")
    dictionary_id = getattr(cv2.aruco, args.aruco_dictionary, None)
    if dictionary_id is None:
        raise ValueError(f"未知 ArUco 字典：{args.aruco_dictionary}")
    dictionary = cv2.aruco.getPredefinedDictionary(dictionary_id)
    board = cv2.aruco.CharucoBoard(
        (args.charuco_squares_x, args.charuco_squares_y),
        args.charuco_square_size_m,
        args.charuco_marker_size_m,
        dictionary,
    )
    image = cv2.imread(str(image_path))
    if image is None:
        return False, []
    gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
    marker_corners, marker_ids, _ = cv2.aruco.detectMarkers(gray, dictionary)
    if marker_ids is None or len(marker_ids) == 0:
        return False, []
    _, charuco_corners, charuco_ids = cv2.aruco.interpolateCornersCharuco(
        marker_corners,
        marker_ids,
        gray,
        board,
    )
    if charuco_ids is None or charuco_corners is None:
        return False, []
    result: list[dict[str, Any]] = []
    for index, corner in zip(charuco_ids.flatten(), charuco_corners.reshape(-1, 2)):
        result.append({"id": int(index), "u_px": float(corner[0]), "v_px": float(corner[1])})
    return bool(result), result


def generate_image_detections(
    image_index_paths: list[Path],
    dataset_root: Path,
    args: argparse.Namespace,
) -> dict[str, Any]:
    if not image_index_paths:
        raise ValueError("生成图像检测结果需要至少一个 --image-index。")
    cv2 = import_cv2()
    if args.target_type == "checkerboard":
        if args.checkerboard_cols <= 0 or args.checkerboard_rows <= 0:
            raise ValueError("棋盘格检测必须设置 --checkerboard-cols 和 --checkerboard-rows。")
        pattern_size = (args.checkerboard_cols, args.checkerboard_rows)
    else:
        pattern_size = (0, 0)

    first_index = load_yaml(image_index_paths[0])
    detections: list[dict[str, Any]] = []
    for index_path in image_index_paths:
        image_index = load_yaml(index_path)
        frames = image_index.get("frames", [])
        if not isinstance(frames, list):
            continue
        for frame in frames:
            if not isinstance(frame, dict):
                continue
            source_file = str(frame.get("file", "") or "")
            image_path = resolve_path(source_file, dataset_root)
            if args.target_type == "checkerboard":
                found, corners = detect_checkerboard(cv2, image_path, pattern_size)
            else:
                found, corners = detect_charuco(cv2, image_path, args)
            detections.append({
                "stamp_us": as_int(frame.get("stamp_us")),
                "source_file": source_file,
                "target_detected": found,
                "quality_score": 1.0 if found else 0.0,
                "corners": corners,
            })

    return {
        "schema_version": 1,
        "sensor_id": first_index.get("sensor_id", ""),
        "sensor_type": first_index.get("sensor_type", ""),
        "frame_id": first_index.get("frame_id", ""),
        "target_board_id": first_index.get("target_board_id", ""),
        "target_type": args.target_type,
        "generated_timestamp_us": now_us(),
        "generated_by": "sensor_preprocess_pipeline.py",
        "detections": detections,
    }


def update_dataset_index(
    dataset_index_path: Path,
    dataset_root: Path,
    report_path: Path,
    detection_path: Path | None,
) -> None:
    if not dataset_index_path.exists():
        return
    index = load_yaml(dataset_index_path)
    metadata = index.setdefault("metadata", {})
    metadata["sensor_preprocess_report_file"] = relative_to_root(report_path, dataset_root)
    metadata["sensor_preprocess_state"] = "completed"
    metadata["sensor_preprocess_generated_timestamp_us"] = now_us()
    if detection_path is not None:
        data_inputs = index.setdefault("data_inputs", {})
        sensor_inputs = data_inputs.setdefault("sensor", {})
        detections = sensor_inputs.setdefault("target_detection_files", [])
        if isinstance(detections, str):
            detections = [detections]
        detection_file = relative_to_root(detection_path, dataset_root)
        if detection_file not in detections:
            detections.append(detection_file)
        sensor_inputs["target_detection_files"] = detections
        metadata["sensor_preprocess_target_detection_file"] = detection_file
    write_yaml(dataset_index_path, index)


def resolved_files(raw_files: list[str], dataset_root: Path) -> list[Path]:
    return [resolve_path(raw_file, dataset_root) for raw_file in raw_files]


def main() -> int:
    args = parse_args()
    dataset_root = Path(args.dataset_root).expanduser().resolve(strict=False)
    output_report = (
        Path(args.output_report).expanduser().resolve(strict=False)
        if args.output_report
        else (dataset_root / "sensor" / "preprocess_quality_report.json").resolve(strict=False)
    )

    image_index_paths = resolved_files(args.image_index, dataset_root)
    target_detection_paths = resolved_files(args.target_detections, dataset_root)
    generated_detection_path: Path | None = None
    generated_detection: dict[str, Any] | None = None

    if args.generate_image_detections:
        generated_detection = generate_image_detections(image_index_paths, dataset_root, args)
        generated_detection_path = (
            Path(args.output_target_detections).expanduser().resolve(strict=False)
            if args.output_target_detections
            else (dataset_root / "sensor" / "preprocessed_target_detections.json").resolve(strict=False)
        )
        write_json(generated_detection_path, generated_detection)
        target_detection_paths.append(generated_detection_path)

    image_summaries = [read_image_index(path, dataset_root) for path in image_index_paths]
    detection_summaries = [read_detection_file(path) for path in target_detection_paths]
    imu_sample_summaries = [read_imu_samples(path, args) for path in resolved_files(args.imu_samples, dataset_root)]
    imu_segment_summaries = [read_imu_segments(path) for path in resolved_files(args.imu_segments, dataset_root)]
    pointcloud_summaries = [read_frame_index(path, "frames") for path in resolved_files(args.pointcloud_index, dataset_root)]
    laser_scan_summaries = [read_frame_index(path, "scans") for path in resolved_files(args.laser_scan_index, dataset_root)]
    robot_pose_summaries = [read_robot_poses(path) for path in resolved_files(args.robot_poses, dataset_root)]

    warnings: list[str] = []
    for summary in image_summaries:
        if summary["missing_file_ratio"] > args.max_image_missing_ratio:
            warnings.append(
                f"{summary['file']} 图像缺失比例 {summary['missing_file_ratio']:.3f} 超过阈值。"
            )
    for summary in imu_sample_summaries:
        if summary["saturation_count"] > 0:
            warnings.append(f"{summary['file']} 存在 IMU 饱和样本：{summary['saturation_count']}。")

    report = {
        "schema_version": 1,
        "generated_timestamp_us": now_us(),
        "generated_by": "sensor_preprocess_pipeline.py",
        "dataset_root": str(dataset_root),
        "image_samples": image_summaries,
        "target_detections": detection_summaries,
        "imu_samples": imu_sample_summaries,
        "imu_segments": imu_segment_summaries,
        "pointcloud_samples": pointcloud_summaries,
        "laser_scan_samples": laser_scan_summaries,
        "robot_poses": robot_pose_summaries,
        "warnings": warnings,
        "passed": not warnings,
    }
    if generated_detection_path is not None:
        report["generated_target_detection_file"] = str(generated_detection_path)
        report["generated_target_detection_count"] = len(generated_detection["detections"] if generated_detection else [])

    write_json(output_report, report)
    if args.dataset_index and args.update_dataset_index:
        update_dataset_index(
            Path(args.dataset_index).expanduser().resolve(strict=False),
            dataset_root,
            output_report,
            generated_detection_path,
        )
    print(f"[OK] 已生成传感器预处理报告：{output_report}")
    if generated_detection_path is not None:
        print(f"[OK] 已生成相机检测结果：{generated_detection_path}")
    if warnings:
        for warning in warnings:
            print(f"[WARN] {warning}")
    return 0


if __name__ == "__main__":
    try:
        raise SystemExit(main())
    except Exception as exc:
        print(f"[错误] {exc}", file=sys.stderr)
        raise SystemExit(1)
