import rawpy
import numpy as np
import glob, os
import colour_demosaicing
import imageio
import argparse
from PIL import Image as PILImage
import scipy.io as scio

parser = argparse.ArgumentParser(description="data preprocess")

parser.add_argument("--camera", type=str, default="NIKON_D700", help="Camera Name")
parser.add_argument(
    "--Bayer_Pattern", type=str, default="RGGB", help="Bayer Pattern of RAW"
)
parser.add_argument(
    "--JPEG_Quality", type=int, default=90, help="Jpeg Quality of the ground truth."
)

args = parser.parse_args()
camera_name = args.camera
Bayer_Pattern = args.Bayer_Pattern
JPEG_Quality = args.JPEG_Quality

dng_path = sorted(glob.glob("/mnt/nvme2n1/hyz/data/" + camera_name + "/DNG/*.cr2"))
rgb_target_path = "/mnt/nvme2n1/hyz/data/" + camera_name + "/RGB/"
raw_input_path = "/mnt/nvme2n1/hyz/data/" + camera_name + "/RAW/"
if not os.path.isdir(rgb_target_path):
    os.mkdir(rgb_target_path)
if not os.path.isdir(raw_input_path):
    os.mkdir(raw_input_path)


def flip(raw_img, flip):
    if flip == 3:
        raw_img = np.rot90(raw_img, k=2)
    elif flip == 5:
        raw_img = np.rot90(raw_img, k=1)
    elif flip == 6:
        raw_img = np.rot90(raw_img, k=3)
    else:
        pass
    return raw_img


for path in dng_path:
    print("Start Processing %s" % os.path.basename(path))
    raw = rawpy.imread(path)
    file_name = path.split("/")[-1].split(".")[0]
    im = raw.postprocess(use_camera_wb=True, no_auto_bright=True)
    flip_val = raw.sizes.flip
    cwb = raw.camera_whitebalance
    raw_img = raw.raw_image_visible
    if camera_name == "Canon_EOS_5D":
        raw_img = np.maximum(raw_img - 127.0, 0)
    de_raw = colour_demosaicing.demosaicing_CFA_Bayer_bilinear(raw_img, Bayer_Pattern)
    de_raw = flip(de_raw, flip_val)
    rgb_img = PILImage.fromarray(im).save(
        rgb_target_path + file_name + ".jpg", quality=JPEG_Quality, subsampling=1
    )
    np.savez(raw_input_path + file_name + ".npz", raw=de_raw, wb=cwb)