import math
import sys
from argparse import ArgumentParser
from pathlib import Path

import cv2
import onnxruntime
from config import (CLASS_COLORS, CLASS_NAMES, ModelType, YOLOv5_ANCHORS,
                    YOLOv7_ANCHORS)
from cv2_nms import non_max_suppression
from numpy_coder import Decoder
from preprocess import Preprocess
from tqdm import tqdm

# Add __FILE__  to sys.path
sys.path.append(str(Path(__file__).resolve().parents[0]))

IMG_EXTENSIONS = ('.jpg', '.jpeg', '.png', '.ppm', '.bmp', '.pgm', '.tif',
                  '.tiff', '.webp')


def path_to_list(path: str):
    path = Path(path)
    if path.is_file() and path.suffix in IMG_EXTENSIONS:
        res_list = [str(path.absolute())]
    elif path.is_dir():
        res_list = [
            str(p.absolute()) for p in path.iterdir()
            if p.suffix in IMG_EXTENSIONS
        ]
    else:
        raise RuntimeError
    return res_list


def parse_args():
    parser = ArgumentParser()
    parser.add_argument(
        'img', help='Image path, include image file, dir and URL.')
    parser.add_argument('onnx', type=str, help='Onnx file')
    parser.add_argument('--type', type=str, help='Model type')
    parser.add_argument(
        '--img-size',
        nargs='+',
        type=int,
        default=[640, 640],
        help='Image size of height and width')
    parser.add_argument(
        '--out-dir', default='./output', type=str, help='Path to output file')
    parser.add_argument(
        '--show', action='store_true', help='Show the detection results')
    parser.add_argument(
        '--score-thr', type=float, default=0.3, help='Bbox score threshold')
    parser.add_argument(
        '--iou-thr', type=float, default=0.7, help='Bbox iou threshold')
    args = parser.parse_args()
    return args


def main():
    args = parse_args()
    out_dir = Path(args.out_dir)
    model_type = ModelType(args.type.lower())

    if not args.show:
        out_dir.mkdir(parents=True, exist_ok=True)

    files = path_to_list(args.img)
    session = onnxruntime.InferenceSession(
        args.onnx, providers=['CPUExecutionProvider'])
    preprocessor = Preprocess(model_type)
    decoder = Decoder(model_type, model_only=True)
    if model_type == ModelType.YOLOV5:
        anchors = YOLOv5_ANCHORS
    elif model_type == ModelType.YOLOV7:
        anchors = YOLOv7_ANCHORS
    else:
        anchors = None

    for file in tqdm(files):
        image = cv2.imread(file)
        image_h, image_w = image.shape[:2]
        img, (ratio_w, ratio_h) = preprocessor(image, args.img_size)
        features = session.run(None, {'images': img})
        decoder_outputs = decoder(
            features,
            args.score_thr,
            num_labels=len(CLASS_NAMES),
            anchors=anchors)
        nmsd_boxes, nmsd_scores, nmsd_labels = non_max_suppression(
            *decoder_outputs, args.score_thr, args.iou_thr)
        for box, score, label in zip(nmsd_boxes, nmsd_scores, nmsd_labels):
            x0, y0, x1, y1 = box
            x0 = math.floor(min(max(x0 / ratio_w, 1), image_w - 1))
            y0 = math.floor(min(max(y0 / ratio_h, 1), image_h - 1))
            x1 = math.ceil(min(max(x1 / ratio_w, 1), image_w - 1))
            y1 = math.ceil(min(max(y1 / ratio_h, 1), image_h - 1))
            cv2.rectangle(image, (x0, y0), (x1, y1), CLASS_COLORS[label], 2)
            cv2.putText(image, f'{CLASS_NAMES[label]}: {score:.2f}',
                        (x0, y0 - 5), cv2.FONT_HERSHEY_SIMPLEX, 0.5,
                        (0, 255, 255), 2)
        if args.show:
            cv2.imshow('result', image)
            cv2.waitKey(0)
        else:
            cv2.imwrite(f'{out_dir / Path(file).name}', image)


if __name__ == '__main__':
    main()