import os
from tqdm import tqdm
import cv2
import numpy as np
import albumentations as A
import random
import shutil
import argparse

# path parameters
parser = argparse.ArgumentParser()

parser.add_argument('--data_dir',
                    type=str,
                    help='Raw training data.',
                    default="raw_data")

transform = A.Compose([
    A.OneOf([
            A.ISONoise(p=0.4),
            A.JpegCompression(quality_lower=50, quality_upper=70,
                              always_apply=False, p=0.8),
            ], p=0.6),
    A.OneOf([
            A.MotionBlur(blur_limit=11, p=.8),
            A.MedianBlur(blur_limit=3, p=0.75),
            A.GaussianBlur(blur_limit=7, p=0.75),
            ], p=0.8),
    A.OneOf([
            A.RandomBrightnessContrast(
                brightness_limit=0.3, contrast_limit=0.3, p=0.75),
            A.RandomShadow(num_shadows_lower=1,
                           num_shadows_upper=18, shadow_dimension=6, p=0.85),
            ], p=0.8),
])


def getListOfFiles(dirName):
    print(dirName)
    # create a list of file and sub directories
    # names in the given directory
    listOfFile = os.listdir(dirName)
    allFiles = list()
    # Iterate over all the entries
    for entry in listOfFile:
        allFiles.append(entry)
    return allFiles


def ImageResize(image, factor=0.6):
    width = int(image.shape[1] * factor)
    height = int(image.shape[0] * factor)
    dim = (width, height)
    # print(image.shape)
    resized = cv2.resize(image, dim, interpolation=cv2.INTER_LANCZOS4)
    # print(resized.shape)
    return resized


def GetOverlappingBlocks(im, M=256, N=256, Part=8):
    tiles = []
    tile = np.zeros((M, N, 3), dtype=np.uint8)
    #tile[:,:,2] = 255

    x = 0
    y = 0
    x_start = 0
    y_start = 0
    while y < im.shape[0]:
        while x < im.shape[1]:
            if(x != 0):
                x_start = x - int(N/Part)
            if(y != 0):
                y_start = y - int(M/Part)
            if(y_start+M > im.shape[0]):
                if(x_start+N > im.shape[1]):
                    tile[0:im.shape[0]-y_start, 0:im.shape[1]-x_start, :] = im[y_start:im.shape[0], x_start:im.shape[1], :]
                else:
                    tile[0:im.shape[0]-y_start, 0:N, :] = im[y_start:im.shape[0], x_start:x_start+N, :]
            else:
                if(x_start+N > im.shape[1]):
                    tile[0:M, 0:im.shape[1]-x_start, :] = im[y_start:y_start+M, x_start:im.shape[1], :]
                else:
                    tile[0:M, 0:N, :] = im[y_start:y_start + M, x_start:x_start+N, :]

            #pre_tile = cv2.cvtColor(PreProcessInput(cv2.cvtColor(tile, cv2.COLOR_RGB2BGR)), cv2.COLOR_BGR2RGB)
            # tiles.append(load_tf_img(pre_tile,M))

            # tiles.append(load_tf_img(tile,M))
            tiles.append(tile)

            tile = np.zeros((M, N, 3), dtype=np.uint8)
            #tile[:,:,2] = 255
            x = x_start + N
        y = y_start + M
        x = 0
        x_start = 0
    return tiles


def GenerateTrainingBlocks(data_folder, gt_folder, dataset_path='./dataset', M=256, N=256):
    print(data_folder)
    print('Generating training blocks!!!')
    train_path = dataset_path + '/' + data_folder + '_Trainblocks'

    if not os.path.exists(train_path):
        os.makedirs(train_path)

    train_filenames = train_path + '/train_block_names.txt'
    f = open(train_filenames, 'w')

    data_path = data_folder
    gt_path = gt_folder
    # data_path = dataset_path + '/' + data_folder
    # gt_path = dataset_path + '/' + gt_folder

    print(data_path)

    filenames = getListOfFiles(data_path)
    cnt = 0
    print(filenames)
    for name in tqdm(filenames):
        print(name)
        if name == '.DS_Store':
            continue
        arr = name.split(".")
        gt_filename = gt_path + '/' + arr[0] + "_mask."+arr[1]
        in_filename = data_path + '/' + name
        print(gt_filename)
        print(in_filename)
        gt_image_initial = cv2.imread(gt_filename)
        in_image_initial = cv2.imread(in_filename)
        if gt_image_initial.shape[0] + gt_image_initial.shape[1] > in_image_initial.shape[0]+in_image_initial.shape[1]:
            gt_image_initial = cv2.resize(gt_image_initial, (in_image_initial.shape[1], in_image_initial.shape[0]))
        else:
            in_image_initial = cv2.resize(in_image_initial, (gt_image_initial.shape[1], gt_image_initial.shape[0]))
        print(gt_image_initial.shape, in_image_initial.shape)
        # cv2.imshow("img", in_image_initial)
        # cv2.imshow("gt", gt_image_initial)
        # cv2.waitKey(0)
        # cv2.destroyAllWindows()
        for scale in [0.7, 1.0, 1.4]:
            gt_image = ImageResize(gt_image_initial, scale)
            in_image = ImageResize(in_image_initial, scale)
            h, w, c = in_image.shape
            gt_img = GetOverlappingBlocks(gt_image, Part=8)
            in_img = GetOverlappingBlocks(in_image, Part=8)
            for i in range(len(gt_img)):
                train_img_path = train_path + '/block_' + str(cnt) + '.png'
                gt_img_path = train_path + '/gtblock_' + str(cnt) + '.png'
                cv2.imwrite(train_img_path, in_img[i])
                # cv2.imwrite(train_img_path,PreProcessInput(in_img[i]))
                cv2.imwrite(gt_img_path, gt_img[i])
                t_name = 'block_' + str(cnt) + '.png'
                f.write(t_name)
                f.write('\n')
                cnt += 1
            Random_Block_Number_PerImage = int(len(gt_img)/5)
            for i in range(Random_Block_Number_PerImage):

                if(in_image.shape[0]-M > 1 and in_image.shape[1]-N > 1):
                    y = random.randint(1, in_image.shape[0]-M)
                    x = random.randint(1, in_image.shape[1]-N)
                    in_part_img = in_image[y:y+M, x:x+N, :].copy()
                    gt_part_img = gt_image[y:y+M, x:x+N, :].copy()
                    train_img_path = train_path + '/block_' + str(cnt) + '.png'
                    gt_img_path = train_path + '/gtblock_' + str(cnt) + '.png'
                    in_part_img = cv2.cvtColor(in_part_img, cv2.COLOR_BGR2RGB)
                    augmented_image = transform(image=in_part_img)['image']
                    augmented_image = cv2.cvtColor(
                        augmented_image, cv2.COLOR_RGB2BGR)

                    cv2.imwrite(train_img_path, augmented_image)
                    cv2.imwrite(gt_img_path, gt_part_img)
                    t_name = 'block_' + str(cnt) + '.png'
                    f.write(t_name)
                    f.write('\n')
                    cnt += 1
                else:
                    break
                    in_part_img = np.zeros((M, N, 3), dtype=np.uint8)
                    gt_part_img = np.zeros((M, N, 3), dtype=np.uint8)
                    in_part_img[:, :, :] = 255
                    gt_part_img[:, :, :] = 255

                    if(in_image.shape[0]-M <= 1 and in_image.shape[1]-N > 1):
                        y = 0
                        x = random.randint(1, in_image.shape[1]-N)
                        in_part_img[:h, :, :] = in_image[:, x:x+N, :].copy()
                        gt_part_img[:h, :, :] = gt_image[:, x:x+N, :].copy()
                    if(in_image.shape[0]-M > 1 and in_image.shape[1]-N <= 1):
                        x = 0
                        y = random.randint(1, in_image.shape[0]-M)
                        in_part_img[:, :w, :] = in_image[y:y+M, :, :].copy()
                        gt_part_img[:, :w, :] = gt_image[y:y+M, :, :].copy()

                    train_img_path = train_path + '/block_' + str(cnt) + '.png'
                    gt_img_path = train_path + '/gtblock_' + str(cnt) + '.png'
                    in_part_img = cv2.cvtColor(in_part_img, cv2.COLOR_BGR2RGB)
                    augmented_image = transform(image=in_part_img)['image']
                    augmented_image = cv2.cvtColor(
                        augmented_image, cv2.COLOR_RGB2BGR)

                    cv2.imwrite(train_img_path, augmented_image)
                    cv2.imwrite(gt_img_path, gt_part_img)
                    t_name = 'block_' + str(cnt) + '.png'
                    f.write(t_name)
                    f.write('\n')
                    cnt += 1
        # print(cnt)

    f.close()

    print('Total number of training blocks generated: ', cnt)

    return train_path, train_filenames

def CombineToImage(imgs,h,w,ch,Part=8):
    Image = np.zeros((h,w,ch),dtype=np.float32)
    Image_flag = np.zeros((h,w),dtype=bool)
    i = 0
    j = 0
    i_end = 0
    j_end = 0
    for k in range(len(imgs)):
        #part_img = np.copy(imgs[k,:,:,:])
        part_img = np.copy(imgs[k])
        hh,ww,cc = part_img.shape
        i_end = min(h,i + hh)
        j_end = min(w,j + ww)
        
        
        for m in range(hh):
            for n in range(ww):
                if(i+m<h):
                    if(j+n<w):
                        if(Image_flag[i+m,j+n]):
                            
                            Image[i+m,j+n,:] = (Image[i+m,j+n,:] + part_img[m,n,:])/2
                        else:
                            Image[i+m,j+n,:] = np.copy(part_img[m,n,:])

        Image_flag[i:i_end,j:j_end] = True
        j =  min(w-1,j + ww - int(ww/Part))
        #print(i,j,w)
        #print(k,len(imgs))
        if(j_end==w):
            j = 0
            i = min(h-1,i + hh - int(hh/Part))
    Image = Image*255.0
    return Image.astype(np.uint8)
if __name__ == "__main__":
    # img = cv2.imread("raw_data/gt/189.jpg")
    args = parser.parse_args()
    data_folder = f"{args.data_dir}/imgs"
    gt_folder = f"{args.data_dir}/gt"
    dataset = "dataset"
    shutil.rmtree(dataset, ignore_errors=True)
    os.mkdir(dataset)
    train_path, train_filenames = GenerateTrainingBlocks(
        data_folder=data_folder, gt_folder=gt_folder, dataset_path=dataset)