import os
import glob
import argparse
from shutil import rmtree
from PIL import Image
from multiprocessing import Pool, cpu_count
from dotenv import load_dotenv, find_dotenv
import subprocess
load_dotenv(find_dotenv())

# This script generates crops with a specific aspect ratio from a 360 video.
# It creates three sequences (identified by "--label")

parser = argparse.ArgumentParser()
parser.add_argument('--folder', default="./results/wood/")
parser.add_argument('--out_dir', default="./thirds/")
parser.add_argument('--label', required=True)
parser.add_argument('--vertical_offset', type=int, default=256)
# parser.add_argument('--dst_width', type=int, default=1024)
# parser.add_argument('--dst_height', type=int, default=512)
parser.add_argument('--count', type=int, default=3)
parser.add_argument('--max', type=int, default=99997)
parser.add_argument('--crop_aspect', type=float, default=1.5)
parser.add_argument('--aspect', type=float, default=3.0)
# parser.add_argument('--folder_id', type=int, required=True)
parser.add_argument('--margin', type=int, default=16)
parser.add_argument('--no_clobber', action='store_false')
opt = parser.parse_args()

src_width = 1024
src_height = 512

dst_width = int(src_width / opt.count)
dst_height = int(dst_width / opt.aspect)

crop_width = dst_width
crop_height = int(dst_width / opt.aspect * opt.crop_aspect)

x_margin = opt.margin
y_margin = opt.margin * opt.crop_aspect

crop_size = (crop_width, crop_height,)
dst_size = (dst_width + opt.margin * 2, dst_height + opt.margin * 2,)
y0 = opt.vertical_offset - crop_height / 2 - y_margin

crops = []
paths = []
out_path = os.path.join(opt.out_dir, opt.label)

if not opt.no_clobber:
  if os.path.exists(out_path):
    rmtree(out_path)

for i in range(opt.count):
  x = int(i / opt.count * src_width) - x_margin
  w = x + crop_width + x_margin + x_margin
  if w > src_width:
    w = src_width
    x = src_width - (crop_width + x_margin + x_margin)
  if x < 0:
    x = 0
  crops.append((x, y0, w, y0 + crop_height + y_margin + y_margin,))
  letter = chr(97 + i)
  path = os.path.join(out_path, letter)
  os.makedirs(path)
  paths.append(path)

max_i = opt.max
dataset = []
for i, fn in enumerate(sorted(glob.glob(os.path.join(opt.folder, '*.png')))):
  if i >= max_i:
    break
  out_fn = "frame_{:05d}.png".format(i + 1)
  if opt.no_clobber and os.path.exists(os.path.join(paths[0], out_fn)):
    continue
  dataset.append((i, fn,))

def build_thumbnail(i, fn):
  out_fn = "frame_{:05d}.png".format(i + 1)
  if (i % 100) == 0:
    print("{}...".format(i))

  image = Image.open(fn)

  for n, path in enumerate(paths):
    image.crop(crops[n]).resize(dst_size).save(os.path.join(path, out_fn))

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