import argparse
import glob
import os

import cv2 as cv
import numpy as np

here = os.path.dirname(os.path.abspath(__file__))

sz = 300
layers = [
    (30, 60, [2], 8, 38),
    (60, 111, [2, 3], 16, 19),
    (111, 162, [2, 3], 32, 10),
    (162, 213, [2, 3], 64, 5),
    (213, 264, [2], 100, 5),
    (264, 315, [2], 300, 5),
]
var = [0.1, 0.1, 0.2, 0.2]


def build_priors():
    p = []
    for mn, mx, ars, step, fm in layers:
        ratios = [1.0]
        for a in ars:
            ratios += [a, 1.0 / a]
        for y in range(fm):
            for x in range(fm):
                cx = (x + 0.5) * step
                cy = (y + 0.5) * step
                boxes = [(mn, mn), ((mn * mx) ** 0.5, (mn * mx) ** 0.5)]
                for a in ratios[1:]:
                    boxes.append((mn * a ** 0.5, mn / a ** 0.5))
                for bw, bh in boxes:
                    p.append([cx, cy, bw, bh])
    return np.array(p, np.float32)


def default_model():
    files = [f for f in glob.glob(os.path.join(here, "*.onnx")) if "known_good" not in os.path.basename(f)]
    return files[0] if files else os.path.join(here, "opencv_face_detector_uint8.onnx")


def main():
    parser = argparse.ArgumentParser(description="OpenCV SSD face detector (ONNX) demo")
    parser.add_argument("--model", default=default_model())
    parser.add_argument("--image", default=os.path.join(here, "example_outputs", "input_image.png"))
    parser.add_argument("--output", default=os.path.join(here, "example_outputs", "output_image.png"))
    parser.add_argument("--conf", type=float, default=0.4)
    args = parser.parse_args()

    img = cv.imread(args.image)
    if img is None:
        raise SystemExit("could not read image: %s" % args.image)

    inp = cv.resize(img, (sz, sz)).astype(np.float32) - np.array([104.0, 177.0, 123.0], np.float32)

    net = cv.dnn.readNetFromONNX(args.model)
    onames = net.getUnconnectedOutLayersNames()
    net.setInput(inp[None])
    res = net.forward(onames)
    loc = res[[i for i, n in enumerate(onames) if "mbox_loc" in n][0]].reshape(-1, 4)
    conf = res[[i for i, n in enumerate(onames) if "mbox_conf" in n][0]].reshape(-1, 2)

    priors = build_priors()
    pcx = priors[:, 0] / sz
    pcy = priors[:, 1] / sz
    pw = priors[:, 2] / sz
    ph = priors[:, 3] / sz

    e = np.exp(conf - conf.max(1, keepdims=True))
    sm = e / e.sum(1, keepdims=True)
    scores = sm[:, 1]

    cx = pcx + loc[:, 0] * var[0] * pw
    cy = pcy + loc[:, 1] * var[1] * ph
    bw = pw * np.exp(loc[:, 2] * var[2])
    bh = ph * np.exp(loc[:, 3] * var[3])
    boxes = np.stack([cx - bw / 2, cy - bh / 2, cx + bw / 2, cy + bh / 2], 1)

    keep = scores > args.conf
    boxes = boxes[keep]
    scores = scores[keep]
    order = scores.argsort()[::-1]
    pick = []
    while order.size:
        i = order[0]
        pick.append(i)
        xx1 = np.maximum(boxes[i, 0], boxes[order[1:], 0])
        yy1 = np.maximum(boxes[i, 1], boxes[order[1:], 1])
        xx2 = np.minimum(boxes[i, 2], boxes[order[1:], 2])
        yy2 = np.minimum(boxes[i, 3], boxes[order[1:], 3])
        inter = np.maximum(0, xx2 - xx1) * np.maximum(0, yy2 - yy1)
        ai = (boxes[i, 2] - boxes[i, 0]) * (boxes[i, 3] - boxes[i, 1])
        aj = (boxes[order[1:], 2] - boxes[order[1:], 0]) * (boxes[order[1:], 3] - boxes[order[1:], 1])
        iou = inter / (ai + aj - inter + 1e-9)
        order = order[1:][iou <= 0.3]

    h, w = img.shape[:2]
    for i in pick:
        x1, y1, x2, y2 = boxes[i]
        cv.rectangle(img, (int(x1 * w), int(y1 * h)), (int(x2 * w), int(y2 * h)), (0, 255, 0), 2)
    cv.imwrite(args.output, img)
    print("opencv_face_detector_uint8", len(pick), "faces")


if __name__ == "__main__":
    main()