#!/usr/bin/env python
# encoding: utf-8
import spaces
import torch
import argparse
from transformers import AutoModel, AutoTokenizer
import gradio as gr
from PIL import Image
import os
import traceback
import re

# Argparser
parser = argparse.ArgumentParser(description='Food Calcium Analysis Demo')
parser.add_argument('--device', type=str, default='cuda', help='cuda or mps')
args = parser.parse_args()
device = args.device
assert device in ['cuda', 'mps']

# Load model
model_path = 'openbmb/MiniCPM-V-2_6'
model = AutoModel.from_pretrained(model_path, trust_remote_code=True, torch_dtype=torch.bfloat16)
model = model.to(device=device)
tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
model.eval()

ERROR_MSG = "Error, please retry"
model_name = 'Food Calcium Analyzer'

def encode_image(image):
    if not isinstance(image, Image.Image):
        image = Image.open(image).convert("RGB")
    max_size = 448*16 
    if max(image.size) > max_size:
        w, h = image.size
        if w > h:
            new_w, new_h = max_size, int(h * max_size / w)
        else:
            new_h, new_w = max_size, int(w * max_size / h)
        image = image.resize((new_w, new_h), resample=Image.BICUBIC)
    return image

@spaces.GPU(duration=120)
def chat(image, prompt, params=None):
    try:
        system_prompt = """You are an AI assistant specialized in analyzing food images and estimating calcium content. Follow these steps:
1. Identify the food in the image accurately.
2. List the main ingredients.
3. Estimate the calcium content for each main ingredient.
4. Calculate the total estimated calcium content.

Here's an example of a good analysis:
อาหาร: ข้าวกะเพราหมูสับไข่ดาว
ส่วนประกอบหลัก: ข้าวสวย, หมูสับผัดกะเพรา, ไข่ดาว
ปริมาณแคลเซียมโดยประมาณ:
- ข้าวสวย (1 จาน): 10 มิลลิกรัม
- หมูสับผัดกะเพรา (100 กรัม): 15 มิลลิกรัม
- ไข่ดาว (1 ฟอง): 25 มิลลิกรัม
ปริมาณแคลเซียมรวมโดยประมาณ: 50 มิลลิกรัม

Now, analyze the given food image in a similar manner, being as accurate as possible."""
        
        user_message = f"{system_prompt}\n\nUser: {prompt}"
        
        messages = [
            {"role": "user", "content": [encode_image(image), user_message]}
        ]
        
        answer = model.chat(
            image=None,
            msgs=messages,
            tokenizer=tokenizer,
            **params
        )
        
        return process_answer(answer)
    except Exception as e:
        print(e)
        traceback.print_exc()
        return ERROR_MSG

def process_answer(answer):
    # Extract food name
    food_name_match = re.search(r'อาหาร: (.*?)(?:\n|$)', answer)
    food_name = food_name_match.group(1) if food_name_match else "ไม่สามารถระบุชื่ออาหารได้"

    # Extract main ingredients
    ingredients_match = re.search(r'ส่วนประกอบหลัก: (.*?)(?:\n|$)', answer)
    ingredients = ingredients_match.group(1) if ingredients_match else "ไม่สามารถระบุส่วนประกอบได้"

    # Extract calcium content for each component
    calcium_contents = re.findall(r'- (.*?): (\d+) มิลลิกรัม', answer)

    # Extract total calcium
    total_calcium_match = re.search(r'ปริมาณแคลเซียมรวมโดยประมาณ: (\d+) มิลลิกรัม', answer)
    total_calcium = total_calcium_match.group(1) if total_calcium_match else "ไม่สามารถระบุปริมาณรวมได้"

    # Format the result
    result = f"อาหาร: {food_name}\n"
    result += f"ส่วนประกอบหลัก: {ingredients}\n\n"
    result += "ปริมาณแคลเซียมโดยประมาณ:\n"
    for component, calcium in calcium_contents:
        result += f"- {component}: {calcium} มิลลิกรัม\n"
    result += f"\nปริมาณแคลเซียมรวมโดยประมาณ: {total_calcium} มิลลิกรัม"

    return result


def analyze_food(image):
    if image is None:
        return "กรุณาอัปโหลดรูปภาพอาหาร"
    
    params = {
        'sampling': True,
        'top_p': 0.8,
        'top_k': 100,
        'temperature': 0.7,
        'repetition_penalty': 1.05,
        "max_new_tokens": 2048,
        "max_inp_length": 4352
    }

    result = chat(image, "วิเคราะห์รูปภาพอาหารนี้และประมาณปริมาณแคลเซียม", params)
    return result

css = """
.example label { font-size: 16px; }
"""

introduction = """
## วิเคราะห์ปริมาณแคลเซียมในอาหาร

แอปพลิเคชันนี้วิเคราะห์รูปภาพอาหารเพื่อประมาณปริมาณแคลเซียม

ฟีเจอร์:
1. อัปโหลดรูปภาพอาหาร 1 รูป
2. รับข้อมูลชื่ออาหาร ส่วนประกอบ และปริมาณแคลเซียม
3. การวิเคราะห์ที่แม่นยำด้วย AI

อัปโหลดรูปภาพอาหารของคุณด้านล่างเพื่อเริ่มการวิเคราะห์
"""

with gr.Blocks(css=css) as demo:
    gr.Markdown(value=introduction)
    
    with gr.Row():
        with gr.Column():
            image_input = gr.Image(type="pil", label="อัปโหลดรูปภาพอาหาร")
            analyze_button = gr.Button("วิเคราะห์อาหาร")
        
        with gr.Column():
            result_output = gr.Textbox(label="ผลการวิเคราะห์", lines=10)
    
    analyze_button.click(analyze_food, inputs=[image_input], outputs=[result_output])

    with gr.Accordion("วิธีใช้งาน", open=False):
        gr.Markdown("""
        1. คลิกที่พื้นที่ 'อัปโหลดรูปภาพอาหาร' หรือลากและวางรูปภาพของคุณ
        2. เมื่ออัปโหลดรูปภาพแล้ว ให้คลิกปุ่ม 'วิเคราะห์อาหาร'
        3. รอสักครู่ให้ AI ประมวลผลรูปภาพ
        4. ดูผลลัพธ์ที่แสดง ซึ่งประกอบด้วย:
           - ชื่ออาหาร
           - ปริมาณแคลเซียมในแต่ละส่วนประกอบ
           - ปริมาณแคลเซียมรวมโดยประมาณ

        หมายเหตุ: เพื่อผลลัพธ์ที่ดีที่สุด ใช้รูปภาพที่ชัดเจนของอาหารจานเดียวหรือมื้อเดียว
        """)

demo.launch(share=True)