#!/usr/bin/env python3

import argparse

parser = argparse.ArgumentParser(description="Dump Example Images")
parser.add_argument("-c", "--camera",  help="camera channel name", default=None)
parser.add_argument("-i", "--input",   help="SwissFEL data file", default=None)
parser.add_argument("-o", "--output",  help="example images as numpy dump", default="example_images.npy")
parser.add_argument("-n", "--nimages", help="number of images", type=int, default=1000)

clargs = parser.parse_args()


import numpy as np
from sfdata import SFDataFiles


FPIC = ":FPICTURE"


def main(fn_input, fn_output, camera, nimages):
    if fn_input is None:
        imgs = mk_rand(nimages)
    else:
        imgs = mk_real(fn_input, camera, nimages)

    print("Writing to:", fn_output)
    np.save(fn_output, imgs)


def mk_rand(n):
    print("No input file specified, will write random data")
    return np.random.random((n, 100, 100))


def mk_real(fn, ch, n):
    ch = harmonize_channel(ch)
    with SFDataFiles(fn) as f:
        if ch is None:
            names = list(f.names)
            raise SystemExit(f"No camera or channel name specificed. Choose from: {names}")
        c = f[ch]
        print("Shape of raw images:", c.shape)
        imgs = c[:n]
    nreal = len(imgs)
    if nreal != n:
        print(f"Warning: Got only {nreal} images from channel {ch} in {fn}")
    return imgs

def harmonize_channel(ch):
    if ch is None: return
    if not ch.endswith(FPIC):
        ch += FPIC
    return ch



if __name__ == "__main__":
    main(clargs.input, clargs.output, clargs.camera, clargs.nimages)