#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
从 NYU 训练集 795 张中按场景均匀抽取量化校准图片。

抽样规则：
    - 按场景类型（scene 去掉 _0001 后缀，如 kitchen、bedroom）分组；
    - 每个场景至少 1 张，其余名额按场景图片数比例分配（最大余数法）；
    - 场景内按等距索引选取，结果确定可复现。

输出：
    <out-dir>/images/*.png        校准图片
    <out-dir>/selection.json      每张图的来源 mat_index / scene / 源文件
    <out-dir>/<name>.tar          给 pulsar2 的校准包，内部为 <name>/xxx.png

用法：
    ../../venv/bin/python src/tool/make_nyu_calibration.py \
        --source-dir datasets/nyu_train795 --count 48 \
        --out-dir datasets/nyu_calibration_48
"""

import argparse
import json
import shutil
import subprocess
from collections import defaultdict
from pathlib import Path

import cv2
import numpy as np


def scene_type(scene):
    return scene.rsplit('_', 1)[0]


def allocate(groups, count):
    """最大余数法分配名额，每个场景至少 1 张。"""
    total = sum(len(v) for v in groups.values())
    if count < len(groups):
        raise ValueError(f'抽取数量 {count} 少于场景数 {len(groups)}，无法保证每场景 1 张')
    quota = {k: 1 for k in groups}
    left = count - len(groups)
    exact = {k: left * len(v) / total for k, v in groups.items()}
    for k, v in exact.items():
        quota[k] += int(np.floor(v))
    remain = count - sum(quota.values())
    order = sorted(exact, key=lambda k: exact[k] - np.floor(exact[k]), reverse=True)
    for k in order[:remain]:
        quota[k] += 1
    return quota


def main():
    parser = argparse.ArgumentParser(description='按场景均匀抽取 NYU 校准集')
    parser.add_argument('--source-dir', required=True, help='export_nyu_train795.py 的输出目录')
    parser.add_argument('--count', type=int, default=48, help='抽取张数，默认 48')
    parser.add_argument('--out-dir', required=True, help='校准集输出目录')
    parser.add_argument('--name', help='tar 包内目录名，默认取 out-dir 名')
    parser.add_argument('--ext', choices=('png', 'jpg'), default='png',
                        help='校准图片格式，默认 png；pulsar2 读取异常时可改 jpg')
    args = parser.parse_args()

    source = Path(args.source_dir)
    split = json.load(open(source / 'split.json'))
    groups = defaultdict(list)
    for record in split:
        groups[scene_type(record['scene'])].append(record)

    quota = allocate(groups, args.count)
    selected = []
    for scene, records in sorted(groups.items()):
        records = sorted(records, key=lambda r: r['file'])
        idx = np.linspace(0, len(records) - 1, quota[scene]).round().astype(int)
        selected.extend(records[i] for i in sorted(set(idx.tolist())))

    if len(selected) != args.count:
        raise RuntimeError(f'实际选出 {len(selected)} 张，与目标 {args.count} 不符')

    out_dir = Path(args.out_dir)
    image_dir = out_dir / 'images'
    image_dir.mkdir(parents=True, exist_ok=True)
    for record in selected:
        src = source / 'images' / record['file']
        if args.ext == 'png':
            dst = image_dir / record['file']
            shutil.copy2(src, dst)
        else:
            dst = image_dir / (Path(record['file']).stem + '.jpg')
            image = cv2.imread(str(src), cv2.IMREAD_COLOR)
            if image is None:
                raise RuntimeError(f'读取失败: {src}')
            cv2.imwrite(str(dst), image, [cv2.IMWRITE_JPEG_QUALITY, 95])

    with open(out_dir / 'selection.json', 'w') as fp:
        json.dump(selected, fp, indent=2)

    name = args.name or out_dir.name
    tar_path = out_dir / f'{name}.tar'
    subprocess.run(
        ['tar', '-cf', str(tar_path), '--transform', f's,^,{name}/,', '-C', str(image_dir)]
        + sorted(p.name for p in image_dir.iterdir()),
        check=True)

    print(f'已抽取 {len(selected)} 张，覆盖 {len(groups)} 个场景类型')
    for scene in sorted(groups):
        print(f'  {scene:<16} {quota[scene]:>3} 张 / 共 {len(groups[scene])}')
    print(f'校准包: {tar_path}')


if __name__ == '__main__':
    main()