import math
import multiprocessing
import re
import subprocess as sp
import time
import ffmpeg
import numpy as np
import torch
import cv2

from .bg import DEVICE, Net, iter_frames, remove_many


def worker(worker_nodes,
           worker_index,
           result_dict,
           model_name,
           gpu_batchsize,
           total_frames,
           frames_dict):
    print(F"WORKER {worker_index} ONLINE")

    output_index = worker_index + 1
    base_index = worker_index * gpu_batchsize
    net = Net(model_name)
    script_net = None
    for fi in (list(range(base_index + i * worker_nodes * gpu_batchsize,
                          min(base_index + i * worker_nodes * gpu_batchsize + gpu_batchsize, total_frames)))
               for i in range(math.ceil(total_frames / worker_nodes / gpu_batchsize))):
        if not fi:
            break

        # are we processing frames faster than the frame ripper is saving them?
        last = fi[-1]
        while last not in frames_dict:
            time.sleep(0.1)

        input_frames = [frames_dict[index] for index in fi]
        # print(np.stack(input_frames).shape)#(1, 320, 568, 3)
        if script_net is None:
            script_net = torch.jit.trace(net,
                                         torch.as_tensor(np.stack(input_frames), dtype=torch.float32, device=DEVICE))


        result_dict[output_index] = remove_many(input_frames, script_net) #(1,320,568)
        # _, result_dict[output_index][0] = cv2.threshold(result_dict[output_index][0], 120, 255, cv2.THRESH_BINARY)
        # cv2.imwrite("/Users/zihao/Desktop/zero/video/output/" + str(uuid.uuid3(uuid.NAMESPACE_DNS, "random")) + ".jpg", np.array(input_frames)[0])
        original_img = np.copy(input_frames) # (1, 320, 556, 3)
        output_copy = np.copy(result_dict[output_index][0]) # (320, 568) -> (320, 556)
        cv2.resize(output_copy, (original_img.shape[2], original_img.shape[1]))
        contours, _ = cv2.findContours(output_copy, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)
        x, y, w, h = cv2.boundingRect(contours[0])
        # upper_left = [x, y]
        # lower_right = [x + w, y + h]
        # print("ul: ")
        # print(upper_left)
        # print("lr: ")
        # print(lower_right)
        color = (255, 0, 0)
        thickness = 2
        cv2.rectangle(original_img[0], (x, y), (x+w, y+h), color, thickness,lineType=cv2.LINE_AA)
        # result_dict[output_index] = original_img
        # cv2.imwrite("/Users/zihao/Desktop/zero/video/output/" + str(uuid.uuid3(uuid.NAMESPACE_DNS, "random")) + ".jpg", original_img[0])
        # cv2.imwrite("/Users/zihao/Desktop/zero/video/output/" + str(uuid.uuid3(uuid.NAMESPACE_DNS, "random1")) + ".jpg", output_copy)
        # print(result_dict[output_index].shape) #(1, 320, 568)
        result_dict[output_index] = original_img

        # clean up the frame buffer
        for fdex in fi:
            del frames_dict[fdex]
        output_index += worker_nodes


def capture_frames(file_path, frames_dict, prefetched_samples, total_frames):
    print(F"WORKER FRAMERIPPER ONLINE")
    for idx, frame in enumerate(iter_frames(file_path)):
        frames_dict[idx] = frame
        while len(frames_dict) > prefetched_samples:
            time.sleep(0.1)
        if idx > total_frames:
            break


def parallel_greenscreen(file_path,
                         worker_nodes,
                         gpu_batchsize,
                         model_name,
                         frame_limit=-1,
                         prefetched_batches=4,
                         framerate=-1):
    manager = multiprocessing.Manager()

    results_dict = manager.dict()
    frames_dict = manager.dict()

    print(file_path)

    info = ffmpeg.probe(file_path)
    total_frames = int(info["streams"][0]["nb_frames"])

    if frame_limit != -1:
        total_frames = min(frame_limit, total_frames)

    fr = info["streams"][0]["r_frame_rate"]

    if framerate == -1:
        print(F"FRAME RATE DETECTED: {fr} (if this looks wrong, override the frame rate)")
        framerate = math.ceil(eval(fr))

    print(F"FRAME RATE: {framerate} TOTAL FRAMES: {total_frames}")

    p = multiprocessing.Process(target=capture_frames,
                                args=(file_path, frames_dict, gpu_batchsize * prefetched_batches, total_frames))
    p.start()

    # note I am deliberatley not using pool
    # we can't trust it to run all the threads concurrently (or at all)
    workers = [multiprocessing.Process(target=worker,
                                       args=(worker_nodes, wn, results_dict, model_name, gpu_batchsize, total_frames,
                                             frames_dict))
               for wn in range(worker_nodes)]
    for w in workers:
        w.start()

    command = None
    proc = None
    frame_counter = 0
    for i in range(math.ceil(total_frames / worker_nodes)):
        for wx in range(worker_nodes):

            hash_index = i * worker_nodes + 1 + wx

            while hash_index not in results_dict:
                time.sleep(0.1)

            frames = results_dict[hash_index]
            # dont block access to it anymore
            del results_dict[hash_index]

            for frame in frames:
                if command is None:
                    command = ['ffmpeg',
                               '-y',
                               '-f', 'rawvideo',
                               '-vcodec', 'rawvideo',
                               '-s', F"{frame.shape[1]}x320",
                               '-pix_fmt', 'rgb24',
                               '-r', F"{framerate}",
                               '-i', '-',
                               '-an',
                               '-vcodec', 'mpeg4',
                               '-b:v', '2000k',
                               re.sub(r"\.(mp4|mov|avi)", r".matte.\1", file_path, flags=re.I)]

                    proc = sp.Popen(command, stdin=sp.PIPE)

                proc.stdin.write(frame.tostring())
                frame_counter = frame_counter + 1

                if frame_counter >= total_frames:
                    p.join()
                    for w in workers:
                        w.join()
                    proc.stdin.close()
                    proc.wait()
                    print(F"FINISHED ALL FRAMES ({total_frames})!")
                    return

    p.join()
    for w in workers:
        w.join()
    proc.stdin.close()
    proc.wait()
