import os
import cv2
import numpy as np
from matplotlib import pyplot as plt
from PIL import Image

'''
First, please install dcraw according to the command in requirements.txt.
Second, modify the path of examples, mask_check_file and mask_file.
'''

example = "./examples/iphonepm/1.dng" 
mask_check_file = "./mask_check/iphonepm/1_mask_check.npy"
mask_file = "./mask/iphonepm/1_mask.npy"

'''
Please modify the black level (k) and saturation (S) depending on the mobile camera.
-k 256 -S 4095: Huawei P30Pro, Huawei Mate30, Xiaomi11 Pro
-k 63 -S 1023: Vivo iqoo neo5, Xiaomi13
-k 0 -S 65535: iPhone14 ProMax
'''

os.system("dcraw -v -o 0 -r 1 1 1 1 -q 3 -k 256 -S 4095 -4 -T " + example)

'''
Save dng.npy
'''
tiff = example.replace('dng', 'tiff')
raw = Image.open(tiff)
raw = np.array(raw)
save_path = tiff.replace('tiff', 'dng.npy')
np.save(save_path, raw)

'''
Read grey_block from mask_check_mat
Please modify mask_check_file
'''
#mask_check_file = "./mask_check/xiaomi11pro/1_mask_check.npy"
img = np.load(mask_check_file, allow_pickle=True)

'''
Read mask_mat to get colorchecker range
Please modify mask_file
'''
#mask_file = "./mask/xiaomi11pro/1_mask.npy"
ma = np.load(mask_file, allow_pickle=True)

'''
Resize raw and img to mask same size, get grey-block coordinates
'''
h, w, _ = ma.shape
raw = cv2.resize(raw, (w, h))
img = cv2.resize(img, (w, h))
idx1, idx2, idx3 = np.where(img == 0)
tx, ty, tz = np.where(ma == 1.0)
left, right = ty.min(), ty.max()
top, bot = tx.min(), tx.max()
new1, new2 = [], []

for i in range(len(idx1)):
    if ma[idx1[i]][idx2[i]] == 1.0 and img[idx1[i]][idx2[i]][0] == 0 and \
img[idx1[i]][idx2[i]][0] == 0 and img[idx1[i]][idx2[i]][0] == 0:
        img[idx1[i]][idx2[i]] = [255,255,255]
        new1.append(idx1[i])
        new2.append(idx2[i])
#plt.imshow(img)

'''
Get groundtruth
'''
gt = np.mean(raw[new1, new2, :], 0)
gt = gt / np.linalg.norm(gt)
print(gt)