from typing import *
from pathlib import Path

import cv2
from inference.openvino import OpenVINOInferencer
import numpy as np

inferencer = None
last_model_path: str = None

def predict_openvino(
        image: Union[str, Path, np.ndarray],
        model_path: Union[str, Path],
        device: str,
        threshold: float) -> Dict[str, np.ndarray]:
    global inferencer, last_model_path
    if not isinstance(inferencer, OpenVINOInferencer) or last_model_path != str(model_path):
        inferencer = OpenVINOInferencer(
            path = model_path,
            device=device.upper()
        )
        last_model_path = str(model_path)
    inferencer.metadata["pixel_threshold"] = threshold
    return inferencer(image)

def draw_contour(image, mask, color=(255,0,0), thickness=3):
    contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE)
    return cv2.drawContours(image.copy(), contours, -1, color, thickness)

def convert_to_heatmap(heatmap):
    heatmap = cv2.applyColorMap((heatmap*255).astype("uint8"), cv2.COLORMAP_HSV)
    return heatmap

def predict_fn(image, threshold, device, model_dir):
    outputs = predict_openvino(image, model_dir, device, threshold)
    out_image = draw_contour(image, outputs["pred_mask"])
    heatmap = convert_to_heatmap(outputs["anomaly_map"])
    return out_image, heatmap