import torch
import numpy as np
from model_loader import models

class Inferencer:

    def __init__(self):
        self.clip_model = models.clip_model
        self.svm_model = models.svm_model

    @torch.no_grad()
    def predict(self, image_tensor:torch.Tensor) -> dict:
        """
        Takes a preprocessed image tensor and returns the classification result.

        Args:
            image_tensor (torch.Tensor): The preprocessed image tensor.

        Returns:
            dict: A dictionary containing the classification label and confidence score.
        """

        image_features = self.clip_model.encode_image(image_tensor)
        image_features_np = image_features.cpu().numpy()

        prediction = self.svm_model.predict(image_features_np)[0]

        if hasattr(self.svm_model, "predict_proba"):
            # If yes, use predict_proba for a true confidence score
            confidence_scores = self.svm_model.predict_proba(image_features_np)[0]
            confidence = float(np.max(confidence_scores))
        else:
            # If no, use decision_function as a fallback confidence measure.
            # The absolute value of the decision function score indicates confidence.
            # We can apply a sigmoid function to scale it to a [0, 1] range for consistency.
            decision_score = self.svm_model.decision_function(image_features_np)[0]
            confidence = 1 / (1 + np.exp(-np.abs(decision_score)))
            confidence = float(confidence)

        label_map = {0: 'real', 1: 'fake'}
        classification_label = label_map.get(prediction, "unknown")

        return {
            "classification": classification_label,
            "confidence": confidence
        }

inferencer = Inferencer()