#!/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 image is not None:
        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
    return None

@spaces.GPU(duration=120)
def chat(image, food_name, params=None):
    try:
        system_prompt = """You are an AI assistant specialized in analyzing food and estimating calcium content. Follow these steps:
1. Identify the food (if image is provided) or use the given food name.
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 in a similar manner, being as accurate as possible."""

        if image:
            user_message = f"{system_prompt}\n\nUser: Analyze this food image and estimate its calcium content."
            content = [encode_image(image), user_message]
        else:
            user_message = f"{system_prompt}\n\nUser: Analyze this food and estimate its calcium content: {food_name}"
            content = user_message

        messages = [
            {"role": "user", "content": content}
        ]
        
        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, food_name):
    if image is None and not food_name:
        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, food_name, params)
    return result

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

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

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

ฟีเจอร์:
1. อัปโหลดรูปภาพอาหาร หรือ
2. พิมพ์ชื่ออาหาร
3. รับข้อมูลชื่ออาหาร ส่วนประกอบ และปริมาณแคลเซียม
4. การวิเคราะห์ที่แม่นยำด้วย 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="อัปโหลดรูปภาพอาหาร (ไม่บังคับ)")
            food_name_input = gr.Textbox(label="หรือพิมพ์ชื่ออาหาร (ถ้าไม่มีรูปภาพ)")
            analyze_button = gr.Button("วิเคราะห์อาหาร")
        
        with gr.Column():
            result_output = gr.Textbox(label="ผลการวิเคราะห์", lines=10)
    
    analyze_button.click(analyze_food, inputs=[image_input, food_name_input], outputs=[result_output])

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

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

demo.launch()