# -*- coding: utf-8 -*-
"""
@file : fem_openaiclipvit.py
@author: hsj100
@license: A2D2
@brief: OpenAI Clip Model (Image & Text Embedding)

@section Modify History
- 2025-06-27 오후 2:05 hsj100 base

"""

"""
Basic logger: custom color logger
"""
try:
    from custom_logger.main_log import vactor_logger as log
except ImportError:
    import logging as log; log.basicConfig(level=log.DEBUG)


"""
Package: 3rd party
"""

"""
Package: standard
"""
import torch
import os, sys

from transformers import CLIPProcessor, CLIPModel
from huggingface_hub import login as huggingface_login

"""
Package: custom
"""
# NOTE(hsj100): 코드 단독 실행을 위한 기본 경로 지정 (일반적으로는 본 코드 보다는 모듈 실행 방법을 권장)

if __name__ == '__main__':
    # 현재 스크립트의 디렉토리를 기준으로 부모 디렉토리의 경로를 가져옴
    parent_dir = os.path.abspath(os.path.join(os.path.dirname(__file__), '..'))
    if parent_dir not in sys.path:
        sys.path.append(parent_dir)

from custom_apps.FEATURE_VECTOR_SIMILARITY_FAISS import feature_extraction_model as FEM
from custom_apps.FEATURE_VECTOR_SIMILARITY_FAISS.const import *


"""
Definition
"""


"""
Implementation
"""
# 특징 벡터 추출 모델 개별 정의
class FEOpenAIClipViT(FEM.FeatureExtractionModel):
    def __init__(self, arguments=FEM.FEMArguments()):
        """
        Constructor
        :param arguments: parameter of this model
        """
        # model interface
        self.model = None
        self.processor = None

        # model parameters
        self.device = None
        self.trained_model = None  # 학습된 모델 결과물이 있는 경로
        self.huggingface_token = None  # huggingface token


        self.fem_arguments = None

        self.load_model(fem_arguments=arguments)

    def __del__(self):
        """
        Destructor
        :return:
        """
        self.release_model()

    def load_model(self, fem_arguments=None):
        """
        추상 인터페이스 구현: 모델 초기화 및 메모리 적재
        :param fem_arguments: 모델 인자 정보
        :return:
        """
        assert fem_arguments is not None, 'failed init_model: fem_arguments is None'

        # model parameters
        self.device = fem_arguments.device
        self.trained_model = fem_arguments.trained_model
        self.huggingface_token = fem_arguments.token

        # device (CPU/GPU)
        if self.device == FEM.MODEL_DEVICE_GPU:
            if not torch.cuda.is_available():
                self.device = FEM.MODEL_DEVICE_CPU
                log.warning('Not found GPU device, use CPU instead')

        # huggingface token
        if self.huggingface_token:
            huggingface_login(fem_arguments.token)

        # model path check
        if not os.path.exists(fem_arguments.trained_model):
            self._download_model(fem_arguments.trained_model)
        
        # model interface
        self.model = CLIPModel.from_pretrained(fem_arguments.trained_model)
        self.processor = CLIPProcessor.from_pretrained(fem_arguments.trained_model)

        if self.device == FEM.MODEL_DEVICE_GPU:
            self.model.to(self.device)

        # inference mode
        self.model.eval()

        log.info(f'load complete:{self.__class__.__name__}')
        
    def _download_model(self, pretrained_model_path):
        """
        모델 다운로드 (로컬에 없는 경우)
        :param pretrained_model: 모델 이름
        :return:
        """
        MODEL_LIST = [ViTB32, ViTB16, ViTL14, ViTL14_336]
        
        model_name = os.path.basename(pretrained_model_path)
        
        if model_name not in MODEL_LIST:
            raise ValueError(f"Unsupported model name: {model_name}. Supported models are: {MODEL_LIST}")
        
        model = CLIPModel.from_pretrained(f"openai/{model_name}")
        processor = CLIPProcessor.from_pretrained(f"openai/{model_name}")

        os.makedirs(pretrained_model_path, exist_ok=True)

        model.save_pretrained(pretrained_model_path)
        processor.save_pretrained(pretrained_model_path)
        
        log.info(f'Model downloaded and saved to {pretrained_model_path}')

    def release_model(self):
        """
        추상 인터페이스 구현: 모델 메모리 해제
        모델 해제
        :return:
        """
        # GPU 메모리 캐시 비우기 (GPU 메모리 해제에 도움)
        if torch.cuda.is_available():
            torch.cuda.empty_cache()

        log.info(f'release complete:{self.__class__.__name__}')

    def image_embedding(self, image_np):
        """
        추상 인터페이스 구현 : 이미지 임베딩 (특징 벡터 추출)
        :param image_np: 이미지 데이터 (numpy)
        :return: 특징 벡터 (numpy)
        """
        result = None
        # 이미지 로딩 및 전처리
        # image_np = Image.open(CLIPModel.TEST_IMAGE_PATH if image_path is None else image_path).convert('RGB')
        inputs_tensor = self.processor(images=image_np, return_tensors='pt')

        if self.device == FEM.MODEL_DEVICE_GPU:
            inputs_tensor = inputs_tensor.to(self.device)

        # 특징 벡터 추출
        with torch.no_grad():
            outputs = self.model.get_image_features(**inputs_tensor)  # shape: (1, 512)
            features = outputs / outputs.norm(p=2, dim=-1, keepdim=True)  # 임베딩 정규화 L2

        # log.debug(f'feature vector shape: {features.shape}')

        if self.device == FEM.MODEL_DEVICE_GPU:
            result = features.cpu().numpy()
        else:
            result = features.numpy()

        # log.debug(f'embedding complete:{self.__class__.__name__}')
        return result

def test():
    """
    module test function
    """
    log.info('\nModule: fem_openaiclipvit.py')


if __name__ == '__main__':
    """
    test module
    """
    test()
