from typing import List, Tuple

import numpy as np
import pandas as pd
import torch
from tqdm import tqdm
from transformers import AutoModel, AutoTokenizer

from autorag.nodes.passagereranker.base import passage_reranker_node
from autorag.utils.util import flatten_apply, sort_by_scores, select_top_k


@passage_reranker_node
def colbert_reranker(
	queries: List[str],
	contents_list: List[List[str]],
	scores_list: List[List[float]],
	ids_list: List[List[str]],
	top_k: int,
	batch: int = 64,
	model_name: str = "colbert-ir/colbertv2.0",
) -> Tuple[List[List[str]], List[List[str]], List[List[float]]]:
	"""
	Rerank a list of contents with Colbert rerank models.
	You can get more information about a Colbert model at https://huggingface.co/colbert-ir/colbertv2.0.
	It uses BERT-based model, so recommend using CUDA gpu for faster reranking.

	:param queries: The list of queries to use for reranking
	:param contents_list: The list of lists of contents to rerank
	:param scores_list: The list of lists of scores retrieved from the initial ranking
	:param ids_list: The list of lists of ids retrieved from the initial ranking
	:param top_k: The number of passages to be retrieved
	:param batch: The number of queries to be processed in a batch
	    Default is 64.
	:param model_name: The model name for Colbert rerank.
	    You can choose colbert model for reranking.
	    Default is "colbert-ir/colbertv2.0".
	:return: Tuple of lists containing the reranked contents, ids, and scores
	"""
	device = "cuda" if torch.cuda.is_available() else "cpu"
	model = AutoModel.from_pretrained(model_name).to(device)
	tokenizer = AutoTokenizer.from_pretrained(model_name)

	# get query and content embeddings
	query_embedding_list = get_colbert_embedding_batch(queries, model, tokenizer, batch)
	content_embedding_list = flatten_apply(
		get_colbert_embedding_batch,
		contents_list,
		model=model,
		tokenizer=tokenizer,
		batch_size=batch,
	)

	del model
	if torch.cuda.is_available():
		torch.cuda.empty_cache()

	df = pd.DataFrame(
		{
			"ids": ids_list,
			"query_embedding": query_embedding_list,
			"contents": contents_list,
			"content_embedding": content_embedding_list,
		}
	)
	temp_df = df.explode("content_embedding")
	temp_df["score"] = temp_df.apply(
		lambda x: get_colbert_score(x["query_embedding"], x["content_embedding"]),
		axis=1,
	)
	df["scores"] = temp_df.groupby(level=0, sort=False)["score"].apply(list).tolist()
	df[["contents", "ids", "scores"]] = df.apply(
		sort_by_scores, axis=1, result_type="expand"
	)
	results = select_top_k(df, ["contents", "ids", "scores"], top_k)

	return (
		results["contents"].tolist(),
		results["ids"].tolist(),
		results["scores"].tolist(),
	)


def get_colbert_embedding_batch(
	input_strings: List[str], model, tokenizer, batch_size: int
) -> List[np.array]:
	encoding = tokenizer(
		input_strings,
		return_tensors="pt",
		padding=True,
		truncation=True,
		max_length=model.config.max_position_embeddings,
	)

	input_batches = slice_tokenizer_result(encoding, batch_size)
	result_embedding = []
	with torch.no_grad():
		for encoding_batch in tqdm(input_batches):
			result_embedding.append(model(**encoding_batch).last_hidden_state)
	total_tensor = torch.cat(
		result_embedding, dim=0
	)  # shape [batch_size, token_length, embedding_dim]
	tensor_results = list(total_tensor.chunk(total_tensor.size()[0]))

	if torch.cuda.is_available():
		return list(map(lambda x: x.detach().cpu().numpy(), tensor_results))
	else:
		return list(map(lambda x: x.detach().numpy(), tensor_results))


def slice_tokenizer_result(tokenizer_output, batch_size):
	input_ids_batches = slice_tensor(tokenizer_output["input_ids"], batch_size)
	attention_mask_batches = slice_tensor(
		tokenizer_output["attention_mask"], batch_size
	)
	token_type_ids_batches = slice_tensor(
		tokenizer_output.get("token_type_ids", None), batch_size
	)
	return [
		{
			"input_ids": input_ids,
			"attention_mask": attention_mask,
			"token_type_ids": token_type_ids,
		}
		for input_ids, attention_mask, token_type_ids in zip(
			input_ids_batches, attention_mask_batches, token_type_ids_batches
		)
	]


def slice_tensor(input_tensor, batch_size):
	# Calculate the number of full batches
	num_full_batches = input_tensor.size(0) // batch_size

	# Slice the tensor into batches
	tensor_list = [
		input_tensor[i * batch_size : (i + 1) * batch_size]
		for i in range(num_full_batches)
	]

	# Handle the last batch if it's smaller than batch_size
	remainder = input_tensor.size(0) % batch_size
	if remainder:
		tensor_list.append(input_tensor[-remainder:])

	device = "cuda" if torch.cuda.is_available() else "cpu"
	tensor_list = list(map(lambda x: x.to(device), tensor_list))

	return tensor_list


def get_colbert_score(query_embedding: np.array, content_embedding: np.array) -> float:
	if query_embedding.ndim == 3 and content_embedding.ndim == 3:
		query_embedding = query_embedding.reshape(-1, query_embedding.shape[-1])
		content_embedding = content_embedding.reshape(-1, content_embedding.shape[-1])

	sim_matrix = np.dot(query_embedding, content_embedding.T) / (
		np.linalg.norm(query_embedding, axis=1)[:, np.newaxis]
		* np.linalg.norm(content_embedding, axis=1)
	)
	max_sim_scores = np.max(sim_matrix, axis=1)
	return float(np.mean(max_sim_scores))
