"""Generate fixed descriptors for saved RGB splits, preserving exact labels/order."""
import argparse
import hashlib
import json
from pathlib import Path

import numpy as np
from features import descriptor

ap = argparse.ArgumentParser()
ap.add_argument('source', type=Path)
ap.add_argument('output', type=Path)
a = ap.parse_args()
a.output.mkdir(exist_ok=True, parents=True)
for kind in ('hog', 'hybrid', 'sparse_hog'):
    arrays, labels = {}, {}
    for split in ('train', 'val', 'test'):
        raw = np.load(a.source / f'{split}-original.npz')
        arrays[split] = descriptor(raw['rgb'], hybrid=kind == 'hybrid', sparse=kind == 'sparse_hog')
        labels[split] = raw['labels']
        arrays[split].tofile(a.output / f'{kind}_{split}.bin')
        np.savetxt(a.output / f'{kind}_{split}.target', labels[split], fmt='%d')
    for suffix, split in [('validation', 'val'), ('train_test', 'test')]:
        np.concatenate([arrays['train'], arrays[split]]).tofile(a.output / f'{kind}_{suffix}.bin')
        np.savetxt(a.output / f'{kind}_{suffix}.target', np.concatenate([labels['train'], labels[split]]), fmt='%d')
    metadata = {'encoding': kind, 'features': 1940 if kind == 'hybrid' else 1764, 'padded_to': 2048,
                'fixed_untrained_transform': True,
                'sizes': {s: len(x) for s, x in arrays.items()},
                'sha256': {s: hashlib.sha256(x.tobytes()).hexdigest() for s, x in arrays.items()}}
    (a.output / f'{kind}-manifest.json').write_text(json.dumps(metadata, indent=2) + '\n')
    print(json.dumps(metadata), flush=True)