import os
import numpy as np
import random
import torch
from torch.utils.data import Dataset
from torchvision import transforms
import cv2
import sys
from pathlib import Path
import glob
sys.path.append('..')
# from utils_ours import util
from dataloader.data_utils import rggb_raw, random_crop, Rawread, path_replace, bayer2raw, read_rawpng
from natsort import ns, natsorted
import time
from tqdm import tqdm
from multiprocessing import Pool
import pdb
import json
from fractions import Fraction
from pathlib import Path
from json import JSONEncoder
from exifread.utils import Ratio
    
class imageSet(Dataset):
    def __init__(self,args):
        super(imageSet).__init__()
        self.args = args
        self.test_dir= args.test_dir
        self.debug = args.debug

        self.paths = []
        for file in os.listdir(self.test_dir):
            if '.png' in file:
                self.img_path = os.path.join(self.test_dir, file)
                self.json_path = self.img_path.replace('png', 'json')
                self.paths.append(dict([(f"{'img_path'}", self.img_path), (f"{'json_path'}", self.json_path)]))
      
    def __getitem__(self, index):
        img_path = self.paths[index]['img_path']
        json_path = self.paths[index]['json_path']

        metadata = json_read(json_path, object_hook=fraction_from_json)

        input_img = read_rawpng(img_path, metadata)

        return {'input': input_img, 'json_path': json_path}
    
    def __len__(self):
        return len(self.paths)


def normalize(raw_image, black_level, white_level):
    if type(black_level) is list and len(black_level) == 1:
        black_level = float(black_level[0])
    if type(white_level) is list and len(white_level) == 1:
        white_level = float(white_level[0])
    black_level_mask = black_level
    if type(black_level) is list and len(black_level) == 4:
        if type(black_level[0]) is Ratio:
            black_level = ratios2floats(black_level)
        if type(black_level[0]) is Fraction:
            black_level = fractions2floats(black_level)
        black_level_mask = np.zeros(raw_image.shape)
        idx2by2 = [[0, 0], [0, 1], [1, 0], [1, 1]]
        step2 = 2
        for i, idx in enumerate(idx2by2):
            black_level_mask[idx[0]::step2, idx[1]::step2] = black_level[i]
    normalized_image = raw_image.astype(np.float32) - black_level_mask
    # if some values were smaller than black level
    normalized_image[normalized_image < 0] = 0
    normalized_image = normalized_image / (white_level - black_level_mask)
    return normalized_image

def ratios2floats(ratios):
    floats = []
    for ratio in ratios:
        floats.append(float(ratio.num) / ratio.den)
    return floats

def fractions2floats(fractions):
    floats = []
    for fraction in fractions:
        floats.append(float(fraction.numerator) / fraction.denominator)
    return floats
    
def json_read(fname, **kwargs):
    with open(fname) as j:
        data = json.load(j, **kwargs)
    return data

def fraction_from_json(json_object):
    if 'Fraction' in json_object:
        return Fraction(*json_object['Fraction'])
    return json_object