import os
import sys

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

import time
from abc import ABC, abstractmethod

import numpy as np
import torch

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


class BaseModelImageSimilarity(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 get_similarity(self, image1: np.ndarray, image2: np.ndarray) -> dict:
        image1_input = image_augmentations()(image=image1)["image"]
        image2_input = image_augmentations()(image=image2)["image"]

        image1_input = image1_input.unsqueeze(axis=0).to(self.device)
        image2_input = image2_input.unsqueeze(axis=0).to(self.device)

        with torch.no_grad():
            start_time = time.perf_counter()
            image1_input = self.model(image1_input)
            image2_input = self.model(image2_input)
            end_time = time.perf_counter() - start_time

        image1_input = image1_input.detach().cpu().numpy()
        image2_input = image2_input.detach().cpu().numpy()
        similarity = euclidean_distance_normalized(image1_input, image2_input)
        result_similarity = (
            "same image"
            if similarity
            > configs.NAME_MODELS[self.name_model]["image_similarity_threshold"]
            else "not same image"
        )
        return {
            "similarity": similarity,
            "result_similarity": result_similarity,
            "inference_time": end_time,
        }