import time
import numpy as np
from torch import cuda
from sklearn.pipeline import make_pipeline
from sklearn.decomposition import TruncatedSVD
from sklearn.feature_extraction.text import TfidfVectorizer
from umap import UMAP

if cuda.is_available():
    torch_device = "gpu"
else: 
    torch_device =  "cpu"

def make_or_load_embeddings(docs, file_list, data_file_name_no_ext, embedding_model, return_intermediate_files, embeddings_super_compress, low_resource_mode_opt):

    embeddings_file_names = [string.lower() for string in file_list if "embedding" in string.lower()]  

    if embeddings_file_names:
        print("Loading embeddings from file.")
        embeddings_out = np.load(embeddings_file_names[0])['arr_0']

        # If embedding files have 'super_compress' in the title, they have been multiplied by 100 before save
        if "compress" in embeddings_file_names[0]:
            embeddings_out /= 100

        # print("embeddings loaded: ", embeddings_out)

    if not embeddings_file_names:
        tic = time.perf_counter()
        print("Starting to embed documents.")

        # Custom model
        # If on CPU, don't resort to embedding models
        if low_resource_mode_opt == "Yes":
            print("Creating simplified 'sparse' embeddings based on TfIDF")
            embedding_model = make_pipeline(
            TfidfVectorizer(),
            TruncatedSVD(100)
            )

            embeddings_out = embedding_model.encode(sentences=docs, show_progress_bar = True, batch_size = 32)

        elif low_resource_mode_opt == "No":
            print("Creating dense embeddings based on transformers model")

            embeddings_out = embedding_model.encode(sentences=docs, max_length=1024, show_progress_bar = True, batch_size = 32) # For Jina # # 

            #import torch
            #from torch.nn.utils.rnn import pad_sequence

            # Assuming embeddings_out is a list of tensors
            #embeddings_out = [torch.tensor(embedding) for embedding in embeddings_out]

            # Pad the sequences
            # Set batch_first=True if you want the batch dimension to be the first dimension
            #embeddings_out = pad_sequence(embeddings_out, batch_first=True, padding_value=0)


        toc = time.perf_counter()
        time_out = f"The embedding took {toc - tic:0.1f} seconds"
        print(time_out)

        # If you want to save your files for next time
        if return_intermediate_files == "Yes":
            if embeddings_super_compress == "No":
                semantic_search_file_name = data_file_name_no_ext + '_' + 'embeddings.npz'
                np.savez_compressed(semantic_search_file_name, embeddings_out)
            else:
                semantic_search_file_name = data_file_name_no_ext + '_' + 'embedding_compress.npz'
                embeddings_out_round = np.round(embeddings_out, 3) 
                embeddings_out_round *= 100 # Rounding not currently used
                np.savez_compressed(semantic_search_file_name, embeddings_out_round)

    # Pre-reduce embeddings for visualisation purposes
    reduced_embeddings = UMAP(n_neighbors=15, n_components=2, min_dist=0.0, metric='cosine', random_state=42).fit_transform(embeddings_out)


    return embeddings_out, reduced_embeddings