import gradio as gr
import os

os.system('pip install dashscope -U')
import tempfile
from pathlib import Path
import secrets
import dashscope
from dashscope import MultiModalConversation, Generation

YOUR_API_TOKEN = os.getenv('YOUR_API_TOKEN')
dashscope.api_key = YOUR_API_TOKEN
math_messages = []
def process_image(image):
    global math_messages
    math_messages = [] # reset when upload image
    uploaded_file_dir = os.environ.get("GRADIO_TEMP_DIR") or str(
        Path(tempfile.gettempdir()) / "gradio"
    )
    os.makedirs(uploaded_file_dir, exist_ok=True)
    
    name = f"tmp{secrets.token_hex(20)}.jpg"
    filename = os.path.join(uploaded_file_dir, name)
    image.save(filename)

    
    # Use qwen-vl-max-0809 for OCR
    messages = [{
        'role': 'system',
        'content': [{'text': 'You are a helpful assistant.'}]
    }, {
        'role': 'user',
        'content': [
            {'image': f'file://{filename}'},
            {'text': 'Please describe the math-related content in this image, ensuring that any LaTeX formulas are correctly transcribed. Non-mathematical details do not need to be described.'}
        ]
    }]
    
    response = MultiModalConversation.call(model='qwen-vl-max-0809', messages=messages)
    
    os.remove(filename)
    
    return response.output.choices[0]["message"]["content"]

def get_math_response(image_description, user_question):
    global math_messages
    if not math_messages:
        math_messages.append({'role': 'system', 'content': 'You are a helpful math assistant.'})
    math_messages = math_messages[:1] + math_messages[1:][-4:]
    if image_description is not None:
        content = f'Image description: {image_description}\n\n'
    else:
        content = ''
    query = f"{content}User question: {user_question}"
    math_messages.append({'role': 'user', 'content': query})
    response = Generation.call(	
        model="qwen2-math-72b-instruct",
        messages=math_messages,	
        result_format='message',
        stream=True
    )
    answer = None
    for resp in response:
        if resp.output is None:
            continue
        answer = resp.output.choices[0].message.content
        yield answer.replace("\\", "\\\\")
    print(f'query: {query}\nanswer: {answer}')
    if answer is None:
        math_messages.pop()
    else:
        math_messages.append({'role': 'assistant', 'content': answer})


def math_chat_bot(image, question):
    if image is not None:
        image_description = process_image(image)
    else:
        image_description = None
    yield from get_math_response(image_description, question)

css = """
#qwen-md .katex-display { display: inline; }
#qwen-md .katex-display>.katex { display: inline; }
#qwen-md .katex-display>.katex>.katex-html { display: inline; }
"""

# Create interface
iface = gr.Interface(
    css=css,
    fn=math_chat_bot,
    inputs=[
        gr.Image(type="pil", label="upload image"),
        gr.Textbox(label="input your question")
    ],
    outputs=gr.Markdown(label="answer", latex_delimiters=[
        {"left": "\\(", "right": "\\)", "display": True},
        {"left": "\\begin\{equation\}", "right": "\\end\{equation\}", "display": True},
        {"left": "\\begin\{align\}", "right": "\\end\{align\}", "display": True},
        {"left": "\\begin\{alignat\}", "right": "\\end\{alignat\}", "display": True},
        {"left": "\\begin\{gather\}", "right": "\\end\{gather\}", "display": True},
        {"left": "\\begin\{CD\}", "right": "\\end\{CD\}", "display": True},
        {"left": "\\[", "right": "\\]", "display": True}
    ], elem_id="qwen-md"),
    title="📖 Qwen2 Math Demo",
    allow_flagging='never',
    description="""\
<p align="center"><img src="https://modelscope.oss-cn-beijing.aliyuncs.com/resource/qwen.png" style="height: 80px"/><p>"""
    """\
<center><font size=3>This WebUI is based on Qwen2-VL for OCR and Qwen2-Math for mathematical reasoning. You can input either images or texts of mathematical or arithmetic problems.</center>"""
)    

# Launch gradio application
iface.launch()