import torch
from model import DeepfakeDetector
import os

DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
MODEL_PATH = os.path.join(os.path.dirname(__file__), "model.pth")

_model = None


def load_model():
    global _model

    if _model is not None:
        return _model

    model = DeepfakeDetector()

    checkpoint = torch.load(MODEL_PATH, map_location=DEVICE)
    model.load_state_dict(checkpoint['model_state_dict'])

    model.to(DEVICE)
    model.eval()

    _model = model
    return model


def run_inference(frames_tensor):
    model = load_model()

    if frames_tensor.dim() != 5:
        raise ValueError("Expected shape: (B, T, C, H, W)")

    frames_tensor = frames_tensor.to(DEVICE)

    with torch.no_grad():
        output = model(frames_tensor)
        prediction = torch.argmax(output, dim=1).item()

    return prediction