import torch # type: ignore
import cv2 # type: ignore
from transformers import VideoMAEForVideoClassification, VideoMAEImageProcessor
import numpy as np # type: ignore

# Set device to GPU if available
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print("Using device:", device)

# Load model and processor
model = VideoMAEForVideoClassification.from_pretrained("OPear/videomae-large-finetuned-UCF-Crime")
processor = VideoMAEImageProcessor.from_pretrained("OPear/videomae-large-finetuned-UCF-Crime")
model = model.to(device)
model.eval()

# Load frames from video
def load_video(path, max_frames=16, sample_every_n_frames=4):
    cap = cv2.VideoCapture(path)
    frames = []
    frame_count = 0

    while len(frames) < max_frames:
        ret, frame = cap.read()
        if not ret:
            break
        if frame_count % sample_every_n_frames == 0:
            frame = cv2.resize(frame, (224, 224))
            frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
            frames.append(frame)
        frame_count += 1

    cap.release()
    return frames

# Classify video
def classify_video(path):
    video = load_video(path)
    if len(video) == 0:
        return "Error: No frames extracted from video."

    inputs = processor(video, return_tensors="pt")
    inputs = {k: v.to(device) for k, v in inputs.items()}

    with torch.no_grad():
        outputs = model(**inputs)

    predicted_class = outputs.logits.argmax().item()
    label = model.config.id2label[predicted_class]
    return label

# Classify video
def classify_video_from_path(video_path):
    result = classify_video(video_path)
    return result