import bz2
import io

import click
from PIL import Image

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

log = Logger.getLogger()

pose_choices = {
'fa':0, 'fb':0, 'hl':67.5, 'hr':-67.5, 'pl':90, 'pr':-90, 
'ql':22.5, 'qr':-22.5, 'ra':45, 'rb':15, 'rc':-15, 'rd':-45, 're':-75}

poses_left = ['hl', 'ql', 'pl', 'ra', 'rb']
poses_right = ['hr', 'qr', 'pr', 'rc', 're', 're']

@click.command()
@click.option('-i', '--input', 'opt_fp_in', required=True,
  help='Input directory')
@click.option('-o', '--output', 'opt_fp_out', required=True,
  help='Output directory')
@click.option('-a', '--angle', 'opt_angle', type=(float, float), default=(0,0),
  help='Min/max face angles')
@click.option('-t', '--threads', 'opt_threads', default=8,
  help='Number of threads')
@click.option('--flip', 'opt_flip', type=click.Choice(['r', 'l']), 
  help='Flip profile images to the R or L')
@click.pass_context
def cli(ctx, opt_fp_in, opt_fp_out, opt_angle, opt_threads, opt_flip):
  """Extracts FERET images"""
  
  from glob import glob
  from os.path import join
  from pathlib import Path
  import time
  from tqdm import tqdm
  from multiprocessing.dummy import Pool as ThreadPool
  from functools import partial

  from PIL import ImageOps
  from app.utils import file_utils

  # filter angles
  poses = [k for k, v in pose_choices.items() if \
    abs(v) >= opt_angle[0] and abs(v) <= opt_angle[1]]
  
  # glob images dir for all *ppm.bz2
  fp_ims = []
  for pose in poses:
    log.info('globbing pose: {}'.format(pose))
    fp_ims += glob(join(opt_fp_in, '**/*_{}.ppm.bz2').format(pose))
  log.info('Processing: {:,} files'.format(len(fp_ims)))
  
  # convert bz2 to png
  def pool_func(fp_im, opt_fp_out, opt_flip):
    try:
      pbar.update(1)
      im_pil = bz2_to_pil(fp_im)
      fpp_im = Path(fp_im)
      fp_out = join(opt_fp_out, '{}.png'.format(fpp_im.stem))
      fp_out = fp_out.replace('.ppm','')  # remove ppm
      if opt_flip:
        pose_code = fpp_im.stem.split('_')[-1][:2]
        # log.debug('opt_flip: {}, found: {}'.format(opt_flip, pose_code))
        if opt_flip == 'r' and pose_code in poses_right \
          or opt_flip == 'l' and pose_code in poses_left:
            im_pil = ImageOps.mirror(im_pil)
      im_pil.save(fp_out)
      return True
    except Exception as e:
      log.error('Error processing: {}, error: {}'.format(fp_im, e))
      return False

  # make output directory
  file_utils.mkdirs(opt_fp_out)

  # setup multithreading
  pbar = tqdm(total=len(fp_ims))
  pool_resize = partial(pool_func, opt_fp_out=opt_fp_out, opt_flip=opt_flip)
  pool = ThreadPool(opt_threads) 
  with tqdm(total=len(fp_ims)) as pbar:
    results = pool.map(pool_resize, fp_ims)
  pbar.close()

  # results
  log.info('Converted: {} / {} images'.format(results.count(True), len(fp_ims)))

  
# ------------------------------------------------------------------
# local utils

def bz2_to_pil(fp_src):
  with open(fp_src, 'rb') as fp:
    im_raw = bz2.decompress(fp.read())
  im_pil = Image.open(io.BytesIO(im_raw))
  return im_pil



"""

A breakdown of the images by pose is:
  Pose Angle   Images Subjects
  fa       0     1364      994
  fb       0     1358      993
  hl   +67.5     1267      917
  hr   -67.5     1320      953
  pl     +90     1312      960
  pr     -90     1363      994
  ql   +22.5      761      501
  qr   -22.5      761      501
  ra     +45      321      261
  rb     +15      321      261
  rc     -15      610      423
  rd     -45      290      236
  re     -75      290      236

  There are 13 different poses. (The orientation "right" means
facing the photographer's right.)
  fa  regular frontal image
  fb  alternative frontal image, taken shortly after the
      corresponding fa image
  pl  profile left
  hl  half left - head turned about 67.5 degrees left
  ql  quarter left - head turned about 22.5 degrees left
  pr  profile right
  hr  half right - head turned about 67.5 degrees right
  qr  quarter right - head turned about 22.5 degrees right
  ra  random image - head turned about 45 degree left
  rb  random image - head turned about 15 degree left
  rc  random image - head turned about 15 degree right
  rd  random image - head turned about 45 degree right
  re  random image - head turned about 75 degree right

"""