import os
import sys

sys.path.append(os.path.join(os.path.dirname(__file__), "..", ".."))

from abc import ABC, abstractmethod

import numpy as np
from PIL import Image
from pytorch_grad_cam import GradCAM
from pytorch_grad_cam.utils.image import show_cam_on_image
from pytorch_grad_cam.utils.model_targets import ClassifierOutputTarget

from utils import configs
from utils.functional import (
    check_data_type_variable,
    get_device,
    image_augmentations,
    normalize_image_to_zero_one,
    reshape_transform,
)


class BaseModelGradCAM(ABC):
    def __init__(
        self,
        name_model: str,
        freeze_model: bool,
        pretrained_model: bool,
        support_set_method: str,
    ):
        self.name_model = name_model
        self.freeze_model = freeze_model
        self.pretrained_model = pretrained_model
        self.support_set_method = support_set_method
        self.model = None
        self.device = get_device()

        self.check_arguments()

    def check_arguments(self):
        check_data_type_variable(self.name_model, str)
        check_data_type_variable(self.freeze_model, bool)
        check_data_type_variable(self.pretrained_model, bool)
        check_data_type_variable(self.support_set_method, str)

        old_name_model = self.name_model
        if self.name_model == configs.CLIP_NAME_MODEL:
            old_name_model = self.name_model
            self.name_model = "clip"
        if self.name_model not in tuple(configs.NAME_MODELS.keys()):
            raise ValueError(f"Model {self.name_model} not supported")
        if self.support_set_method not in configs.SUPPORT_SET_METHODS:
            raise ValueError(
                f"Support set method {self.support_set_method} not supported"
            )
        self.name_model = old_name_model

    @abstractmethod
    def init_model(self):
        pass

    def set_grad_cam(self):
        if self.name_model == "resnet50":
            self.target_layers = (self.model.model.layer4[-1],)
        elif self.name_model == "vgg16":
            self.target_layers = (self.model.model.features[-1],)
        elif self.name_model == "inception_v4":
            self.target_layers = (self.model.model.features[-1],)
        elif self.name_model == "efficientnet_b4":
            self.target_layers = (self.model.model.blocks[-1],)
        elif self.name_model == "mobilenetv3_large_100":
            self.target_layers = (self.model.model.blocks[-1],)
        elif self.name_model == "densenet121":
            self.target_layers = (self.model.model.features[-1],)
        elif self.name_model == "vit_base_patch16_224_dino":
            self.target_layers = (self.model.model.blocks[-1].norm1,)
        elif self.name_model == "clip":
            self.target_layers = (
                self.model.vision_model.encoder.layers[-1].layer_norm1,
            )
        else:
            self.target_layers = (self.model.model.features[-1],)

        if self.name_model in ("vit_base_patch16_224_dino", "clip"):
            self.gradcam = GradCAM(
                model=self.model,
                target_layers=self.target_layers,
                reshape_transform=reshape_transform,
                use_cuda=True if self.device.type == "cuda" else False,
            )
        else:
            self.gradcam = GradCAM(
                model=self.model,
                target_layers=self.target_layers,
                use_cuda=True if self.device.type == "cuda" else False,
            )

    def get_grad_cam(self, image: np.ndarray) -> np.ndarray:
        image = np.array(
            Image.fromarray(image).resize((configs.SIZE_IMAGES, configs.SIZE_IMAGES))
        )
        image_input = image_augmentations()(image=image)["image"]
        image_input = image_input.unsqueeze(axis=0).to(self.device)
        gradcam = self.gradcam(image_input)
        gradcam = gradcam[0, :]
        gradcam = show_cam_on_image(
            normalize_image_to_zero_one(image), gradcam, use_rgb=True
        )
        return gradcam

    def get_grad_cam_with_output_target(
        self, image: np.ndarray, index_class: int
    ) -> np.ndarray:
        image = np.array(
            Image.fromarray(image).resize((configs.SIZE_IMAGES, configs.SIZE_IMAGES))
        )
        image_input = image_augmentations()(image=image)["image"]
        image_input = image_input.unsqueeze(axis=0).to(self.device)
        targets = (ClassifierOutputTarget(index_class),)
        gradcam = self.gradcam(image_input, targets=targets)
        gradcam = gradcam[0, :]
        gradcam = show_cam_on_image(
            normalize_image_to_zero_one(image), gradcam, use_rgb=True
        )
        return gradcam