import os
import cv2
import numpy as np
from pathlib import Path

def load_templates(cursor_dir):
    templates = []
    for fname in os.listdir(cursor_dir):
        if fname.endswith('.png'):
            templates.append(cv2.imread(os.path.join(cursor_dir, fname), cv2.IMREAD_UNCHANGED))
    return templates

def track_cursor(frames_dir, cursor_dir, output_json):
    import json
    templates = load_templates(cursor_dir)
    results = []
    for frame_file in sorted(os.listdir(frames_dir)):
        if not frame_file.endswith('.png'):
            continue
        frame_path = os.path.join(frames_dir, frame_file)
        frame = cv2.imread(frame_path)
        best_match = None
        best_val = -np.inf
        for template in templates:
            res = cv2.matchTemplate(frame, template[..., :3], cv2.TM_CCOEFF_NORMED)
            min_val, max_val, min_loc, max_loc = cv2.minMaxLoc(res)
            if max_val > best_val:
                best_val = max_val
                best_match = (max_loc, template.shape[:2], max_val)
        if best_match:
            (x, y), (h, w), score = best_match
            results.append({"frame": frame_file, "cursor_pos": [int(x), int(y)], "score": float(score)})
    with open(output_json, "w") as f:
        json.dump(results, f, indent=2)
    print(f"Cursor tracking complete. Results saved to {output_json}")

if __name__ == "__main__":
    track_cursor("frames", "cursors", "annotations/cursor_positions.json")