import numpy as np
from PIL import Image
from blockprocessing import BlockProcessing
from argparse import ArgumentParser
from utils import save
from pathlib import Path
import imageio


def parse_args():
    p = ArgumentParser()
    p.add_argument("--img", required=True, help="PNG or TIF image path")
    p.add_argument("--out-fname", required=True, help="HDF5 output path")
    p.add_argument("--sigma", required=True, type=float, help="scale of image noise")
    return p.parse_args()


def add_noise(img, out_fname, sigma):
    """Take png or tif image, produce HDF5 file noisy version."""
    ext = Path(img).suffix
    if ext == ".png":
        img = np.array(Image.open(img), dtype=np.float32)
    elif ext == ".tif":
        img = np.array(imageio.imread(img)[:,:,0], dtype=np.float32)
    else:
        raise ValueError(f"{ext} format is not supported")
    noisy = img + np.random.normal(scale=sigma, size=img.shape)
    save(out_fname, {"data": noisy, "original": img})


if __name__ == "__main__":
    args = parse_args()
    add_noise(args.img, args.out_fname, args.sigma)
