import streamlit as st
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
from typing import Optional
from my_model.image_captioning import ImageCaptioningModel
from my_model.object_detection import ObjectDetector


class KBVQA():

    def __init__(self):
        self.kbvqa_model_name = "m7mdal7aj/fine_tunned_llama_2_merged"
        self.quantization='4bit'
        self.bnb_config = self.create_bnb_config()
        self.max_context_window = 4096
        self.add_eos_token = False
        self.trust_remote = False
        self.use_fast = True
        self.kbvqa_tokenizer = None
        self.captioner = None
        self.detector = None
        self.kbvqa_model = None
      #  self.kbvqa_model_loaded = self.all_models_loaded() 

 
    def create_bnb_config(self) -> BitsAndBytesConfig:
    """
    Creates a BitsAndBytes configuration based on the quantization setting.
    Returns:
        BitsAndBytesConfig: Configuration for BitsAndBytes optimized model.
    """
        if self.quantization == '4bit':
            return BitsAndBytesConfig(
                load_in_4bit=True,
                bnb_4bit_use_double_quant=True,
                bnb_4bit_quant_type="nf4",
                bnb_4bit_compute_dtype=torch.bfloat16
            )
        elif self.quantization == '8bit':
            return BitsAndBytesConfig(
                load_in_8bit=True,
                bnb_8bit_use_double_quant=True,
                bnb_8bit_quant_type="nf4",
                bnb_8bit_compute_dtype=torch.bfloat16
            )


    def load_caption_model(self):
        self.captioner = ImageCaptioningModel(model_type='i_blip')
        self.captioner.load_model()

    def get_caption(self, img):

        return self.captioner.generate_caption(img)

    def load_detector(self, model):

        self.detector = ObjectDetector()
        self.detector.load_model(model)

    def detect_objects(self, img, threshold=0.2):
        image = self.detector.process_image(img)
        detected_objects_string, detected_objects_list = self.detector.detect_objects(image, threshold=threshold)
        image_with_boxes = self.detector.draw_boxes(img, detected_objects_list)
        return image_with_boxes, detected_objects_string

    def load_fine_tuned_model(self):

        self.kbvqa_model = AutoModelForCausalLM.from_pretrained(self.kbvqa_model_name, device_map="auto", quantization_config=self.bnb_config)
        self.kbvqa_tokenizer = AutoTokenizer.from_pretrained(self.kbvqa_model_name, use_fast=self.use_fast, trust_remote_code=self.trust_remote, add_eos_token=self.add_eos_token)


    @property
    def all_models_loaded(self):
        return self.kbvqa_model is not None and self.captioner is not None and self.detector is not None



    def format_prompt(self, current_query, history = None , sys_prompt=None, caption=None, objects=None):

        if sys_prompt is None:
            sys_prompt = "You are a helpful, respectful and honest assistant for visual question answering. you are provided with a caption of an image and a list of objects detected in the image along with their bounding boxes and level of certainty, you will output an answer to the given questions in no more than one sentence. Use logical reasoning to reach to the answer, but do not output your reasoning process unless asked for it. If provided, you will use the [CAP] and [/CAP] tags to indicate the begining and end of the caption respectively. If provided you will use the [OBJ] and [/OBJ] tags to indicate the begining and end of the list of detected objects in the image along with their bounding boxes respectively.if provided, you will use [QES] and [/QES] tags to indicate the begining and end of the question respectively."
    
        B_SENT = '<s>'
        E_SENT = '</s>'
        B_INST = '[INST]'
        E_INST = '[/INST]'
        B_SYS = '<<SYS>>\n'
        E_SYS = '\n<</SYS>>\n\n'
        B_CAP = '[CAP]'
        E_CAP = '[/CAP]'
        B_QES = '[QES]'
        E_QES = '[/QES]'
        B_OBJ = '[OBJ]'
        E_OBJ = '[/OBJ]'
    
    
        current_query = current_query.strip()
        sys_prompt = sys_prompt.strip()
        
        if history is None:
            if objects is None:
                p = f"""{B_SENT}{B_INST} {B_SYS}{sys_prompt}{E_SYS}{B_CAP}{caption}{E_CAP}{B_QES}{current_query}{E_QES}{E_INST}"""
            else:
              p = f"""{B_SENT}{B_INST} {B_SYS}{sys_prompt}{E_SYS}{B_CAP}{caption}{E_CAP}{B_OBJ}{objects}{E_OBJ}{B_QES}taking into consideration the objects with high certainty, {current_query}{E_QES}{E_INST}"""
        else:
            p = f"""{history}\n{B_SENT}{B_INST} {B_QES}{current_query}{E_QES}{E_INST}"""
        
        
        return p
       

    def generate_answer(self, question, image):
        caption = self.get_caption(image)
        image_with_boxes, detected_objects_str = self.detect_objects(image)
        prompt = self.format_prompt(question, caption=caption, objects=detected_objects_str)
        num_tokens = len(self.kbvqa_tokenizer.tokenize(prompt))
        if num_tokens > self.max_context_window:
            st.write(f"Prompt too long with {num_tokens} tokens, consider increasing the confidence threshold for the object detector")
            return

        model_inputs = self.kbvqa_tokenizer(prompt, add_special_tokens=False, return_tensors="pt").to('cuda')
        input_ids = model_inputs["input_ids"]
        output_ids = self.kbvqa_model.generate(input_ids)
        index = input_ids.shape[1] # needed to avoid printing the input prompt
        history = self.kbvqa_tokenizer.decode(output_ids[0], skip_special_tokens=False)
        output_text = self.kbvqa_tokenizer.decode(output_ids[0][index:], skip_special_tokens=True)

        return output_text.capitalize()

def prepare_kbvqa_model(detection_model):
    kbvqa = KBVQA()
    # Progress bar for model loading
    with st.spinner('Loading models...'):
        progress_bar = st.progress(0)
        kbvqa.load_fine_tuned_model()
        progress_bar.progress(33)
        kbvqa.load_caption_model()
        progress_bar.progress(66)
        kbvqa.load_detector(detection_model)  # Replace with your model
        progress_bar.progress(100)
    
    if kbvqa.all_models_loaded:
        st.success('Model loaded successfully!')
        kbvqa.kbvqa_model.eval()
        return kbvqa