"""
File: app.py

Description: Visual question answering (with pixtral-12b-2409)

Author: Didier Guillevic
Date: 2024-09-29
"""
import logging
logger = logging.getLogger(__name__)

import gradio as gr
import PIL
import base64
from io import BytesIO

import os
from mistralai import Mistral

#
# Mistral AI client
#
api_key = os.environ["MISTRAL_API_KEY"]
client = Mistral(api_key=api_key)
model_id = "pixtral-12b-2409"

#
# Utils
#
def image_to_base64(filepath):
    file_format = filepath.split('.')[-1].lower()
    with open(filepath, "rb") as image_file:
        base64_string = base64.b64encode(image_file.read()).decode('utf-8')
    return base64_string, file_format

def pil_to_base64(pil_image, image_format='jpeg'):  # Default to 'png' if format not known
    # Save the PIL image to a BytesIO object in the specified format
    buffered = BytesIO()
    pil_image.save(buffered, format=image_format)
    
    # Convert the image to a base64 string
    base64_string = base64.b64encode(buffered.getvalue()).decode('utf-8')
    
    return base64_string, image_format


#
# Generate a response
#
def answer_question(text, image, temperature=0.0, max_tokens=1_024):
    # Convert image to base64 string
    #base64_string, file_format = image_to_base64(image)
    base64_string, file_format = pil_to_base64(image)

    messages = [
        {
            "role": "user",
            "content": [
                {
                    "type": "text",
                    "text": text
                },
                {
                    "type": "image_url",
                    "image_url": f"data:image/{file_format};base64,{base64_string}" 
                }
            ]
        }
    ]
    
    chat_response = client.chat.complete(
        model=model_id,
        messages=messages,
        temperature=temperature,
        max_tokens=max_tokens
    )

    return chat_response.choices[0].message.content


def clear_all():
    # Clear question, image, output_text
    return ('', None, '')


#
# User interface
#
with gr.Blocks() as demo:
    
    gr.Markdown("""
        ## Visual question answering / image captioning
    """)
    with gr.Row():
        with gr.Column():
            question = gr.Textbox(
                placeholder="Ask anything about the image..",
                lines=2,
                render=True)
            image = gr.Image(type="pil") # return type: 'numpy', 'pil', 'filepath'
            with gr.Row():
                temperature = gr.Slider(
                    label="Temperature",
                    minimum=0.0,
                    maximum=1.0,
                    value=0.0,
                    step=0.1
                )
                max_tokens = gr.Slider(
                    label="Max tokens",
                    minimum=128,
                    maximum=2_048,
                    value=1_024,
                    step=128
                )
        with gr.Column():
            output_text = gr.Textbox(
                lines=10,
                label="Pixtral 12B",
                render=True
            )
    with gr.Row():
        clear_btn = gr.Button("Clear", variant="secondary")
        submit_btn = gr.Button(value="Submit", variant="primary")

    # Examples
    examples = gr.Examples(
        [
            [
                ('Can you describe this image? If this is the image a '
                 'government issued ID, extract the information and format '
                 'it as a JSON document.'),
                './sample_ID.jpeg'
            ],
            [
                'In which city and country was  this picture taken?',
                './16-park-ave-milton-park-1024x683.jpg'
            ],
            [   (
                    'Please extract the identification information present in the image. '
                    'Format the response in a JSON format such as "{name_of_element : [value]}".'
                ),
                './driver_license.png'
            ],
        ],
        inputs=[question, image],
        outputs=[output_text,],
        fn=answer_question,
        cache_examples=False,
        label="Examples"
    )

    # Click actions
    clear_btn.click(
        clear_all,
        inputs=[],
        outputs=[question, image, output_text]
    )
    submit_btn.click(
        fn=answer_question,
        inputs=[question, image, temperature, max_tokens],
        outputs=[output_text,]
    )

    # Documentation
    with gr.Accordion("Documentation", open=False):
        gr.Markdown("""
            - model: Serving pixtral-12b (https://mistral.ai/news/pixtral-12b/)
            - temperature: 0.0=little variable in output, 1.0=lots of variability in output
            - max_tokens: maximum number of tokens to generate as output (1 word approx 1.5 tokens)
        """)

demo.launch(show_api=False)