# SPDX-License-Identifier: MIT
# Copyright (C) Intel Corporation
#
# Heatmap Generation: OpenVINO Runtime sample.
# See README.md for prerequisites and inputs. Run: python3 openvino_sample.py

import cv2
import numpy as np
import openvino as ov

CONF_THRESHOLD = 0.4
# Select the model: "yolo26n" (default) or "detr-resnet-50".
MODEL = "yolo26n"
INPUT_SIZE = 800 if MODEL.startswith("detr") else 640
HEATMAP_ALPHA = 0.55
MATCH_DIST = 80.0   # max px between frames to treat detections as the same object
MAX_SPEED = 20.0    # px/frame that maps to full red


def render_speed_heatmap(frame, speed_sum, count, alpha):
    """Color traffic by average speed: blue = slow, yellow = medium,
    red = fast. Only regions where vehicles were seen are tinted, so
    empty background keeps its original color."""
    avg = np.zeros_like(speed_sum)
    seen = count > 0
    avg[seen] = speed_sum[seen] / count[seen]
    avg = cv2.GaussianBlur(avg, (0, 0), sigmaX=15)
    presence = cv2.GaussianBlur(seen.astype(np.float32), (0, 0), sigmaX=15)
    norm = np.clip(avg / MAX_SPEED, 0, 1)  # 0 = slow (blue), 1 = fast (red)
    color = cv2.applyColorMap((norm * 255).astype(np.uint8), cv2.COLORMAP_JET)
    weight = (np.clip(presence, 0, 1) * alpha)[..., np.newaxis]
    overlay = frame.astype(np.float32) * (1 - weight) + color.astype(np.float32) * weight
    return overlay.astype(np.uint8), color


core = ov.Core()
model = core.read_model(f"{MODEL}_openvino_model/{MODEL}.xml")

# Change device to "GPU" or "NPU" to run on integrated GPU or NPU.
compiled = core.compile_model(model, "CPU")

cap = cv2.VideoCapture("test_video.mp4")
fps = cap.get(cv2.CAP_PROP_FPS) or 30.0
width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
writer = cv2.VideoWriter(
    "output_openvino.mp4", cv2.VideoWriter_fourcc(*"mp4v"), fps, (width, height))

speed_sum = np.zeros((height, width), dtype=np.float32)
count = np.zeros((height, width), dtype=np.float32)
prev_centroids = []
heatmap_color = None
frame_idx = 0
total_dets = 0

while True:
    ok, frame = cap.read()
    if not ok:
        break
    frame_idx += 1
    h0, w0 = frame.shape[:2]
    sx, sy = w0 / INPUT_SIZE, h0 / INPUT_SIZE

    blob = cv2.resize(frame, (INPUT_SIZE, INPUT_SIZE))
    blob = cv2.cvtColor(blob, cv2.COLOR_BGR2RGB).astype(np.float32) / 255.0
    blob = blob.transpose(2, 0, 1)[np.newaxis, ...]

    output = compiled([blob])[compiled.output(0)][0]
    dets = output[output[:, 4] >= CONF_THRESHOLD]
    total_dets += len(dets)

    cur_centroids = []
    for det in dets:
        x1, y1 = int(det[0] * sx), int(det[1] * sy)
        x2, y2 = int(det[2] * sx), int(det[3] * sy)
        x1, x2 = max(0, x1), min(width, x2)
        y1, y2 = max(0, y1), min(height, y2)
        cx, cy = (x1 + x2) / 2, (y1 + y2) / 2
        cur_centroids.append((cx, cy))

        # Speed = displacement from the nearest detection in the previous frame.
        speed = 0.0
        if prev_centroids:
            d = min(np.hypot(cx - px, cy - py) for px, py in prev_centroids)
            if d <= MATCH_DIST:
                speed = d
        speed_sum[y1:y2, x1:x2] += speed
        count[y1:y2, x1:x2] += 1.0
    prev_centroids = cur_centroids

    overlay, heatmap_color = render_speed_heatmap(frame, speed_sum, count, HEATMAP_ALPHA)
    cv2.putText(overlay, f"Detections: {len(dets)}", (10, 30),
                cv2.FONT_HERSHEY_SIMPLEX, 0.8, (255, 255, 255), 2)
    writer.write(overlay)

cap.release()
writer.release()

if heatmap_color is not None:
    cv2.imwrite("heatmap.jpg", heatmap_color)
    print("Saved: heatmap.jpg")

print(f"Processed {frame_idx} frames, {total_dets} total detections", flush=True)