import numpy as np


prompt_mapping = {
	"qwen": (
		"<|im_start|>system\nPlease reason step by step, and put your final answer within \\boxed{{}}.<|im_end|>\n"
		"<|im_start|>user\n{input}<|im_end|>\n"
		"<|im_start|>assistant\n"
	),
	"qwen_deepseek_distill": (
		"<｜begin▁of▁sentence｜>Please reason step by step, and put your final answer within \\boxed{{}}"
		"<｜User｜>{input}"
		"<｜Assistant｜><think>",   
		"{COT}</think>",
		"{output}<｜end▁of▁sentence｜>"
	),

	"qwen_deepseek_distill_endThink": (
		"<｜begin▁of▁sentence｜>Please reason step by step, and put your final answer within \\boxed{{}}"
		"<｜User｜>{input}"
		"<｜Assistant｜><think></think>",   
		"{COT}",
		"{output}<｜end▁of▁sentence｜>"
	),
}


def ensure_add_token(tokenizer, llm, token="<T>"):
    if "<T>" not in tokenizer.get_vocab():
        tokenizer.add_special_tokens({"additional_special_tokens": [token]})
        llm.resize_token_embeddings(len(tokenizer))
    else:
        print(f"{token} already in tokenizer")

def get_precise_equidistant_points(lst, n=7):
	length = len(lst)
	if length <= n + 1:
		return lst.copy()
	interval_size = (length - 1) / n
	indices = [0]
	for i in range(1, n):
		pos = i * interval_size
		idx = int(round(pos))
		if idx <= indices[-1]:
			idx = indices[-1] + 1
		indices.append(idx)
	indices.append(length - 1)
	return [lst[i].cpu() for i in indices]

def split_list(data_list, n):
    return np.array_split(data_list, n)
