from llama_index.core.prompts import PromptTemplate
from llama_index.llms.huggingface import HuggingFaceLLM
from llama_index.core import Settings
import torch
import os
import spaces
import torch
from transformers import BitsAndBytesConfig

@spaces.GPU
def setGPU():
    torch.cuda.empty_cache()
    os.environ["CUDA_VISIBLE_DEVICES"] = "3,4"
    os.environ['CUDA_LAUNCH_BLOCKING'] = '1'

    if torch.cuda.is_available():
        print(f"Numero di GPU disponibili: {torch.cuda.device_count()}")
        for i in range(torch.cuda.device_count()):
            print(f"GPU {i}: {torch.cuda.get_device_name(i)}")
        current_device = torch.cuda.current_device()
        print(f"GPU in uso: {current_device}, {torch.cuda.get_device_name(current_device)}")
    else:
        print("CUDA non disponibile. Utilizzando CPU.")

        
def setLLM():
    # Define the quantization configuration
    quantization_config = BitsAndBytesConfig(
        load_in_4bit=True,
        bnb_4bit_compute_dtype=torch.float16,
        bnb_4bit_quant_type="nf4",
        bnb_4bit_use_double_quant=True, 
    )
        
    # Define the generate_kwargs
    generate_kwargs = {
        "do_sample": True,
        "min_length": 50,  # Lunghezza minima della risposta generata
        "no_repeat_ngram_size": 5,  # Evita la ripetizione di n-grammi
    }
    
    # Define the prompt template
    prompt_template = PromptTemplate("<s> [INST] {query_str} [/INST] ")
    
    # Load the HuggingFaceLLM with specified configurations
    llm = HuggingFaceLLM(
        model_name="swap-uniba/LLaMAntino-2-chat-7b-hf-UltraChat-ITA",
        tokenizer_name="swap-uniba/LLaMAntino-2-chat-7b-hf-UltraChat-ITA",
        query_wrapper_prompt=prompt_template,
        context_window=3900,
        max_new_tokens=512,
        generate_kwargs=generate_kwargs,
        model_kwargs={"quantization_config": quantization_config},
        # tokenizer_kwargs={"token": hf_token},
        device_map="auto", # Automatically allocate the model to GPU if available
    )
    
    return llm

def setPromptTemplate():
    text_qa_template_str = (
        "Sei un chatbot in grado di rispondere solo alle domande su bandi regionali e avvisi della regione Puglia. Le informazioni di contesto recuperate da diverse sorgenti sono qua sotto.\n"
        "---------------------\n"
        "{context_str}\n"
        "---------------------\n"
        "Usa le informazioni di contesto sopra fornite e non la tua conoscenza pregressa per rispondere, l'unica regione che conosci è la regione Puglia. Se le informazioni di contesto non sono utili rispondi usando la tua conoscenza pregressa senza comunicare informazioni false."
        "rispondi alla seguente query usando le informazioni dei bandi della regione Puglia \n"
        "Query: {query_str}\n"
        "Risposta: "
    )

    refine_template_str = (
        "La domanda orginale è la seguente: {query_str}\n Abbiamo fornito la"
        " seguente risposta: {existing_answer}\nAbbiamo l'opportunità di aggiornare"
        " la risposta (solo se necessario) con il seguente contesto in più"
        " .\n------------\n{context_msg}\n------------\nUsando il nuovo"
        " contesto, aggiorna o ripeti la risposta.\n"
    )


    text_qa_template = PromptTemplate(text_qa_template_str)
    refine_template = PromptTemplate(refine_template_str)

    return text_qa_template, refine_template