"""
this script will convert the original data format to tsv format, in which jpg files are converted to base64 strings and
saved in a resolution of 512(shorter side).
each contains 3 columns: input, edit, seed0_0, seed0_1, ..., seedN_0, seedN_1
each tsv file should contain 1000 samples.
"""

import argparse
import base64
import io
import json
import os
from multiprocessing import Process

from PIL import Image
from tqdm import tqdm


def save_tsv(args, i, sub_seeds_list):
    with open(os.path.join(args.output_dir, f'{str(i).zfill(4)}.tsv'), 'w') as f:
        for name, seeds in tqdm(sub_seeds_list, desc=f'processing {i}th tsv file', leave=False):
            # load prompt
            prompt = json.load(open(os.path.join(args.data_dir, name, 'prompt.json')))
            # load images
            images = [Image.open(os.path.join(args.data_dir, name, f'{seed}_{j}.jpg')).convert('RGB') for seed in seeds
                      for j in range(2)]
            # resize the shorter side to 512
            images = [im.resize((512, int(512 / im.size[0] * im.size[1])) if im.size[0] < im.size[1] else
                                (int(512 / im.size[1] * im.size[0]), 512)) for im in images]
            # encode image using base64
            for j, im in enumerate(images):
                buffer = io.BytesIO()
                im.save(buffer, format='PNG')
                images[j] = base64.b64encode(buffer.getvalue()).decode('utf-8')

            # write to tsv file
            f.write('\t'.join([
                prompt['input'].replace('\t', '').replace('\n', '').strip(),
                prompt['edit'].replace('\t', '').replace('\n', '').strip(),
                *images
            ]) + '\n')


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--data-dir', type=str, default='/path/to/clip-filtered-dataset/')
    parser.add_argument('--output-dir', type=str, default='/path/to/output-dir/')
    parser.add_argument('--max-image', type=int, default=1000)
    parser.add_argument('--num-process', type=int, default=64)

    args = parser.parse_args()

    # load seeds
    seeds_list = json.load(open(os.path.join(args.data_dir, 'seeds.json')))

    # split seeds into 1000 samples per tsv file
    seeds_list = [seeds_list[i:i + args.max_image] for i in range(0, len(seeds_list), args.max_image)]

    # save tsv files
    processes = []
    for i, sub_seeds_list in enumerate(seeds_list):
        p = Process(target=save_tsv, args=(args, i, sub_seeds_list))
        p.start()
        processes.append(p)
        if len(processes) == args.num_process:
            for p in processes:
                p.join()
            processes = []

    for p in processes:
        p.join()

    print('done.')