#!/usr/bin/env python3
# coding: utf-8

import argparse
import glob
import json
import os
from collections import Counter


DEFAULT_JSON_ROOT = "/data1/dongying/Mono3d/G1M3/Testdata_0129"
DEFAULT_GLOB_PATTERN = os.path.join(DEFAULT_JSON_ROOT, "*", "annotations_20260320", "*.json")


def parse_args():
    parser = argparse.ArgumentParser(
        description="统计原始真值中的 camera_mea.cls 类别，并生成映射模板。"
    )
    parser.add_argument(
        "--json-root",
        default=DEFAULT_JSON_ROOT,
        help="原始标注根目录。若未显式传 glob-pattern，则会自动拼接默认 pattern。",
    )
    parser.add_argument(
        "--glob-pattern",
        default=None,
        help="JSON 搜索 pattern，优先级高于 json-root。",
    )
    parser.add_argument(
        "--output-dir",
        default=os.path.join(os.path.dirname(__file__), "category_stats_output"),
        help="统计结果输出目录。",
    )
    parser.add_argument(
        "--limit",
        type=int,
        default=0,
        help="仅统计前 N 个 JSON，0 表示统计全部。",
    )
    parser.add_argument(
        "--sample-per-category",
        type=int,
        default=5,
        help="每个类别保留的样例数量。",
    )
    return parser.parse_args()


def canonicalize_cls(raw_cls):
    if isinstance(raw_cls, str):
        return raw_cls
    if raw_cls is None:
        return "<missing>"
    try:
        return json.dumps(raw_cls, ensure_ascii=False, sort_keys=True)
    except TypeError:
        return str(raw_cls)


def build_mapping_template(categories):
    return {category: None for category in categories}


def main():
    args = parse_args()
    glob_pattern = args.glob_pattern or os.path.join(
        args.json_root, "*", "annotations_20260320", "*.json"
    )

    json_files = sorted(glob.glob(glob_pattern))
    if args.limit > 0:
        json_files = json_files[: args.limit]

    if not json_files:
        print(f"未找到标注文件，pattern: {glob_pattern}")
        return 1

    os.makedirs(args.output_dir, exist_ok=True)

    category_counter = Counter()
    type_counter = Counter()
    file_object_counter = Counter()
    category_samples = {}
    parse_errors = []

    total_objects = 0
    total_valid_camera_mea = 0

    for json_path in json_files:
        if '019b4bbf-a702-7c8d-abee-ec1ec17e66e9' in json_path or '019b4c4f-83c5-7aea-b364-cdaf93be8759' in json_path:
            continue
        try:
            with open(json_path, "r", encoding="utf-8") as f:
                data = json.load(f)
        except Exception as exc:
            parse_errors.append({"file": json_path, "error": str(exc)})
            continue

        objects = data.get("asso_list", [])
        file_object_counter[json_path] = len(objects)

        for index, obj in enumerate(objects):
            total_objects += 1
            camera_mea = obj.get("camera_mea")
            if not isinstance(camera_mea, dict):
                category = "<missing_camera_mea>"
                raw_cls = None
            else:
                total_valid_camera_mea += 1
                raw_cls = camera_mea.get("cls")
                category = canonicalize_cls(raw_cls)

            category_counter[category] += 1
            type_counter[type(raw_cls).__name__] += 1

            if category not in category_samples:
                category_samples[category] = []

            if len(category_samples[category]) < args.sample_per_category:
                category_samples[category].append(
                    {
                        "file": json_path,
                        "index_in_asso_list": index,
                        "raw_cls": raw_cls,
                        "camera_mea_keys": sorted(camera_mea.keys()) if isinstance(camera_mea, dict) else [],
                    }
                )

    sorted_categories = [category for category, _ in category_counter.most_common()]
    mapping_template = build_mapping_template(sorted_categories)

    stats_report = {
        "glob_pattern": glob_pattern,
        "total_json_files": len(json_files),
        "total_objects": total_objects,
        "total_valid_camera_mea": total_valid_camera_mea,
        "category_value_types": dict(type_counter.most_common()),
        "category_counts": [
            {"category": category, "count": count}
            for category, count in category_counter.most_common()
        ],
        "category_samples": category_samples,
        "parse_errors": parse_errors,
    }

    report_path = os.path.join(args.output_dir, "category_stats.json")
    mapping_path = os.path.join(args.output_dir, "category_mapping_template.json")

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

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

    print(f"统计完成，JSON 文件数: {len(json_files)}")
    print(f"目标总数: {total_objects}")
    print(f"有效 camera_mea 数: {total_valid_camera_mea}")
    print("类别统计结果:")
    for category, count in category_counter.most_common():
        print(f"  {category}: {count}")
    print("cls 字段类型统计:")
    for value_type, count in type_counter.most_common():
        print(f"  {value_type}: {count}")
    if parse_errors:
        print(f"解析失败文件数: {len(parse_errors)}")
    print(f"详细统计已写入: {report_path}")
    print(f"映射模板已写入: {mapping_path}")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())
