#!/bin/env python
"""
process images
"""
import os
import sys
import io
import shutil
import base64
import pathlib

import numpy as np
import mediapipe as mp
from PIL import Image, ImageOps

from util import log, Map
from sdapi import postsync


params = Map({
    'src': '', # source folder
    'dst': '', # destination folder
    'extract_face': True, # extract face from image
    'extract_body': True, # extract face from image
    'clear_dst': True, # remove all files from destination at the start
    'target_size': 512, # target resolution
    'square_images': True, # should output images be squared
    'blur_samplesize': 60, # sample size to use for blur detection
    'face_score': 0.6, # min face detection score
    'face_pad': 0.05, # pad face image percentage
    'face_model': 1, # which face model to use 0/close-up 1/standard
    'face_blur_score': 1.2, # max score for face blur detection
    'body_score': 0.9, # min body detection score
    'body_visibility': 0.5, # min visibility score for each detected body part
    'body_parts': 15, # min number of detected body parts with sufficient visibility
    'body_pad': 0.2,  # pad body image percentage
    'body_model': 2, # body model to use 0/low 1/medium 2/high
    'body_blur_score': 1.6, # max score for body blur detection
    'segmentation_face': True, # segmentation enabled
    'segmentation_body': False, # segmentation enabled
    'segmentation_model': 0, # segmentation model 0/general 1/landscape
    'segmentation_background': (192, 192, 192), # segmentation background color
    'interrogate_model': 'clip' # interrogate model
})


def detect_blur(image):
    # based on <https://github.com/karthik9319/Blur-Detection/>
    bw = ImageOps.grayscale(image)
    cx, cy = image.size[0] // 2, image.size[1] // 2
    fft = np.fft.fft2(bw)
    fftShift = np.fft.fftshift(fft)
    fftShift[cy - params.blur_samplesize: cy + params.blur_samplesize, cx - params.blur_samplesize: cx + params.blur_samplesize] = 0
    fftShift = np.fft.ifftshift(fftShift)
    recon = np.fft.ifft2(fftShift)
    magnitude = np.log(np.abs(recon))
    mean = round(np.mean(magnitude), 2)
    return mean


def segmentation(image):
    with mp.solutions.selfie_segmentation.SelfieSegmentation(model_selection=params.segmentation_model) as selfie_segmentation:
        data = np.array(image)
        results = selfie_segmentation.process(data)
        condition = np.stack((results.segmentation_mask,) * 3, axis=-1) > 0.1
        background = np.zeros(data.shape, dtype=np.uint8)
        background[:] = params.segmentation_background
        data = np.where(condition, data, background) # consider using a joint bilateral filter instead of pure combine
        segmented = Image.fromarray(data)
        return segmented


def extract_face(img):
    if not params.extract_face:
        return None, True
    if img.mode == 'RGBA':
        img = img.convert('RGB')
    scale = max(img.size[0], img.size[1]) / params.target_size
    resized = img.copy()
    resized.thumbnail((params.target_size, params.target_size), Image.HAMMING)
    with mp.solutions.face_detection.FaceDetection(min_detection_confidence=params.face_score, model_selection=params.face_model) as face:
        results = face.process(np.array(resized))
    if results.detections is None:
        return None, False
    box = results.detections[0].location_data.relative_bounding_box
    x = (box.xmin - params.face_pad / 2) * resized.width
    y = (box.ymin - params.face_pad / 2)* resized.height
    w = (box.width + params.face_pad) * resized.width
    h = (box.height + params.face_pad) * resized.height
    cx = x + w / 2
    cy = y + h / 2
    l = max(w, h) / 2
    square = [scale * (cx - l), scale * (cy - l), scale * (cx + l), scale * (cy + l)]
    square = [max(square[0], 0), max(square[1], 0), min(square[2], img.width), min(square[3], img.height)]
    cropped = img.crop(tuple(square))
    if cropped.size[0] < params.target_size and cropped.size[1] < params.target_size:
        log.warning({ 'extract face': 'low resolution', 'size': [cropped.size[0], cropped.size[1]] })
        return None, True
    cropped.thumbnail((params.target_size, params.target_size), Image.HAMMING)
    if params.square_images:
        squared = Image.new('RGB', (params.target_size, params.target_size))
        squared.paste(cropped, (0, 0))
        if params.segmentation_face:
           squared = segmentation(squared)
    else:
        squared = cropped
    blur = detect_blur(squared)
    if blur > params.face_blur_score:
        log.warning({ 'extract face': 'blur check fail', 'blur': blur })
        return None, True
    else:
        log.info({ 'extract face blur': blur })

    return squared, True


def extract_body(img):
    if not params.extract_body:
        return None, True
    if img.mode == 'RGBA':
        img = img.convert('RGB')
    scale = max(img.size[0], img.size[1]) / params.target_size
    resized = img.copy()
    resized.thumbnail((params.target_size, params.target_size), Image.HAMMING)
    with mp.solutions.pose.Pose(static_image_mode=True, min_detection_confidence=params.body_score, model_complexity=params.body_model) as pose:
        results = pose.process(np.array(resized))
    if results.pose_landmarks is None:
        return None, False
    x = [resized.width * (i.x - params.body_pad / 2) for i in results.pose_landmarks.landmark if i.visibility > params.body_visibility]
    y = [resized.height * (i.y - params.body_pad / 2) for i in results.pose_landmarks.landmark if i.visibility > params.body_visibility]
    if len(x) < params.body_parts:
        log.warning({ 'extract body': 'insufficient body parts', 'detected': len(x) })
        return None, True
    w = max(x) - min(x) + resized.width * params.body_pad
    h = max(y) - min(y) + resized.height * params.body_pad
    cx = min(x) + w / 2
    cy = min(y) + h / 2
    l = max(w, h) / 2
    square = [scale * (cx - l), scale * (cy - l), scale * (cx + l), scale * (cy + l)]
    square = [max(square[0], 0), max(square[1], 0), min(square[2], img.width), min(square[3], img.height)]
    cropped = img.crop(tuple(square))
    if cropped.size[0] < params.target_size and cropped.size[1] < params.target_size:
        log.warning({ 'extract body': 'low resolution', 'size': [cropped.size[0], cropped.size[1]] })
        return None, True
    cropped.thumbnail((params.target_size, params.target_size), Image.HAMMING)
    if params.square_images:
        squared = Image.new('RGB', (params.target_size, params.target_size))
        squared.paste(cropped, (0, 0))
        if params.segmentation_body:
           squared = segmentation(squared)
    else:
        squared = cropped
    blur = detect_blur(squared)
    if blur > params.body_blur_score:
        log.warning({ 'extract body': 'blur check fail', 'blur': blur })
        return None, True
    else:
        log.info({ 'extract body blur': blur })
    return squared, True


def interrogate(img, fn):
    def encode(f):
        with io.BytesIO() as stream:
            img.save(stream, 'JPEG')
            values = stream.getvalue()
            encoded = base64.b64encode(values).decode()
            return encoded

    if params.interrogate_model is None or params.interrogate_model == '':
        return
    json = Map({ 'image': encode(img), 'model': params.interrogate_model })
    res = postsync('/sdapi/v1/interrogate', json)
    caption = res.caption if 'caption' in res else ''
    log.info({ 'interrogate': caption })
    file = fn.replace('.jpg', '.txt')
    f = open(file, 'w')
    f.write(caption)
    f.close()


i = {}
def process_file(f: str, dst: str = None):
    def save(img, f, what):
        i[what] = i.get(what, 0) + 1
        if dst is None:
            dir = os.path.dirname(f)
        else:
            dir = dst
        base = os.path.basename(f).split('.')[0]
        fn = os.path.join(dir, str(i[what]).rjust(3, '0') + '-' + what + '-' + base + '.jpg')
        # log.debug({ 'save': fn })
        img.save(fn)
        interrogate(img, fn)
        return fn

    log.info({ 'processing': f })
    try:
        image = Image.open(f)
    except Exception as err:
        log.error({ 'image': f, 'error': err })
        return

    if image.width < 512 or image.height < 512:
        log.warning({ 'skip low resolution': [image.width, image.height], 'file': f })
        return
    log.info({ 'resolution': [image.width, image.height], 'mp': round((image.width * image.height) / 1024 / 1024, 1) })

    face, ok = extract_face(image)
    if face is not None:
        fn = save(face, f, 'face')
        log.info({ 'extract face': fn })
    else:
        log.warning({ 'no face': f })

    if not ok:
        return

    body, ok = extract_body(image)
    if body is not None:
        fn = save(body, f, 'body')
        log.info({ 'extract body': fn })   
    else:
        log.warning({ 'no body': f })


def process_images(src: str, dst: str, args = None):
    params.src = src
    params.dst = dst
    if args is not None:
        params.update(args)
    log.info({ 'processing': params })
    if not os.path.isdir(src):
        log.error({ 'process': 'not a folder', 'src': src })
    else:
        if os.path.isdir(dst) and params.clear_dst:
            log.warning({ 'clear dst': dst })
            shutil.rmtree(dst) 
        pathlib.Path(dst).mkdir(parents=True, exist_ok=True)
        for root, _sub_dirs, files in os.walk(src):
            for f in files:
                process_file(os.path.join(root, f), dst)


if __name__ == '__main__':
    # log.setLevel(logging.DEBUG)
    sys.argv.pop(0)
    dst = sys.argv.pop(0)
    params.dst = dst
    log.info({ 'processing': params })
    for loc in sys.argv:
        if os.path.isfile(loc):
            process_file(loc, dst)
        elif os.path.isdir(loc):
            for root, _sub_dirs, files in os.walk(loc):
                for f in files:
                    process_file(os.path.join(root, f), dst)

"""
- add interrogate on save
- write final stats
- create final grid
"""