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=192, type=int, help='Width of frames')
parser.add_argument('--overlap', default=2, type=int, help='Expand all frames by this margin')
parser.add_argument('--ratio', default=2.0, type=float, help='Pre-crop frames to this ratio')
parser.add_argument('--in_dir', help='Directory to process')
parser.add_argument('--out_dir', default='inputs', help='Directory to output crops to')
parser.add_argument('--start', default=0, type=int, help='Minimum frame number')
parser.add_argument('--end', default=99999, type=int, help='Maximum frame number')
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, old_height = img.size
  if args.ratio == (width / old_height):
    height = old_height
    top_offset = 0
  else:
    height = round(width / args.ratio)
    top_offset = round((old_height - height) / 2)

  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:
    fn, ext = os.path.splitext(os.path.basename(raw_fn))
    partz = fn.split('_')
    try:
      num = int(partz[-1].replace('.png', ''))
      if num >= opt.start and num <= opt.end:
        dataset.append((raw_fn, width, height, step, side, overlap, top_offset,))
    except:
      pass

  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, top_offset):
  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 + top_offset, xx + w, yy + h + top_offset))
      crop.save("{}/{}".format(crop_dir, fn))

if __name__ == "__main__":
  crop_dir()