import cv2
import torch

from transformers import AutoImageProcessor, Swinv2ForImageClassification

from cam import ClassActivationMap


class GlaucomaModel(object):
    def __init__(self, 
                 cls_model_path="pamixsun/swinv2_tiny_for_glaucoma_classification", 
                 device=torch.device('cpu')):
        # where to load the model, gpu or cpu ?
        self.device = device
        # classification model for nails disease
        self.cls_extractor = AutoImageProcessor.from_pretrained(cls_model_path)
        self.cls_model = Swinv2ForImageClassification.from_pretrained(cls_model_path).to(device).eval()
        # class activation map
        self.cam = ClassActivationMap(self.cls_model, self.cls_extractor)

        # classification id to label
        self.id2label = self.cls_model.config.id2label

        # number of classes for nails disease
        self.num_diseases = len(self.id2label)

    def glaucoma_pred(self, image):
        """
        Args:
            image: image array in RGB order.
        """
        inputs = self.cls_extractor(images=image.copy(), return_tensors="pt")
        with torch.no_grad():
            inputs.to(self.device)
            outputs = self.cls_model(**inputs).logits
            disease_idx = outputs.cpu()[0, :].detach().numpy().argmax()

        return disease_idx
    
    def process(self, image):
        """
        Args:
            image: image array in RGB order.
        """
        image_shape = image.shape[:2]
        disease_idx = self.glaucoma_pred(image)
        cam = self.cam.get_cam(image, disease_idx)
        cam = cv2.resize(cam, image_shape[::-1])

        return disease_idx, cam