# -*- coding: utf-8 -*-
"""
@file : faiss_similarity_search.py
@author: hsj100
@license: A2D2
@brief: 특징 벡터를 이용한 이미지 유사도 검색

@section Modify History
- 2025-06-19 오후 2:45 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: standard
"""
import sys
import torch
import faiss
import numpy as np

from transformers import CLIPProcessor, CLIPModel
from PIL import Image
from abc import ABC, abstractmethod
from enum import Enum
from datetime import datetime


"""
Package: 3rd party
"""

"""
Package: custom
"""
# NOTE(hsj100): 코드 단독 실행을 위한 기본 경로 지정 (일반적으로는 본 코드 보다는 모듈 실행 방법을 권장)
import os, sys
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.fem_openaiclipvit import FEOpenAIClipViT
from custom_apps.FEATURE_VECTOR_SIMILARITY_FAISS.const import *


"""
Definition
"""
FILE_BASE_PATH = ''
if __name__ == '__main__':
    FILE_BASE_PATH = os.path.dirname(os.path.realpath(__file__)) + '/'

# TRAINED_MODEL_PATH = FILE_BASE_PATH + 'openai/clip-resnet50' # OLD VERSION
# TRAINED_MODEL_PATH = FILE_BASE_PATH + 'openai/clip-vit-base-patch32'  # huggingface 저장소 위치
TRAINED_MODEL_PATH = "./data" # 로컬 다운로드 경로

DEFAULT_INDEX_IMAGE_PATH = FILE_BASE_PATH + 'index_images'

class FEMUsageInfo(Enum):
    """
    특징 추출(이미지 임베딩) 모델 추가시, 생성 정보를 추가해서 사용
    * 변수명 변경되면 인덱스파일 이름 새로 만듬
    """
    openaiclip_vit_b32_l2 = [FEOpenAIClipViT, FEM.FEMArguments(
        trained_model=os.path.join(PRETRAINED_MODEL_PATH,ViTB32), 
        token=HUGGINGFACE_TOKEN,
        index_type=INDEX_TYPE_L2)]
    
    openaiclip_vit_b32_cos = [FEOpenAIClipViT, FEM.FEMArguments(
        trained_model=os.path.join(PRETRAINED_MODEL_PATH,ViTB32), 
        token=HUGGINGFACE_TOKEN,
        index_type=INDEX_TYPE_COSINE)]
    
    openaiclip_vit_b16_l2 = [FEOpenAIClipViT, FEM.FEMArguments(
        trained_model=os.path.join(PRETRAINED_MODEL_PATH,ViTB16), 
        token=HUGGINGFACE_TOKEN,
        index_type=INDEX_TYPE_L2)]
    
    openaiclip_vit_b16_cos = [FEOpenAIClipViT, FEM.FEMArguments(
        trained_model=os.path.join(PRETRAINED_MODEL_PATH,ViTB16), 
        token=HUGGINGFACE_TOKEN,
        index_type=INDEX_TYPE_COSINE)]
    
    openaiclip_vit_l14_l2 = [FEOpenAIClipViT, FEM.FEMArguments(
        trained_model=os.path.join(PRETRAINED_MODEL_PATH,ViTL14), 
        token=HUGGINGFACE_TOKEN,
        index_type=INDEX_TYPE_L2)]

    openaiclip_vit_l14_cos = [FEOpenAIClipViT, FEM.FEMArguments( 
        trained_model=os.path.join(PRETRAINED_MODEL_PATH,ViTL14),
        token=HUGGINGFACE_TOKEN, 
        index_type=INDEX_TYPE_COSINE)]
    
    openaiclip_vit_l14_336_l2 = [FEOpenAIClipViT, FEM.FEMArguments(
        trained_model=os.path.join(PRETRAINED_MODEL_PATH,ViTL14_336), 
        token=HUGGINGFACE_TOKEN,
        index_type=INDEX_TYPE_L2)]
    
    openaiclip_vit_l14_336_cos = [FEOpenAIClipViT, FEM.FEMArguments(
        trained_model=os.path.join(PRETRAINED_MODEL_PATH,ViTL14_336), 
        token=HUGGINGFACE_TOKEN,
        index_type=INDEX_TYPE_COSINE)]
    
    openaicliprn = None
    mobilenetv2 = None


DEFAULT_WEIGHT_PATH = FILE_BASE_PATH + TRAINED_MODEL_PATH

TEST_IMAGES_PATH = FILE_BASE_PATH + 'testimages'

"""
Implementation
"""
# log off
import logging
logging.getLogger('PIL.PngImagePlugin').disabled = True

# Utility Method
def function_name():
    return sys._getframe(1).f_code.co_name


class FeatureExtractionModel(ABC):
    """
    특징 벡터 추출 모델의 추상 인터페이스
    """

    @abstractmethod
    def init_model(self, pretrained_model):
        """
        로드 모델 및 초기화
        :param pretrained_model:
        :return:
        """
        raise NotImplementedError(f'Subclasses must implement this method{function_name()}')

    @abstractmethod
    def image_embedding(self, image_np):
        """
        이미지 임베딩 - 이미지 데이터(numpy)
        :param image_np: 이미지 데이터(numpy array)
        :return:
        """
        raise NotImplementedError(f'Subclasses must implement this method{function_name()}')


class VectorSimilarity:
    """
    벡터 기반 유사성 검색 모델 (FACEBOOK - FAISS)
    """
    def __init__(self, model_info:FEMUsageInfo=None, index_file_path=None, txt_file_path=None):
        self.fem_model:FEM.FeatureExtractionModel = None

        self.model_info = model_info
        self.fem_model_name = ''
        self.fem_model_args = None
        self.index_file_path = index_file_path
        self.txt_file_path = txt_file_path
        self.index_type = model_info.value[1].index_type

        if model_info is not None:
            self.create_feature_extraction_model(model_info)

    def create_feature_extraction_model(self, model_info):
        """
        특징 벡터 추출 모델 초기화
        :return: True/False
        """
        # check FEMUsageInfo
        if isinstance(model_info, FEMUsageInfo):
            if model_info not in FEMUsageInfo or model_info.value is None:
                log.error('invalid FEMArguments')
                return False
        else:
            log.error('invalid FEMArguments')
            return False

        if model_info.value[0] is None or model_info.value[1] is None:
            log.error('invalid FEMArguments')
            return False

        # release model
        if self.fem_model is not None:
            del self.fem_model

        # initialize ai-model
        self.fem_model_name = model_info.name
        self.fem_model_args = model_info.value[1]
        self.fem_model = model_info.value[0](self.fem_model_args)

        log.info(f'created fem-model: {model_info.value[0]}')

    def image_embedding_from_file(self, image_file):
        """
        이미지 파일에서 특징 벡터 추출
        :param image_file: 이미지 파일
        :return: 특징 벡터 or None
        """
        feature_vectors = None

        if not os.path.isfile(image_file):
            log.error(f'not found file[{image_file}]')
            return feature_vectors

        image_data_np = Image.open(image_file).convert('RGB')
        feature_vectors = self.fem_model.image_embedding(image_data_np)
        return feature_vectors

    def image_embedding_from_data(self, image_data_np=None):
        """
        이미지 데이터(numpy)에서 특징 벡터 추출
        :param image_data_np: 이미지 데이터(numpy)
        :return: 특징 벡터 or None
        """
        feature_vectors = None

        if image_data_np is None:
            log.error(f'invalid data[{image_data_np}]')
            return feature_vectors

        feature_vectors = self.fem_model.image_embedding(image_data_np)
        return feature_vectors

    def save_index_from_files(self, images_path=None, save_index_dir=None, save_txt_dir=None, index_type=INDEX_TYPE_L2):
        """
        인덱스 정보 저장
        :param images_path: 이미지 파일 경로
        :param save_index_dir: 인덱스 저장 디렉토리
        :param save_txt_dir: 이미지 경로 저장 디렉토리
        :param index_type: 인덱스 타입 (L2, Cosine)
        :return: 이미지 임베딩 차원 정보
        """
        assert self.fem_model is not None, 'invalid fem_model'
        assert images_path is not None, 'images_path is required'

        if not os.path.exists(images_path):
            log.error(f'image_path={images_path} does not exist')

        image_folder = images_path
        result = None

        # 모든 이미지 임베딩 모으기
        if os.path.exists(image_folder):
            image_paths = [os.path.join(image_folder, image_file) for image_file in os.listdir(image_folder)]
            image_vectors = np.vstack([self.image_embedding_from_file(image_file) for image_file in image_paths])
            log.debug(f'image_vectors.shape={image_vectors.shape}')
            result = image_vectors.shape

            # 인덱스 생성 (L2)
            dimension = image_vectors.shape[1]

            if index_type == INDEX_TYPE_L2:
                index = faiss.IndexFlatL2(dimension)  # L2
            elif index_type == INDEX_TYPE_COSINE:
                index = faiss.IndexFlatIP(dimension) # cosine
                faiss.normalize_L2(image_vectors)
            else:
                raise Exception(f'Invalid index_type={index_type}')

            index.add(image_vectors)

            # 인덱스 저장
            if save_index_dir is None:
                save_index_path = os.path.join(FILE_BASE_PATH,f'{self.fem_model_name}_{os.path.basename(self.fem_model_args.trained_model)}_{index_type}_{DEFAULT_SAVE_INDEX_SUFFIX}')
                self.index_file_path = save_index_path
            else:
                save_index_path = os.path.join(save_index_dir,f'{self.fem_model_name}_{os.path.basename(self.fem_model_args.trained_model)}_{index_type}_{DEFAULT_SAVE_INDEX_SUFFIX}')
                os.makedirs(save_index_dir, exist_ok=True)
            
            # 이미지 경로 파일 저장    
            if save_txt_dir is None:
                save_txt_path = os.path.join(FILE_BASE_PATH, f'{self.fem_model_name}_{os.path.basename(self.fem_model_args.trained_model)}_{index_type}.txt')
                self.txt_file_path = save_txt_path
            else:
                save_txt_path = os.path.join(save_txt_dir, f'{self.fem_model_name}_{os.path.basename(self.fem_model_args.trained_model)}_{index_type}.txt')
                os.makedirs(save_txt_dir, exist_ok=True)

            if not os.path.exists(save_txt_path):
                faiss.write_index(index, save_index_path)
                log.debug(f'save_index_path={save_index_path}')
            
            if not os.path.exists(save_txt_path):
                # 이미지 경로 저장
                with open(save_txt_path, 'w') as f:
                    for path in image_paths:
                        f.write(f"{path}\n")
                log.debug(f'save_txt_path={save_txt_path}')
            
        else:
            log.error(f'Image folder {image_folder} does not exist')
        return result

    def set_index_path(self, index_path):
        assert index_path is not None, 'index_path is required'
        self.index_file_path = index_path

    def load_index(self, index_file_name=None):
        """
        인덱스 정보 불러오기
        :return:
        """
        if index_file_name is not None:
            self.index_file_path = index_file_name

        return faiss.read_index(self.index_file_path)
    
    def time_diff(self, start_time, end_time):
        # 시간 차이 계산
        time_difference = end_time - start_time

        # timedelta 객체에서 분:초:마이크로초 추출
        total_seconds_float = time_difference.total_seconds()
        total_seconds_int = int(total_seconds_float) # 정수 부분의 초

        minutes = total_seconds_int // 60
        seconds = total_seconds_int % 60
        microseconds = time_difference.microseconds # 남은 마이크로초

        return f"{minutes:02d}:{seconds:02d}.{microseconds:06d}"

    def query_faiss(self, query_image=None, top_k=4):

        assert query_image is not None, 'query_image is required'
        
        if os.path.exists(self.txt_file_path):
            with open(self.txt_file_path, 'r') as f:
                image_paths = [line.strip() for line in f.readlines()]
        else:
            logging.error("Image path list TXT file not found.")
            image_paths = []

        start_vector_time = datetime.now()
        index = self.load_index(self.index_file_path)
        query_vector = self.image_embedding_from_file(query_image)
        end_vector_time = datetime.now()
        
        diff_vector_time = self.time_diff(start_vector_time,end_vector_time)
        
        if self.index_type == INDEX_TYPE_COSINE:
            faiss.normalize_L2(query_vector)
        
        start_search_time = datetime.now()
        distances, indices = index.search(query_vector, top_k)
        end_search_time = datetime.now()
        
        diff_search_time = self.time_diff(start_search_time,end_search_time)
        diff_total_time = self.time_diff(start_vector_time,end_search_time)
        inference_times = f"Total time - {diff_total_time}, vector_time - {diff_vector_time}, search_time - {diff_search_time}"
        
        result_img_paths = []
        result_percents = []
        
        # 결과
        # for i in range(top_k):
        #     print(f"{i + 1}: {image_paths[indices[0][i]]}, Distance: {distances[0][i]}")
        for idx, dist in zip(indices[0], distances[0]):
            log.debug(f"{idx} (거리: {dist:.4f})")
            
            result_img_paths.append(image_paths[idx])

            if self.index_type == INDEX_TYPE_COSINE:
                result_percents.append(f"{dist*100:.2f}")
            else:
                result_percents.append(f"{((1 - dist)*100):.2f}")
                
        return inference_times, result_img_paths, result_percents

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

    # index_file_path = f'{FILE_BASE_PATH}openaiclipvit_clip-vit-base-patch32_index.faiss'
    model = FEMUsageInfo.openaiclip_vit_l14_cos
    index_file_path = os.path.join(VECTOR_IMG_LIST_PATH,f'{model.name}_{os.path.basename(model.value[1].trained_model)}_{model.value[1].index_type}_{DEFAULT_INDEX_NAME_SUFFIX}')

    cm = VectorSimilarity(model_info=model, 
                          index_file_path=index_file_path,
                          txt_file_path=os.path.join(VECTOR_IMG_LIST_PATH,f'{model.name}_{os.path.basename(model.value[1].trained_model)}_{model.value[1].index_type}.txt'))
    
    # cm.create_feature_extraction_model(FEMUsageInfo.openaiclipvit)

    cm.save_index_from_files(images_path=INDEX_IMAGES_PATH,
                            save_index_dir=FAISS_VECTOR_PATH,
                            save_txt_dir=VECTOR_IMG_LIST_PATH,
                            index_type=model.value[1].index_type)
                             

    cm.query_faiss(os.path.join(INDEX_IMAGES_PATH,"img1.png"))


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