from transformers import AutoTokenizer, AutoModel
import torch
import torch.nn.functional as F
import hnswlib
import numpy as np
import datetime
from fastapi import FastAPI
from pydantic import BaseModel
from typing import List

if torch.cuda.is_available():
    print("CUDA is available! Inference on GPU!")
else:
    print("CUDA is not available. Inference on CPU.")

seperator = "-HFSEP-"
base_name="intfloat/e5-small-v2"
device="cuda"
max_length=512
max_batch_size = 500
tokenizer = AutoTokenizer.from_pretrained(base_name)
model = AutoModel.from_pretrained(base_name).to(device)

def current_timestamp():
    return datetime.datetime.utcnow().timestamp()

def get_embeddings(input_texts):
    input_texts = input_texts[:max_batch_size]
    batch_dict = tokenizer(
        input_texts,
        max_length=max_length,
        padding=True,
        truncation=True,
        return_tensors='pt'
    ).to(device)

    with torch.no_grad():
        outputs = model(**batch_dict)
        embeddings = _average_pool(
            outputs.last_hidden_state, batch_dict['attention_mask']
        )
        embeddings = F.normalize(embeddings, p=2, dim=1)
        embeddings_np = embeddings.cpu().numpy()
        
        if device == "cuda":
            del embeddings
            torch.cuda.empty_cache()        
        
        return embeddings_np

def _average_pool(
    last_hidden_states,
    attention_mask
):
    last_hidden = last_hidden_states.masked_fill(~attention_mask[..., None].bool(), 0.0)
    return last_hidden.sum(dim=1) / attention_mask.sum(dim=1)[..., None]

def create_hnsw_index(embeddings_np, space='ip', ef_construction=100, M=16):
    index = hnswlib.Index(space=space, dim=len(embeddings_np[0]))
    index.init_index(max_elements=len(embeddings_np), ef_construction=ef_construction, M=M)
    ids = np.arange(embeddings_np.shape[0])
    index.add_items(embeddings_np, ids)
    return index

def preprocess_texts(query, paragraphs):
    query = f'query: {query}'
    paragraphs = [f'passage: {p}' for p in paragraphs]
    return [query]+paragraphs


app = FastAPI()

class EmbeddingsSimilarityReq(BaseModel):
    paragraphs: List[str]
    query: str
    top_k: int

@app.post("/")
async def find_similar_paragraphsitem(req: EmbeddingsSimilarityReq):
    print("Len of batches", len(req.paragraphs))

    print("creating embeddings", current_timestamp())
    inputs = preprocess_texts(req.query, req.paragraphs)
    embeddings_np = get_embeddings(inputs)
    query_embedding, chunks_embeddings = embeddings_np[0], embeddings_np[1:]
    
    print("creating index", current_timestamp())
    search_index = create_hnsw_index(chunks_embeddings)
    print("searching index", current_timestamp())
    labels, _ = search_index.knn_query(query_embedding, k=min(int(req.top_k), len(chunks_embeddings)))
    labels = labels[0].tolist()
    return labels