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)