import numpy as np
import pandas as pd
from sentence_transformers import SentenceTransformer

class ModelFactory():

    def __init__(self):
        pass

    def create_model(self, model_type):
        model = None

        if (model_type=='all-MiniLM-L6-v2'):
            model = MiniLM_L6_v2_Model()

        if (model_type=='sentence_similarity_spanish'):
            model = SentenceSimilaritySpanishModel()

        if (model_type=='multilingual-e5-large-ft-sts-spanish-matryoshka-768-64-5e'):
            model = Multilingual_E5_Large_Ft_Sts_Spanish_Matryoshka()

        return model

class BaseModel():

    def __init__(self):
        pass

    def retrieve_embeddings_from_single_input_text(self, input_text):
        embeddings = self.model.encode(input_text, batch_size=32)
        embeddings = embeddings.astype(np.float16).astype(str).tolist()

        return embeddings
    
    def retrieve_embeddings_from_texts_list(self, input_texts_list, limitnoffeatures=-1):
        all_embeddings_list = []
        for current_input_text_aux in input_texts_list:
            embeddings = self.retrieve_embeddings_from_single_input_text(current_input_text_aux)
            nof_features = len(embeddings[0])
            all_embeddings_list += [current_input_text_aux.tolist() + embeddings[0]]

        queries_embeddings_df = pd.DataFrame(all_embeddings_list)
        columns_list = ['text'] + [f'feature_{idx}' for idx in range(0, nof_features)]
        queries_embeddings_df.columns = columns_list

        if (limitnoffeatures>-1):
            columns_to_choose = queries_embeddings_df.columns[0:limitnoffeatures+1]
            queries_embeddings_df = queries_embeddings_df[columns_to_choose]

        return queries_embeddings_df

class MiniLM_L6_v2_Model(BaseModel):
    
    def __init__(self):
        self.model = SentenceTransformer('all-MiniLM-L6-v2')

class SentenceSimilaritySpanishModel(BaseModel):

    def __init__(self):
        self.model = SentenceTransformer('hiiamsid/sentence_similarity_spanish_es')

class Multilingual_E5_Large_Ft_Sts_Spanish_Matryoshka(BaseModel):

    def __init__(self):
        self.model = SentenceTransformer('mrm8488/multilingual-e5-large-ft-sts-spanish-matryoshka-768-64-5e')