"""Template Demo for IBM Granite Hugging Face spaces."""

 from collections.abc import Iterator
 from datetime import datetime
 from pathlib import Path
 from threading import Thread

 import gradio as gr
 import spaces
 import torch
 from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer

 from themes.research_monochrome import theme

 # Vision imports
 import random
 from transformers import LlavaNextProcessor, LlavaNextForConditionalGeneration

 today_date = datetime.today().strftime("%B %-d, %Y")  # noqa: DTZ002

 SYS_PROMPT = f"""Knowledge Cutoff Date: April 2024.
 Today's Date: {today_date}.
 You are Granite, developed by IBM. You are a helpful AI assistant"""
 TITLE = "IBM Granite 3.1 8b Instruct & Vision Preview"
 DESCRIPTION = """
 <p>Granite 3.1 8b instruct is an open-source LLM supporting a 128k context window. Start with one of the sample prompts
 or upload an image and ask a question. Keep in mind that AI can occasionally make mistakes.
 <span class="gr_docs_link">
 <a href="https://www.ibm.com/granite/docs/">View Documentation <i class="fa fa-external-link"></i></a>
 </span>
 </p>
 """
 MAX_INPUT_TOKEN_LENGTH = 128_000
 MAX_NEW_TOKENS = 1024
 TEMPERATURE = 0.7
 TOP_P = 0.85
 TOP_K = 50
 REPETITION_PENALTY = 1.05

 if not torch.cuda.is_available():
     print("This demo may not work on CPU.")

 # Text Model and Tokenizer
 text_model = AutoModelForCausalLM.from_pretrained(
     "ibm-granite/granite-3.1-8b-instruct", torch_dtype=torch.float16, device_map="auto"
 )
 text_tokenizer = AutoTokenizer.from_pretrained("ibm-granite/granite-3.1-8b-instruct")
 text_tokenizer.use_default_system_prompt = False

 # Vision Model and Processor
 vision_model_path = "ibm-granite/granite-vision-3.1-2b-preview"
 vision_processor = LlavaNextProcessor.from_pretrained(vision_model_path, use_fast=True)
 vision_model = LlavaNextForConditionalGeneration.from_pretrained(vision_model_path, torch_dtype="auto", device_map="auto")


 @spaces.GPU
 def generate(
     message: str,
     chat_history: list[dict],
     temperature: float = TEMPERATURE,
     repetition_penalty: float = REPETITION_PENALTY,
     top_p: float = TOP_P,
     top_k: float = TOP_K,
     max_new_tokens: int = MAX_NEW_TOKENS,
 ) -> Iterator[str]:
     """Generate function for text chat demo."""
     # Build messages
     conversation = []
     conversation.append({"role": "system", "content": SYS_PROMPT})
     conversation += chat_history
     conversation.append({"role": "user", "content": message})

     # Convert messages to prompt format
     input_ids = text_tokenizer.apply_chat_template(
         conversation,
         return_tensors="pt",
         add_generation_prompt=True,
         truncation=True,
         max_length=MAX_INPUT_TOKEN_LENGTH - max_new_tokens,
     )

     input_ids = input_ids.to(text_model.device)
     streamer = TextIteratorStreamer(text_tokenizer, skip_prompt=True, skip_special_tokens=True)
     generate_kwargs = dict(
         {"input_ids": input_ids},
         streamer=streamer,
         max_new_tokens=max_new_tokens,
         do_sample=True,
         top_p=top_p,
         top_k=top_k,
         temperature=temperature,
         num_beams=1,
         repetition_penalty=repetition_penalty,
     )

     t = Thread(target=text_model.generate, kwargs=generate_kwargs)
     t.start()

     outputs = []
     for text in streamer:
         outputs.append(text)
         yield "".join(outputs)

def get_text_from_content(content):
    texts = []
    for item in content:
        if item["type"] == "text":
            texts.append(item["text"])
        elif item["type"] == "image":
            texts.append("[Image]")
    return " ".join(texts)

@spaces.GPU
def chat_inference(image, text, temperature, top_p, top_k, max_tokens, conversation):
    if conversation is None:
        conversation = []

    user_content = []
    if image is not None:
        user_content.append({"type": "image", "image": image})
    if text and text.strip():
        user_content.append({"type": "text", "text": text.strip()})
    if not user_content:
        return conversation_display(conversation), conversation

    conversation.append({
        "role": "user",
        "content": user_content
    })

    inputs = vision_processor.apply_chat_template(
        conversation,
        add_generation_prompt=True,
        tokenize=True,
        return_dict=True,
        return_tensors="pt"
    ).to("cuda")

    torch.manual_seed(random.randint(0, 10000))

    generation_kwargs = {
        "max_new_tokens": max_tokens,
        "temperature": temperature,
        "top_p": top_p,
        "top_k": top_k,
        "do_sample": True,
    }

    output = vision_model.generate(**inputs, **generation_kwargs)
    assistant_response = vision_processor.decode(output[0], skip_special_tokens=True)

    conversation.append({
        "role": "assistant",
        "content": [{"type": "text", "text": assistant_response.strip()}]
    })

    return conversation_display(conversation), conversation

def conversation_display(conversation):
    chat_history = []
    for msg in conversation:
        if msg["role"] == "user":
            user_text = get_text_from_content(msg["content"])
            chat_history.append({"role": "user", "content": user_text})
        elif msg["role"] == "assistant":
            assistant_text = msg["content"][0]["text"].split("<|assistant|>")[-1].strip()
            chat_history.append({"role": "assistant", "content": assistant_text})
    return chat_history

def clear_chat():
    return [], [], "", None

css_file_path = Path(Path(__file__).parent / "app.css")
head_file_path = Path(Path(__file__).parent / "app_head.html")

# Advanced settings (displayed in Accordion) - Common settings for both models
temperature_slider = gr.Slider(
    minimum=0, maximum=1.0, value=TEMPERATURE, step=0.1, label="Temperature", elem_classes=["gr_accordion_element"]
)
top_p_slider = gr.Slider(
    minimum=0, maximum=1.0, value=TOP_P, step=0.05, label="Top P", elem_classes=["gr_accordion_element"]
)
top_k_slider = gr.Slider(
    minimum=0, maximum=100, value=TOP_K, step=1, label="Top K", elem_classes=["gr_accordion_element"]
)

# Advanced settings specific to Text model
repetition_penalty_slider = gr.Slider(
    minimum=0,
    maximum=2.0,
    value=REPETITION_PENALTY,
    step=0.05,
    label="Repetition Penalty (Text Model)",
    elem_classes=["gr_accordion_element"],
)
max_new_tokens_slider = gr.Slider(
    minimum=1,
    maximum=2000,
    value=MAX_NEW_TOKENS,
    step=1,
    label="Max New Tokens (Text Model)",
    elem_classes=["gr_accordion_element"],
)

# Advanced settings specific to Vision model
max_tokens_slider_vision = gr.Slider(
    minimum=10,
    maximum=300,
    value=128,
    step=1,
    label="Max Tokens (Vision Model)",
    elem_classes=["gr_accordion_element"],
)

chat_interface_accordion = gr.Accordion(label="Advanced Settings", open=False)

with gr.Blocks(fill_height=True, css_paths=css_file_path, head_paths=head_file_path, theme=theme, title=TITLE) as demo:
    gr.HTML(f"<h1>{TITLE}</h1>", elem_classes=["gr_title"])
    gr.HTML(DESCRIPTION)

    state = gr.State([]) # State for vision chat history
    chat_history_state = gr.State([]) # State for text chat history

    with gr.Row():
        with gr.Column(scale=2):
            image_input = gr.Image(type="pil", label="Upload Image (optional)")
            with gr.Accordion(label="Vision Model Settings", open=False):
                max_tokens_input_vision = max_tokens_slider_vision
            with gr.Accordion(label="Text Model Settings", open=False):
                repetition_penalty_input = repetition_penalty_slider
                max_new_tokens_input = max_new_tokens_slider
            with chat_interface_accordion: # Common Settings
                temperature_input = temperature_slider
                top_p_input = top_p_slider
                top_k_input = top_k_slider

        with gr.Column(scale=3):
            chatbot = gr.Chatbot(label="Chat History", elem_id="chatbot", type='messages')
            text_input = gr.Textbox(lines=2, placeholder="Enter your message here", label="Message")
            with gr.Row():
                send_button = gr.Button("Chat")
                clear_button = gr.Button("Clear Chat")

    def process_chat(image_input, text_input, temperature_input, top_p_input, top_k_input, repetition_penalty_input, max_new_tokens_input, max_tokens_input_vision, state, chat_history_state):
        if image_input:
            # Use Vision model
            return chat_inference(image_input, text_input, temperature_input, top_p_input, top_k_input, max_tokens_input_vision, state)
        else:
            # Use Text model
            return generate(text_input, chat_history_state, temperature_input, repetition_penalty_input, top_p_input, top_k_input, max_new_tokens_input), None # Return None for state as text model doesn't use it

    def process_chat_wrapper(image_input_val, text_input_val, temperature_input_val, top_p_input_val, top_k_input_val, repetition_penalty_input_val, max_new_tokens_input_val, max_tokens_input_vision_val, state_val, chat_history_state_val):
        if image_input_val:
            chatbot_output, updated_state = process_chat(image_input_val, text_input_val, temperature_input_val, top_p_input_val, top_k_input_val, repetition_penalty_input_val, max_new_tokens_input_val, max_tokens_input_vision_val, state_val, chat_history_state_val)
            return chatbot_output, updated_state, chat_history_state_val # Return vision state and keep text state unchanged
        else:
            chatbot_output_generator, _ = process_chat(image_input_val, text_input_val, temperature_input_val, top_p_input_val, top_k_input_val, repetition_penalty_input_val, max_new_tokens_input_val, max_tokens_input_vision_val, state_val, chat_history_state_val)
            updated_chat_history = []
            full_response = ""
            for response_chunk in chatbot_output_generator:
                full_response = response_chunk
            if chat_history_state_val is None:
                updated_chat_history = []
            else:
                updated_chat_history = chat_history_state_val

            updated_chat_history.append({"role": "user", "content": text_input_val})
            updated_chat_history.append({"role": "assistant", "content": full_response})

            return updated_chat_history, state_val, updated_chat_history # Return text chat history, keep vision state unchanged, return updated text history for chatbot display


    send_button.click(
        process_chat_wrapper,
        inputs=[image_input, text_input, temperature_input, top_p_input, top_k_input, repetition_penalty_input, max_new_tokens_input, max_tokens_input_vision, state, chat_history_state],
        outputs=[chatbot, state, chat_history_state] # Keep both states as output
    )

    clear_button.click(
        clear_chat,
        inputs=None,
        outputs=[chatbot, state, text_input, image_input] # clear_chat clears vision state and input. Need to clear text state also.
    )

    gr.Examples(
        examples=[
            ["Explain the concept of quantum computing to someone with no background in physics or computer science."],
            ["What is OpenShift?"],
            ["What's the importance of low latency inference?"],
            ["Help me boost productivity habits."],
            [
                """Explain the following code in a concise manner:

```java
import java.util.ArrayList;
import java.util.List;

public class Main {

    public static void main(String[] args) {
        int[] arr = {1, 5, 3, 4, 2};
        int diff = 3;
        List<Pair> pairs = findPairs(arr, diff);
        for (Pair pair : pairs) {
            System.out.println(pair.x + " " + pair.y);
        }
    }

    public static List<Pair> findPairs(int[] arr, int diff) {
        List<Pair> pairs = new ArrayList<>();
        for (int i = 0; i < arr.length; i++) {
            for (int j = i + 1; j < arr.length; j++) {
                if (Math.abs(arr[i] - arr[j]) < diff) {
                    pairs.add(new Pair(arr[i], arr[j]));
                }
            }
        }

        return pairs;
    }
}

class Pair {
    int x;
    int y;
    public Pair(int x, int y) {
        this.x = x;
        this.y = y;
    }
}
```"""
            ],
            [
                """Generate a Java code block from the following explanation:

The code in the Main class finds all pairs in an array whose absolute difference is less than a given value.

The findPairs method takes two arguments: an array of integers and a difference value. It iterates over the array and compares each element to every other element in the array. If the absolute difference between the two elements is less than the difference value, a new Pair object is created and added to a list.

The Pair class is a simple data structure that stores two integers.

The main method creates an array of integers, initializes the difference value, and calls the findPairs method to find all pairs in the array. Finally, the code iterates over the list of pairs and prints each pair to the console."""  # noqa: E501
            ],
             ["https://raw.githubusercontent.com/gradio-app/gradio/main/test/test_files/bus.png", "What is this?"] # Vision example
        ],
        inputs=[text_input, text_input, text_input, text_input, text_input, text_input, image_input, image_input] , # Duplicated text_input to match example count, last two are image_input for vision example
        examples_per_page=7
    )

if __name__ == "__main__":
    demo.queue().launch()