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('--step', default=256, type=int)
parser.add_argument('--overlap', default=2, type=int)
parser.add_argument('--in_dir', help='Directory to process')
parser.add_argument('--out_dir', default='inputs', help='Directory to output to')
args = parser.parse_args()

def crop_dir():
  step = args.step
  overlap = args.overlap
  side = step + overlap * 2
  files = glob.glob('{}/*.png'.format(args.in_dir))
  # print(files)

  img = Image.open(files[0])
  width, height = img.size
  print("{}x{} {}".format(width, height, step))
  for x in range(0, width, step):
    for y in range(0, height, step):
      xx = max(0, x - overlap)
      yy = max(0, y - overlap)
      w = min(side, width - xx)
      h = min(side, height - yy)
      crop_dir = "{}/crop_{}_{}_{}_{}".format(args.out_dir, x, y, w, h)
      if os.path.exists(crop_dir):
        rmtree(crop_dir)
      os.makedirs(crop_dir)

  dataset = []
  for raw_fn in files:
    dataset.append((raw_fn, width, height, step, side, overlap,))

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

def crop_image(raw_fn, width, height, step, side, overlap):
  img = Image.open(raw_fn)
  fn = os.path.basename(raw_fn)
  for x in range(0, width, step):
    for y in range(0, height, step):
      xx = max(0, x - overlap)
      yy = max(0, y - overlap)
      w = min(side, width - xx)
      h = min(side, height - yy)
      crop_dir = "{}/crop_{}_{}_{}_{}".format(args.out_dir, x, y, w, h)
      # print("{}x{} {}x{}".format(x, y, w, h))
      crop = img.crop((xx, yy, xx + w, yy + h))
      crop.save("{}/{}".format(crop_dir, fn))

if __name__ == "__main__":
  crop_dir()