import numpy as np
import h5py
from blockprocessing import BlockProcessing
from argparse import ArgumentParser
from utils import save


def parse_args():
    p = ArgumentParser()
    p.add_argument("--in-fname", required=True, help="noisy HDF5 file")
    p.add_argument("--out-fname", required=True, help="HDF5 output path")
    p.add_argument(
        "--patch-size", required=True, type=int, help="height/width of square image patches"
    )
    return p.parse_args()


def patchify(in_fname, out_fname, patch_size):
    """Take png image, produce HDF5 file with patchified version."""

    with h5py.File(in_fname, "r") as in_f:
        data = in_f["data"][...]
    bp = BlockProcessing(data, mask=np.ones_like(data), patchheight=patch_size, patchwidth=patch_size)
    bp.im2bl()
    data = bp.Y.T.astype(np.float32)
    save(out_fname, {"data": data})


if __name__ == "__main__":
    args = parse_args()
    patchify(args.in_fname, args.out_fname, args.patch_size)
