# Copyright (C) 2022-present Naver Corporation. All rights reserved.
# Licensed under CC BY-NC-SA 4.0 (non-commercial use only).
# 
# --------------------------------------------------------
# Extracting crops for pre-training
# --------------------------------------------------------

import os
import argparse
from tqdm import tqdm
from PIL import Image
import functools
from multiprocessing import Pool
import math


def arg_parser():
    parser = argparse.ArgumentParser('Generate cropped image pairs from image crop list')

    parser.add_argument('--crops', type=str, required=True, help='crop file')
    parser.add_argument('--root-dir', type=str, required=True, help='root directory')
    parser.add_argument('--output-dir', type=str, required=True, help='output directory')
    parser.add_argument('--imsize', type=int, default=256, help='size of the crops')
    parser.add_argument('--nthread', type=int, required=True, help='number of simultaneous threads')
    parser.add_argument('--max-subdir-levels', type=int, default=5, help='maximum number of subdirectories')
    parser.add_argument('--ideal-number-pairs-in-dir', type=int, default=500, help='number of pairs stored in a dir')
    return parser


def main(args):
    listing_path = os.path.join(args.output_dir, 'listing.txt')

    print(f'Loading list of crops ... ({args.nthread} threads)')
    crops, num_crops_to_generate = load_crop_file(args.crops)

    print(f'Preparing jobs ({len(crops)} candidate image pairs)...')
    num_levels = min(math.ceil(math.log(num_crops_to_generate, args.ideal_number_pairs_in_dir)), args.max_subdir_levels)
    num_pairs_in_dir = math.ceil(num_crops_to_generate ** (1/num_levels))

    jobs = prepare_jobs(crops, num_levels, num_pairs_in_dir)
    del crops

    os.makedirs(args.output_dir, exist_ok=True)
    mmap = Pool(args.nthread).imap_unordered if args.nthread > 1 else map
    call = functools.partial(save_image_crops, args)

    print(f"Generating cropped images to {args.output_dir} ...")
    with open(listing_path, 'w') as listing:
        listing.write('# pair_path\n')
        for results in tqdm(mmap(call, jobs), total=len(jobs)):
            for path in results:
                listing.write(f'{path}\n')
    print('Finished writing listing to', listing_path)


def load_crop_file(path):
    data = open(path).read().splitlines()
    pairs = []
    num_crops_to_generate = 0
    for line in tqdm(data):
        if line.startswith('#'):
            continue
        line = line.split(', ')
        if len(line) < 8:
            img1, img2, rotation = line
            pairs.append((img1, img2, int(rotation), []))
        else:
            l1, r1, t1, b1, l2, r2, t2, b2 = map(int, line)
            rect1, rect2 = (l1, t1, r1, b1), (l2, t2, r2, b2)
            pairs[-1][-1].append((rect1, rect2))
            num_crops_to_generate += 1
    return pairs, num_crops_to_generate


def prepare_jobs(pairs, num_levels, num_pairs_in_dir):
    jobs = []
    powers = [num_pairs_in_dir**level for level in reversed(range(num_levels))]

    def get_path(idx):
        idx_array = []
        d = idx
        for level in range(num_levels - 1):
            idx_array.append(idx // powers[level])
            idx = idx % powers[level]
        idx_array.append(d)
        return '/'.join(map(lambda x: hex(x)[2:], idx_array))

    idx = 0
    for pair_data in tqdm(pairs):
        img1, img2, rotation, crops = pair_data
        if -60 <= rotation and rotation <= 60:
            rotation = 0  # most likely not a true rotation
        paths = [get_path(idx + k) for k in range(len(crops))]
        idx += len(crops)
        jobs.append(((img1, img2), rotation, crops, paths))
    return jobs


def load_image(path):
    try:
        return Image.open(path).convert('RGB')
    except Exception as e:
        print('skipping', path, e)
        raise OSError()


def save_image_crops(args, data):
    # load images
    img_pair, rot, crops, paths = data
    try:
        img1, img2 = [load_image(os.path.join(args.root_dir, impath)) for impath in img_pair]
    except OSError as e:
        return []

    def area(sz):
        return sz[0] * sz[1]

    tgt_size = (args.imsize, args.imsize)

    def prepare_crop(img, rect, rot=0):
        # actual crop
        img = img.crop(rect)

        # resize to desired size
        interp = Image.Resampling.LANCZOS if area(img.size) > 4*area(tgt_size) else Image.Resampling.BICUBIC
        img = img.resize(tgt_size, resample=interp)

        # rotate the image
        rot90 = (round(rot/90) % 4) * 90
        if rot90 == 90:
            img = img.transpose(Image.Transpose.ROTATE_90)
        elif rot90 == 180:
            img = img.transpose(Image.Transpose.ROTATE_180)
        elif rot90 == 270:
            img = img.transpose(Image.Transpose.ROTATE_270)
        return img

    results = []
    for (rect1, rect2), path in zip(crops, paths):
        crop1 = prepare_crop(img1, rect1)
        crop2 = prepare_crop(img2, rect2, rot)

        fullpath1 = os.path.join(args.output_dir,  path+'_1.jpg')
        fullpath2 = os.path.join(args.output_dir,  path+'_2.jpg')
        os.makedirs(os.path.dirname(fullpath1), exist_ok=True)

        assert not os.path.isfile(fullpath1), fullpath1
        assert not os.path.isfile(fullpath2), fullpath2
        crop1.save(fullpath1)
        crop2.save(fullpath2)
        results.append(path)

    return results


if __name__ == '__main__':
    args = arg_parser().parse_args()
    main(args)