"""
Crop images to prepare for training
"""

import click

from app.settings import types
from app.utils import click_utils
from app.settings import app_cfg as cfg

@click.command()
@click.option('-i', '--input', 'opt_fp_in', required=True,
  help='Input CSV')
@click.option('--media', 'opt_dir_media', required=True,
  help='Input image/video directory')
@click.option('-o', '--output', 'opt_dir_out', required=True,
  help='Output directory for extracted ROI images')
@click.option('--size', 'opt_size', 
  type=(int, int), default=(300, 300),
  help='Output image size')
@click.option('--slice', 'opt_slice', type=(int, int), default=(None, None),
  help='Slice list of files')
@click.option('--padding', 'opt_padding', default=0,
  help='Facial padding')
@click.option('--ext', 'opt_ext_out', default='jpg', type=click.Choice(['jpg', 'png']),
  help='Output image type')
@click.pass_context
def cli(ctx, opt_fp_in, opt_dir_media, opt_dir_out, opt_size, opt_slice, 
  opt_padding, opt_ext_out):
  """Extrace ROIs to images"""
  
  import os
  from os.path import join
  from pathlib import Path
  from glob import glob
  
  from tqdm import tqdm
  import numpy as np
  from PIL import Image, ImageOps, ImageFilter, ImageDraw
  import cv2 as cv
  import pandas as pd

  from app.utils import logger_utils, file_utils, im_utils
  from app.models.bbox import BBox
  # -------------------------------------------------
  # process here
  log = logger_utils.Logger.getLogger()
  
  df_rois = pd.read_csv(opt_fp_in)
  if opt_slice:
    df_rois = df_rois[opt_slice[0]:opt_slice[1]]
  
  log.info('Processing {:,} rows'.format(len(df_rois)))

  file_utils.mkdirs(opt_dir_out)

  df_rois_grouped = df_rois.groupby(['fn'])  # group by fn/filename
  groups = df_rois_grouped.groups

  for group in groups:
    
    # get image
    group_rows = df_rois_grouped.get_group(group)

    row = group_rows.iloc[0]
    fp_im = join(opt_dir_media, '{fn}{ext}'.format(**row))  #TODO change to ext
    im = Image.open(fp_im)

    
    for idx, roi in group_rows.iterrows():
      log.info('{}'.format(roi['fn']))
      # get bbox to im dimensions
      xywh = [roi['x'], roi['y'], roi['w'] , roi['h']]
      bbox = BBox.from_xywh(*xywh)
      dim = im.size
      bbox_dim = bbox.to_dim(dim)
      # expand
      bbox_dim_exp = bbox_dim.expand_dim(opt_padding, dim)
      # crop
      x1y2 = bbox_dim_exp.pt_tl + bbox_dim_exp.pt_br
      im_crop = im.crop(box=x1y2)
      # save
      idx_zpad = file_utils.zpad(idx, zeros=3)
      fp_im_out = join(opt_dir_out, '{}_{}.{}'.format(roi['fn'], idx_zpad, opt_ext_out))
      im_crop.save(fp_im_out)