#!/usr/bin/env -S sd nix shell
#!nix-shell -i python3 -p python3Packages.pyvirtualcam python3Packages.opencv4 python3Packages.onnxruntime
# vim:ft=python

import pyvirtualcam
import numpy as np
import time
import cv2
import onnxruntime as ort
import subprocess
import time
import threading
from pathlib import Path

u2net_onnx = subprocess.run([
    "nix-prefetch-url",
    "https://huggingface.co/spaces/Akbartus/U2net-with-rgba/resolve/main/u2net.onnx?download=true",
    "1bvqvh6hl753pdwx6qf2mv99qancs62xlidnxw8ng3wf8pya37ih",
    "--name", "u2net.onnx",
    "--print-path"
], stdout=subprocess.PIPE).stdout.decode('utf-8').split('\n')[1].strip()

haarcascades_xml = str(next(Path(cv2.__file__).parent.parent.parent.parent.glob('**/haarcascade_frontalface_default.xml')))
print('haarcascades', haarcascades_xml)
face_detector = cv2.CascadeClassifier(haarcascades_xml)

u2net = ort.InferenceSession(u2net_onnx)


def adjust_gamma(image, gamma=1.0):
	# build a lookup table mapping the pixel values [0, 255] to
	# their adjusted gamma values
	invGamma = 1.0 / gamma
	table = np.array([((i / 255.0) ** invGamma) * 255
		for i in np.arange(0, 256)]).astype("uint8")
	# apply gamma correction using the lookup table
	return cv2.LUT(image, table)

cap = cv2.VideoCapture(1)
cap.set(cv2.CAP_PROP_BUFFERSIZE, 0)

output_size = (cap.get(cv2.CAP_PROP_FRAME_WIDTH), cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
output_size = [int(x) for x in output_size]

padding = (min(*output_size)//5)

state=dict(
    detected_faces = [],
    target_faces = [],
    last_mask = np.ones(output_size, dtype='uint8')*255,
    last_ai_change = time.time(),
    last_frame = None,
    stop = False,
    disable_bg = False,
    disable_warp = False,
    disable_rectangle=True,
    disable_grayscale=True,
    has_new_frame=True,
    pause_processing=False
)

def update_detector_state(frame, state):
    if len(state['target_faces']) > 0:
        for i in range(4):
            # transição mais suave
            delta =  state['target_faces'][0][i] - state['detected_faces'][0][i]
            if delta == 0:
                continue
            state['detected_faces'][0][i] += delta / (abs(delta)**(1/2))

def update_detector_stuff(state):
    try:
        while not state['stop']:
            time.sleep(.2)

            if state['pause_processing']:
                time.sleep(1)
                continue

            frame = state['last_frame']
            if frame is None:
                continue
            detected_faces_candidates = []
            for face in face_detector.detectMultiScale(frame):
                [xi, yi, wf, hf] = face
                xf = xi + wf
                yf = yi + hf
                detected_faces_candidates.append([xi, yi, xf, yf])
            if len(state['detected_faces']) == 0:
                state['detected_faces'] = detected_faces_candidates
            state['target_faces'] = detected_faces_candidates
    finally:
        state['stop'] = True

def update_ai_stuff(state):
    last_run = time.time()
    try:
        while not state['stop']:
            if state['pause_processing']:
                time.sleep(1)
                continue
            interval = time.time() - last_run
            if not state['has_new_frame']:
                state['has_new_frame'] = False
                continue
            if interval < 1.0:
                time.sleep(1.0 - interval)
                continue
            last_run = time.time()
            print("Running AI round")
            frame = state['last_frame']
            if frame is None:
                continue
            if not state['disable_bg']:
                bgmodel_result = u2net.run(None, {"input_image": np.expand_dims(np.moveaxis(np.array(cv2.resize(frame, (320, 320))/255, dtype='float32'), -1, 0), 0)})
                mask = bgmodel_result[0][0][0,:,:] < 0.1
                mask = cv2.resize(np.array(mask*255, dtype='uint8'), output_size, cv2.INTER_NEAREST)
                state['last_mask'] = mask

    finally:
        state['stop'] = True

def process_image(frame, state):
    if state['pause_processing']:
        return None
    frame = adjust_gamma(frame)
    state['last_frame'] = frame

    update_detector_state(frame, state)
    detected_faces = state['detected_faces']
    last_mask = state['last_mask']
    if len(detected_faces) == 0:
        return frame
    (xi, yi, xf, yf) = detected_faces[0]
    sx = xf - xi
    sy = yf - yi
    propx = sx / output_size[1]
    propy = sy / output_size[0]
    propavg = (propx + propy) / 2
    propavg = min(propavg, 0.7)
    sx = int(propavg * output_size[0])
    sy = int(propavg * output_size[1])
    center_x = int((xi + xf) / 2)
    center_y = int((yi + yf) / 2)

    final_size = [int(sz * propavg) for sz in output_size]

    frame = adjust_gamma(frame)
    if not state['disable_rectangle']:
        cv2.rectangle(frame, (int(xi), int(yi)), (int(xf), int(yf)), (255, 0, 0), 2)
    padding_x = (final_size[0] - sx) // 2
    padding_y = (final_size[1] - sy) // 2
    yi = int(max(center_y - (sy/2), 0))
    yf = int(min(center_y + (sy/2), output_size[1]))
    xi = int(max(center_x - (sx/2), 0))
    xf = int(min(center_x + (sx/2), output_size[0]))
    mask = last_mask
    if not state['disable_warp']:
        frame_cutted = frame[yi:yf, xi:xf,:]
        mask_cutted = last_mask[yi:yf, xi:xf]
        frame = cv2.resize(frame_cutted, output_size)
        mask = cv2.resize(mask_cutted, output_size, cv2.INTER_NEAREST)

    if not state['disable_bg']:
        frame_blur = cv2.blur(frame, (23, 23))
        frame[mask > 128, :] = frame_blur[mask > 128, :]
    if not state['disable_grayscale']:
        frame_gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
        frame = np.zeros(frame.shape, dtype=frame.dtype)
        frame[:,:,0] = frame_gray
        frame[:,:,1] = frame_gray
        frame[:,:,2] = frame_gray
    state['has_new_frame'] = True
    return frame

def update_camera_stuff(state):
    try:
        with pyvirtualcam.Camera(width=output_size[0], height=output_size[1], fps=20, print_fps=True, fmt=pyvirtualcam.PixelFormat.BGR) as cam:
            while not state['stop']:
                ret, frame = cap.read()
                if not ret:
                    break
                new_frame = process_image(frame, state)
                if new_frame is not None:
                    frame = new_frame
                cam.send(frame)
                cam.sleep_until_next_frame()
    finally:
        state['stop'] = True
        cap.release()

t_det = threading.Thread(target=update_detector_stuff, args=[state])
t_det.start()
t_ai = threading.Thread(target=update_ai_stuff, args=[state])
t_ai.start()
t_cam = threading.Thread(target=update_camera_stuff, args=[state])
t_cam.start()

try:
    while True:
        line = input("CMD: ").strip()
        if line in ["bg", "b"]:
            state['disable_bg'] = not state['disable_bg']
        if line in ['warp', 'w']:
            state['disable_warp'] = not state['disable_warp']
        if line in ['rect', 'r']:
            state['disable_rectangle'] = not state['disable_rectangle']
        if line in ['gray', 'g']:
            state['disable_grayscale'] = not state['disable_grayscale']
        if line in ['p', 'pause']:
            state['pause_processing'] = not state['pause_processing']
            
except:
    state['stop'] = True
    t_det.join()
    t_ai.join()
    t_cam.join()
