import os
import shutil
import torch
import faiss

from custom_apps.faiss_imagenet.utils import preprocessing, preprocessing_quary, normalize, get_dataset_list
from custom_apps.faiss_imagenet.const import *
# from rest.app.utils.date_utils import D
# from const import OUTPUT_FOLDER

dataset_list = get_dataset_list(DATASET_TEXT)

def search_idxs(image_path,dataset_bin=DATASET_BIN,index_type="hnsw",search_num=4):
    
    if index_type not in ["hnsw", "l2"]:
        raise ValueError("index_type must be either 'hnsw' or 'l2'")
    
    DIM = 1280
    
    # res = faiss.StandardGpuResources()
    # __res = faiss.StandardGpuResources()
    
    
    dataset_fvces, dataset_index = preprocessing(DIM,dataset_bin,index_type)
    
    org_fvces, org_index = preprocessing_quary(DIM,image_path,index_type)

    # cpu -> gpu
    # gpu_idx = faiss.index_cpu_to_gpu_multiple_py([res], 0, dataset_index)
    # __gpu_idx = faiss.index_cpu_to_gpu_multiple_py([__res], 0, org_index)
    
    dists, idxs = dataset_index.search(normalize(org_fvces), search_num)
    
    index_image_save(image_path, dists[0], idxs[0])
    
    # del dataset_fvces, dataset_index, org_fvces, org_index, dists, idxs
    # del res, __res, gpu_idx, __gpu_idx
    
    # torch.cuda.empty_cache()
    # import gc
    # gc.collect()
    
def index_image_save(query_image_path, dists, idxs):
    directory_path, file = os.path.split(query_image_path)
    _name, extension = os.path.splitext(file)

    if not os.path.exists(directory_path):
        raise ValueError(f"Folder {directory_path} does not exist.")
    
    for dist,index in zip(dists, idxs):
        
        if dist > 1:
            dist = 0
        else:
            dist = 1-dist
        
        origin_file_path = dataset_list[index]
        dest_file_path = os.path.join(directory_path, f"search_{round(float(dist),4)}{extension}")
        shutil.copy(origin_file_path, dest_file_path)
        


    
    