"""
One-time utility (not part of the 01-05 pipeline): extracts a small,
illustrative slice of real measurement data into data/measurement_samples/,
so 05_infer_measurement.py has something to fall back to on a machine
without access to group storage
(/mnt/sls_det_storage/moench_data/MLXID/Samples/Measurement). Needs that
access itself to run. ~30 MB total across both targets and all three
photon counts -- a few tens of thousands of events, nowhere near enough
for a real reconstruction, just enough to demonstrate the script runs and
show roughly the right shape of output.

Needs: numpy, h5py only.
"""
from pathlib import Path
import h5py

SOURCE_DIR = Path('/mnt/sls_det_storage/moench_data/MLXID/Samples/Measurement')
OUTPUT_DIR = Path(__file__).resolve().parent.parent / 'data/measurement_samples'

TARGETS = {
    'Edge2Filters_12keV': '2603MaxIV_Edge2Filters_12keV',
    'Flat2Filters_12keV': '2603MaxIV_Flat2Filters_12keV',
}
### (cluster-file prefix, events to keep) -- sized so 2 targets x 3 photon counts totals ~30 MB
PHOTON_PATTERNS = [
    ('1Photon_CS3', 65_000),
    ('2Photon_CS6', 32_000),
    ('3Photon_CS9', 18_000),
]

if __name__ == '__main__':
    OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
    for target, folder in TARGETS.items():
        out_target_dir = OUTPUT_DIR / target
        out_target_dir.mkdir(exist_ok=True)
        for prefix, n_events in PHOTON_PATTERNS:
            src = SOURCE_DIR / folder / f'{prefix}_chunk0.h5'
            if not src.exists():
                print(f'SKIP (not found): {src}')
                continue
            with h5py.File(src, 'r') as f:
                total = f['clusters'].shape[0]
                n = min(n_events, total)
                clusters = f['clusters'][:n]
                ref_points = f['referencePoint'][:n]

            dst = out_target_dir / f'{prefix}_chunk0.h5'
            with h5py.File(dst, 'w') as out:
                out['clusters'] = clusters
                out['referencePoint'] = ref_points
            print(f'{src} ({total:,} total) -> {dst} ({n:,} events, {dst.stat().st_size / 1e6:.1f} MB)')
    print('Done.')