import faiss
import re
import pandas as pd
import torch
import numpy as np
from .. import esm
from Bio import SeqIO
from math import ceil
from ..model import MSARetrieveModel
from ..vector_construction import fasta2vec
from ..utils import TimeCounter


# 构建Faiss索引
def contruct_faiss_index(vectors, path, dim, measure, param):
	index = faiss.index_factory(dim, param, measure)
		
	if not index.is_trained:
		index.train(vectors)
	index.add(vectors)
	
	faiss.write_index(index, path)


if __name__ == '__main__':
	
	# model_name = "esm1b_t33_650M_UR50S"
	model_name = "RetrieveModel_scope2.08_d1_m1_t62_I663_B256_phase3"
	params = {"device_ids": [0, 1, 2, 3],
	          "model_path": f"../PretrainedModels/{model_name}.pt",
	          "fasta_path": f"/sujin/dataset/uniclust30/UniRef30_2020_06_100_200",
	          "save_path": f"/sujin/dataset/uniclust30/UniRef30_2020_06_100_200",
	          "batch_size": 64,
	          "log_name": None,
	          "cover": True}
	# fasta2vec(**params)
	file_names = ['400_500']
	for name in file_names:
		params['fasta_path'] = f"/sujin/dataset/uniclust30/UniRef30_2020_06_{name}"
		params['save_path'] = f"/sujin/dataset/uniclust30/UniRef30_2020_06_{name}"
		fasta2vec(**params)

	
	# name = '900_1000'
	# path = "/sujin/dataset/uniclust30/UniRef30_2020_06"
	#
	# vectors = np.load(f"{path}_{name}.npy")
	# print(vectors.shape[0])
	#
	# n = ceil(vectors.shape[0] / 256)
	# dim, measure = 1280, faiss.METRIC_L2
	# param = f'IVF{n}, Flat'
	# save_path = f"{path}_{name}_1.index"
	# with TimeCounter("Constructing index..."):
	# 	contruct_faiss_index(vectors, save_path, dim, measure, param)