"""
Face and hand detection inference on multiple videos, forming a mosaic output
"""
import os
import cv2
import numpy as np

from modules.yolov5 import YOLOv5
from modules.inference_utils import preprocess, draw_boxes
from modules.config import CFG


# target resolution and aspect ratio
RES = (1280, 720)
SUB_RES = (RES[0]//2, RES[1]//2)
AR = round(RES[0] / RES[1], 5)


def postprocess(prediction):
    """Postprocess predictions, keeping <0, 1> value range

    Args:
        prediction (tuple): detection outputs - boxes, scores, classes

    Returns:
        np.ndarray, np.array, np.array: postprocessed outputs - boxes, scores, classes
    """
    boxes = prediction[0][0, :]
    scores = prediction[1][0, :]
    classes = prediction[2][0, :]
    # mask predictions with zero score
    mask = scores != 0
    boxes = boxes[mask, :]
    scores = scores[mask]
    classes = classes[mask].astype(np.int32)
    return boxes, scores, classes


def main(model_cfg, vid_paths):
    """Process videos and create an output video mosaic from them

    Args:
        model_cfg (str): model configuration, see model implementations for implemented configurations
        vid_paths (list): list of input video paths
    """
    # load model and trained weights
    model = YOLOv5(
        (CFG.COMBINED_PARAMS.IMG_SHAPE[0], CFG.COMBINED_PARAMS.IMG_SHAPE[1], 3),
        2,
        CFG.COMBINED_PARAMS.YOLO_ANCHORS,
        model=model_cfg,
        training=False)
    model.load_weights(CFG.COMBINED_WEIGHTS_PATH[model_cfg]).expect_partial()
    # make video capture objects for each video
    caps = [cv2.VideoCapture(vid) for vid in vid_paths]
    # construct writer object
    codec = cv2.VideoWriter_fourcc(*'mp4v')
    out_path = os.path.join(CFG.BASE_PATH, 'data', 'videos', f'out_{model_cfg}.mp4')
    out = cv2.VideoWriter(out_path, codec, 60, RES)
    # bbox coefficients to correct bbox' relative positions in mosaic
    bbox_coefs = [
        [0, 0],
        [0.5, 0],
        [0, 0.5],
        [0.5, 0.5],
    ]
    # canvas for frames
    combined_frame = np.zeros((RES[1], RES[0], 3), dtype=np.uint8)
    while True:
        rets = [None] * len(caps)
        # for each capture object
        for i, cap in enumerate(caps):
            # read frame and calculate its position on the mosaic
            ret, frame = cap.read()
            x, y = i % 2, i // 2
            rets[i] = ret
            # if video ended, write black frame
            if not ret:
                frame = np.zeros((SUB_RES[1], SUB_RES[0], 3), dtype=np.uint8)
                boxes, scores, classes = None, None, None
            else:
                # make model prediction
                preprocessed = preprocess(frame, CFG.COMBINED_PARAMS.IMG_SHAPE)
                output = model.predict(preprocessed)
                boxes, scores, classes = postprocess(output)
                # correct aspect ratio with black padding
                if round(frame.shape[1] / frame.shape[0], 5) != AR:
                    if frame.shape[1] > frame.shape[0]:
                        left_right = int(frame.shape[1] / AR / 2)
                        padding = (0, 0, left_right, left_right)
                    else:
                        top_bot = int(frame.shape[0] / AR / 2)
                        padding = (top_bot, top_bot, 0, 0)
                    boxes[:, :2] = boxes[:, :2] * frame.shape[1::-1] + [padding[2], padding[0]]
                    frame = cv2.copyMakeBorder(frame, padding[0], padding[1], padding[2], padding[3], cv2.BORDER_CONSTANT)
                    boxes[:, :2] = boxes[:, :2] / frame.shape[1::-1]
                # correct bbox' positions in mosaic
                boxes[:, :2] = boxes[:, :2] * 0.5 + bbox_coefs[i]
                boxes[:, 2:] *= 0.5
                # scale boxes to image space
                boxes = boxes * np.tile(combined_frame.shape[1::-1], 2)
            # resize image and write it to canvas
            frame = cv2.resize(frame, SUB_RES, cv2.INTER_CUBIC)
            combined_frame[SUB_RES[1]*y:SUB_RES[1]*y+SUB_RES[1], SUB_RES[0]*x:SUB_RES[0]*x+SUB_RES[0], :] = frame
            # draw boxes if any are detected
            if boxes is not None:
                draw_boxes(combined_frame, (boxes, scores, classes))
        # if all videos ended, terminate
        if not any(rets):
            break
        out.write(combined_frame)
    # release resources
    out.release()
    for cap in caps:
        cap.release()


if __name__ == '__main__':
    vid_paths = [
        os.path.join(CFG.BASE_PATH, 'data', 'videos', 'conv_1.mp4'),
        os.path.join(CFG.BASE_PATH, 'data', 'videos', 'conv_2.mp4'),
        os.path.join(CFG.BASE_PATH, 'data', 'videos', 'news_1.mp4'),
        os.path.join(CFG.BASE_PATH, 'data', 'videos', 'news_2.mp4'),
    ]
    main('yolov5m', vid_paths)
