#!/usr/bin/env python3
"""Read one unchanged MATLAB v7.3 I/Q file using the supplied labels."""
import argparse
import csv
from pathlib import Path

import h5py
import numpy as np


def main():
    parser = argparse.ArgumentParser(description='读取教学数据包中的一个I/Q样本')
    parser.add_argument('--root', type=Path, default=Path('.'))
    parser.add_argument('--split', choices=['train', 'student_test'], default='train')
    parser.add_argument('--index', type=int, default=0)
    args = parser.parse_args()
    with (args.root / f'{args.split}_labels.csv').open(encoding='utf-8-sig', newline='') as f:
        rows = list(csv.DictReader(f))
    if not 0 <= args.index < len(rows):
        parser.error(f'index必须在0到{len(rows) - 1}之间')
    row = rows[args.index]
    path = args.root / row['relative_path']
    with h5py.File(path, 'r') as f:
        i = np.asarray(f['I'], dtype=np.float32).reshape(-1)
        q = np.asarray(f['Q'], dtype=np.float32).reshape(-1)
    iq = i.astype(np.complex64)
    iq.imag = q
    print(f'文件: {path}')
    print(f'个体: {row["individual_id"]}; 型号: {row["model_name"]}; 开集标签: {row["open_set_label"]}')
    print(f'I: {i.shape}, {i.dtype}; Q: {q.shape}, {q.dtype}')
    print(f'复数信号 I + 1j * Q: {iq.shape}, {iq.dtype}')


if __name__ == '__main__':
    main()