import os
import argparse
import glob
import operator
from multiprocessing import Pool, cpu_count
from shutil import rmtree
from PIL import Image

parser = argparse.ArgumentParser()
parser.add_argument('--overlap', default=2, type=int)
parser.add_argument('--scale', default=4, type=int)
parser.add_argument('dir', metavar='dir', help='Directory to process')
args = parser.parse_args()

def split_crop_dir(fn):
  crop, x, y, w, h = os.path.basename(fn).split("_")
  return int(x), int(y), int(w), int(h)

def merge_files(file, crop_dir_list, ww, hh, overlap, merge_dir):
    fn = os.path.basename(file)
    print(fn)
    canvas = Image.new('RGB', (ww, hh,))
    for crop_dir_tuple in crop_dir_list:
      x, y, w, h = crop_dir_tuple
      crop_dir = "crop_{}_{}_{}_{}".format(x, y, w, h)
      image = Image.open("{}/{}/{}".format(args.dir, crop_dir, fn))
      crop = image.crop((overlap, overlap, image.size[0] - overlap * 2, image.size[1] - overlap * 2))
      canvas.paste(image, (x * args.scale, y * args.scale,))
    canvas.save("{}/{}".format(merge_dir, fn))

def merge_dir():
  overlap = args.overlap * args.scale
  crop_dirs = glob.glob("{}/crop_*".format(args.dir))
  merge_dir = "{}/merged".format(args.dir)
  os.makedirs(merge_dir, exist_ok=True)
  ww = 0
  hh = 0
  masks = {}
  crop_dir_list = []
  widths = []
  heights = []
  for crop_dir in crop_dirs:
    x, y, w, h = split_crop_dir(crop_dir)
    if x == 0:
      hh = max(y + h - args.overlap * 2, hh)
    if y == 0:
      ww = max(x + w - args.overlap * 2, ww) 
    crop_dir_list.append((x, y, w, h,))
  crop_dir_list = sorted(crop_dir_list, key=operator.itemgetter(0, 1))

  print("{}x{}".format(ww, hh))

  ww *= args.scale
  hh *= args.scale

  files = sorted(glob.glob("{}/*.png".format(crop_dirs[0])))
  dataset = []
  for file in files:
    dataset.append((file, crop_dir_list, ww, hh, overlap, merge_dir,))

  chunksize = 3
  with Pool(processes=cpu_count()) as pool:
    pool.starmap(merge_files, dataset, chunksize)

if __name__ == "__main__":
  merge_dir()