#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
从 nyu_depth_v2_labeled.mat 导出 795 张训练集（trainNdxs）图片。

划分来源与 export_nyu654_test.py 相同：BTS utils/splits.mat。
trainNdxs 与 testNdxs（654）无重叠、并集覆盖 1..1449，可以与测试集严格隔离，
用于量化校准或训练。

只导出 RGB；GT 仍从 mat 按 split.json 的 mat_index 读取。

用法：
    ../../venv/bin/python src/tool/export_nyu_train795.py \
        --mat ../../datasets/nyu_depth_v2_labeled.mat \
        --splits datasets/nyu654_test/splits.mat \
        --out-dir datasets/nyu_train795
"""

import argparse
import json
from pathlib import Path

import cv2
import h5py
import numpy as np
import scipy.io as sio
from tqdm import tqdm


def decode_cell(f, ref):
    return ''.join(chr(c) for c in np.asarray(f[ref]).ravel())


def main():
    parser = argparse.ArgumentParser(description='导出 NYU Depth V2 训练集 795 张')
    parser.add_argument('--mat', required=True, help='nyu_depth_v2_labeled.mat')
    parser.add_argument('--splits', required=True, help='含 trainNdxs 的 splits.mat')
    parser.add_argument('--out-dir', required=True, help='输出目录')
    args = parser.parse_args()

    image_dir = Path(args.out_dir) / 'images'
    image_dir.mkdir(parents=True, exist_ok=True)

    ndxs = sio.loadmat(args.splits)['trainNdxs'].ravel().astype(int)
    if len(ndxs) != 795:
        raise ValueError(f'trainNdxs 应为 795 个，实际 {len(ndxs)}')

    records = []
    with h5py.File(args.mat, 'r') as f:
        for seq, mat_index in enumerate(tqdm(ndxs, desc='export', ncols=80)):
            i = mat_index - 1
            scene = decode_cell(f, f['scenes'][0, i])
            rgb = np.asarray(f['images'][i]).transpose(1, 2, 0)
            name = f'{seq:04d}_{scene}.png'
            if not cv2.imwrite(str(image_dir / name), cv2.cvtColor(rgb, cv2.COLOR_RGB2BGR)):
                raise RuntimeError(f'写入失败: {image_dir / name}')
            records.append({'mat_index': int(i), 'scene': scene, 'file': name})

    with open(Path(args.out_dir) / 'split.json', 'w') as fp:
        json.dump(records, fp, indent=2)
    print(f'已导出 {len(records)} 张到 {image_dir}')


if __name__ == '__main__':
    main()