import torch
import numpy as np
import json
import random
device = "cuda" if torch.cuda.is_available() else "cpu"

def pack_raw(raw):
    # pack Bayer image to 4 channels
    im = raw.raw_image_visible.astype(np.float32)
    im = np.maximum(im - 512, 0) / (16383 - 512)  # subtract the black level

    im = np.expand_dims(im, axis=2)
    img_shape = im.shape
    H = img_shape[0]
    W = img_shape[1]

    out = np.concatenate((im[0:H:2, 0:W:2, :],
                          im[0:H:2, 1:W:2, :],
                          im[1:H:2, 1:W:2, :],
                          im[1:H:2, 0:W:2, :]), axis=2)
    return out

def CV72fillCurve(img, first_index, end_index):
    with open('customs/k_table_AMBA_4k0111_update.json') as f:
        K_tableCV72_load_json = json.load(f)
    with open('customs/gaussian_table_AMBA_4k0111_update.json') as f:
        GS_tableCV72_load_json = json.load(f)
    img = torch.clone(img)
    imgPoisson = img
    imgPoisson = torch.clamp(imgPoisson, 0, 1)
    gain_index = str(random.randrange(int(first_index), int(end_index)))
    h, w, c = img.shape

    for i in range(c):
        syn_noise_R = torch.zeros([h, w], device=img.device)
        channel = str(i)
        pu_K = K_tableCV72_load_json[gain_index][channel]['K']
        pu_std = GS_tableCV72_load_json[gain_index][channel]['std']
        pu_row_std = GS_tableCV72_load_json[gain_index][channel]['row_std']# std of row noise
        generated_row = torch.normal(0, pu_row_std, size=(h, 1))
        for j in range(h):
            gen_row_noisy = torch.normal(float(generated_row[j]), torch.tensor(pu_std), size=(1, w))
            syn_noise_R[j] = gen_row_noisy
        img[:, :, i] = img[:, :, i] + (imgPoisson[:, :, i] - torch.poisson(imgPoisson[:, :, i] / pu_K) * pu_K) ## Don't add 200
        img[:, :, i] += syn_noise_R
        img[:, :, i] = torch.clamp(img[:, :, i], 0, 1)
    return img

def CV72fillCurve_np(img, first_index, end_index):
    with open('customs/k_table_AMBA_4k0111_update.json') as f:
        K_tableCV72_load_json = json.load(f)
    with open('customs/gaussian_table_AMBA_4k0111_update.json') as f:
        GS_tableCV72_load_json = json.load(f)
    img = np.copy(img)
    imgPoisson = img
    imgPoisson = np.clip(imgPoisson, 0, 1)
    gain_index = str(random.randrange(int(first_index), int(end_index)))
    h, w, c = img.shape

    for i in range(c):
        syn_noise_R = np.zeros([h, w])
        channel = str(i)
        pu_K = K_tableCV72_load_json[gain_index][channel]['K']
        pu_std = GS_tableCV72_load_json[gain_index][channel]['std']
        pu_row_std = GS_tableCV72_load_json[gain_index][channel]['row_std']# std of row noise
        generated_row = np.random.normal(0, pu_row_std, size=(h, 1))
        for j in range(h):
            gen_row_noisy = np.random.normal(float(generated_row[j]), pu_std, size=(1, w))
            syn_noise_R[j] = gen_row_noisy
        img[:, :, i] = img[:, :, i] + (imgPoisson[:, :, i] - np.random.poisson(imgPoisson[:, :, i] / pu_K) * pu_K) ## Don't add 200
        img[:, :, i] += syn_noise_R
        img[:, :, i] = np.clip(img[:, :, i], 0, 1)
    return img

def rgb2rggb(img):
    red = img[:, :, 0]
    green1 = img[:, :, 1]
    green2 = img[:, :, 1]
    blue = img[:, :, 2]
    rggb = torch.stack([red, green1, green2, blue], dim=2)
    return rggb

def rggb2rgb(img):
    red = img[:, :, 0]
    green = img[:, :, 1] / 2 + img[:, :, 2] / 2
    blue = img[:, :, 3]
    rgb = torch.stack([red, green, blue], dim=2)
    return rgb

def rggb2rgb_np(img):
    red = img[:, :, 0]
    green = img[:, :, 1] / 2 + img[:, :, 2] / 2
    blue = img[:, :, 3]
    rgb = np.stack([red, green, blue], axis=2)
    return rgb

def rgb2rggb_np(img):
    red = img[:, :, 0]
    green1 = img[:, :, 1]
    green2 = img[:, :, 1]
    blue = img[:, :, 2]
    rggb = np.stack([red, green1, green2, blue], axis=2)
    return rggb