import torch
import time
import gradio as gr
from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig, GenerationConfig, pipeline, AutoModelForSpeechSeq2Seq, AutoProcessor, TextStreamer
from langchain import PromptTemplate
from langchain import HuggingFacePipeline
from langchain.embeddings import HuggingFaceEmbeddings
from langchain.document_loaders import PyPDFLoader, PyPDFDirectoryLoader
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain.vectorstores import Chroma
from langchain.chains import RetrievalQA, ConversationalRetrievalChain
from langchain.memory import ConversationBufferMemory
from diffusers import LCMScheduler, AutoPipelineForText2Image
import tensorflow as tf
import keras_cv
from keras.applications.imagenet_utils import preprocess_input
import fitz
from PIL import Image

import warnings
warnings.filterwarnings("ignore")

model_name="SrikanthChellappa/llama2-7b-sharded-finetuning-medical100k-chatbot"
#model_name="SrikanthChellappa/mistral-7b-sharded-finetuning-medical100k-chatbot"
#model_name="SrikanthChellappa/orca2-7b-sharded-finetuning-medical100k-chatbot"
#model_name="SrikanthChellappa/falcon-7b-sharded-finetuning-medical100k-chatbot"
bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_use_double_quant=True,
    bnb_4bit_compute_dtype=torch.bfloat16,
)
#bnb_4bit_compute_dtype=torch.float16,
tokenizer = AutoTokenizer.from_pretrained(model_name, use_fast=True, token="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")
tokenizer.pad_token = tokenizer.eos_token
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    quantization_config=bnb_config,
    device_map="auto",
    trust_remote_code=True,
    token="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"
)
#model = AutoModelForCausalLM.from_pretrained('./mistral_model_7B_8bit_Q/', quantization_config=quantization_config)
generation_config = GenerationConfig(max_new_tokens=1024, pad_token_id = tokenizer.eos_token_id, \
                 eos_token_id = tokenizer.eos_token_id,\
                temperature=0.0001, top_p=0.95, repetition_penalty=1.15, do_sample = True, num_return_sequences=1,)

streamer = TextStreamer(tokenizer)
pipeline = pipeline(
    "text-generation",
    model=model,
    tokenizer=tokenizer,
    streamer=streamer,
    return_full_text=True,
    generation_config=generation_config,
)

def identify_image(chat_history,image_path):
    #image_path=r"D:\Generative AI\img-3.jpg"
    imgmodel = tf.keras.applications.resnet50.ResNet50()
    #imgmodel = model_resnet50
    image_size=(224,224)
    tf_image = tf.io.read_file(image_path)
    # Decode the above `tf_image` from a Bytes string to a numeric Tensor.
    decoded_image = tf.image.decode_image(tf_image)
    # Resize the image to the spatial size required by the model.
    image_resized = tf.image.resize(decoded_image, image_size)
    # Add a batch dimension to the first axis (required). 
    image_batch = tf.expand_dims(image_resized, axis=0)
    # Pre-process the input image.
    image_batch = preprocess_input(image_batch)
    # Forward pass through the model to make predictions.
    preds = imgmodel.predict(image_batch)
    decoded_preds = tf.keras.applications.imagenet_utils.decode_predictions(
                preds=preds,
                top=5
            )
    label = decoded_preds[0][0][1]
    score = decoded_preds[0][0][2] * 100
    title = 'The image is of a ' + label + ' with a prediction accuracy of ' + str('{:.2f}%'.format(score))
    chat_history.append((title, [image_path,title]))
    return chat_history

llm = HuggingFacePipeline(
    pipeline=pipeline,
)
embeddings = HuggingFaceEmbeddings(
    model_name="thenlper/gte-large",
    model_kwargs={"device": "cuda"},
    encode_kwargs={"normalize_embeddings": True},
)
persist_dir=r"D:\Generative AI\Langchain-RAG(RetrievalAugmentedGen)\davidson-medicine-db"
db = Chroma(persist_directory=persist_dir, embedding_function=embeddings)

def querying(query, chat_history,file):
      return "", chat_history + [[query, None]]
       
def render_pdf(file):
        doc = fitz.open(file.name)
        page = doc[0]
        #Render the page as a PNG image with a resolution of 300 DPI
        pix = page.get_pixmap(matrix=fitz.Matrix(300/72, 300/72))
        image = Image.frombytes('RGB', [pix.width, pix.height], pix.samples)
        return image

def transcribe_audio(chat_history,file):
        audio_result = audio_pipe(file,batch_size=8)
        chat_history.append((file.name, audio_result["text"]))
        return chat_history

def bot(query,history,file):
        #bot_message = random.choice(["How are you?", "I love you", "I'm very hungry"])
        query=history[len(history)-1][0]
        #matches = ["generate", "Generate","GENERATE","image", "images", "IMAGE"]
        matches = ["generate image", "Generate Image","GENERATE IMAGE","create image", "Create Image","CREATE IMAGE","Generate an image", "generate an image", "Generate An Image","GENERATE AN IMAGE","create an image", "Create An Image","CREATE AN IMAGE","generate a image", "Generate A Image","GENERATE A IMAGE","create a image", "Create A Image","CREATE A IMAGE"]
        models = ["mistral", "llama2"]
        if any(x in query for x in matches):
            genimg_model = keras_cv.models.StableDiffusion(img_width=512, img_height=512, jit_compile=True)
            image = genimg_model.text_to_image(query, batch_size=1, num_steps=15, seed=123)
            Image.fromarray(image[0]).save("tmp.png")
            #history.append((query, ["tmp.png",query[:15]]))
            history[len(history)-1][0]=query
            history[len(history)-1][1]=["tmp.png",None]
            yield history
        else:
            if "orca" in model_name:
                template = """
                <|im_start|>system
                You are a helpful, respectful, and honest human assistant. Always answer as 
                helpfully as possible while being safe.  Your answers should not 
                include any harmful, unethical, racist, sexist, toxic, dangerous, or 
                illegal content. Please ensure that your responses are socially 
                unbiased and positive in nature.

                If a question does not make any sense or is not factually coherent, 
                explain why instead of answering something incorrectly. If you don't 
                know the answer to a question, don't share false information.
                <|im_end|>
                <|im_start|>user
                Act as a human assistant. Use the following information to answer the question at the end.
                {context}

                {question} 
                <|im_end|>
                <|im_start|>assistant
                """
            elif "mistral" in model_name:
                template = """
                <s>[INST] <> <<SYS>>
                You are a helpful, respectful, and honest human assistant. Always answer as 
                helpfully as possible while being safe.  Your answers should not 
                include any harmful, unethical, racist, sexist, toxic, dangerous, or 
                illegal content. Please ensure that your responses are socially 
                unbiased and positive in nature.

                If a question does not make any sense or is not factually coherent, 
                explain why instead of answering something incorrectly. If you don't 
                know the answer to a question, don't share false information.
                <</SYS>>
                Act as a human assistant. Use the following information to answer the question at the end.
                <>
                {context}

                {question} [/INST]
                """
            elif "llama2" in model_name:
                template = """
                [INST] <> <<SYS>>
                You are a helpful, respectful, and honest human assistant. Always answer as 
                helpfully as possible while being safe.  Your answers should not 
                include any harmful, unethical, racist, sexist, toxic, dangerous, or 
                illegal content. Please ensure that your responses are socially 
                unbiased and positive in nature.

                If a question does not make any sense or is not factually coherent, 
                explain why instead of answering something incorrectly. If you don't 
                know the answer to a question, don't share false information.
                <</SYS>>
                Act as a human assistant. Use the following information to answer the question at the end.
                <>

                {context}

                {question} [/INST]
                """
            else:
                template = """
                {context}
                {question}
                """
            tmpdb=None
            if file:
                loader = PyPDFLoader(file.name)
                documents = loader.load()
                text_splitter = RecursiveCharacterTextSplitter(chunk_size=1024, chunk_overlap=64)
                texts_chunks = text_splitter.split_documents(documents)
                tmpdb = Chroma.from_documents(texts_chunks, embeddings, persist_directory=f"{file.name}-db")
                
            prompt = PromptTemplate(template=template, input_variables=["context", "question"])
            if tmpdb:
                qa_chain = RetrievalQA.from_chain_type(
                    llm=llm,
                    chain_type="stuff",
                    retriever=tmpdb.as_retriever(search_kwargs={"k": 2}),
                    return_source_documents=True,
                    chain_type_kwargs={"prompt": prompt},
                )
            else:
                qa_chain = RetrievalQA.from_chain_type(
                    llm=llm,
                    chain_type="stuff",
                    retriever=db.as_retriever(search_kwargs={"k": 2}),
                    return_source_documents=True,
                    chain_type_kwargs={"prompt": prompt},
                )
            result_ = qa_chain(
               query
            )
            result = result_["result"].strip()
            history[-1][1] = ""
            for character in result:
                history[-1][1] += str(character)
                time.sleep(0.03)
                yield history

theme = gr.themes.Default(primary_hue="blue").set(
    loader_color="#FF0000",
    slider_color="#FF0000",
)
def vote(data: gr.LikeData):
    if data.liked:
        print("You upvoted this response: " + data.value)
    else:
        print("You downvoted this response: " + data.value)
        
with gr.Blocks(theme=theme) as demo:
    gr.Markdown('Agasthya - Your Bot Assistant')
    with gr.Row():
        with gr.Column(scale=1.0):
            chatbot = gr.Chatbot(value=[], elem_id='chatbot',height=500)
        #with gr.Column(scale=0.30):
            #show_img = gr.Image(label='Upload PDF', height=500)
    with gr.Row():
        with gr.Column(scale=0.70):
            txt = gr.Textbox(
                        show_label=False,
                        placeholder="Enter text and press enter",
                        container=False
                    )
            examples=[
                ["Generate image of a boy playing cricket"],
                ["How do you treat a painful wart?"],
                ["What are Flynn's calssification of computer architecture?"],
            ],
        with gr.Column(scale=0.30):
            submit_btn = gr.Button('submit')
        #with gr.Column(scale=0.10):
            #audio_btn=gr.Audio(sources=["microphone"])
    with gr.Row():
        with gr.Column(scale=0.34):
            imgbtn=gr.UploadButton("Identify My Image", file_types=[".jpg,.png,.jpeg"],interactive=False)
            imgbtn.upload(identify_image, [chatbot,imgbtn],[chatbot])
        with gr.Column(scale=0.34):
            btn=gr.UploadButton("Q&A With My PDF", file_types=[".pdf"],interactive=False)
            #btn.upload(render_pdf, [btn],show_img)
        with gr.Column(scale=0.34):
            transcribe_btn=gr.UploadButton("Transcribe My Audio", file_types=[".ogg,.mp3,.wav"],interactive=False)
            transcribe_btn.upload(transcribe_audio,[chatbot,transcribe_btn],[chatbot])
            
        chatbot.like(vote, None, None)
    
    submit_btn.click(
            fn=querying, 
            inputs=[txt,chatbot,btn], 
            outputs = [txt,chatbot],
            queue=False
    ).then(
        bot, [txt,chatbot,btn], chatbot
    )
    txt.submit(
            fn=querying, 
            inputs=[txt,chatbot,btn], 
            outputs = [txt,chatbot],
            queue=False
    ).then(
        bot, [txt,chatbot,btn], chatbot
    )

demo.queue()
demo.launch(share=True)