import subprocess
# Installing flash_attn
subprocess.run('pip install flash-attn --no-build-isolation', env={'FLASH_ATTENTION_SKIP_CUDA_BUILD': "TRUE"}, shell=True)

import gradio as gr
from PIL import Image
from transformers import AutoModelForCausalLM
from transformers import AutoProcessor
from transformers import TextIteratorStreamer
import time
from threading import Thread
import torch
import spaces

model_id = "microsoft/Phi-3-vision-128k-instruct"
model = AutoModelForCausalLM.from_pretrained(model_id, device_map="cuda", trust_remote_code=True, torch_dtype="auto")
processor = AutoProcessor.from_pretrained(model_id, trust_remote_code=True)
model.to("cuda:0")

# Enhanced Placeholder HTML with instructions and centralization
PLACEHOLDER = """
<div style="padding: 30px; text-align: center; display: flex; flex-direction: column; align-items: center; justify-content: center; background-image: url('https://huggingface.co/spaces/simonraj/PersonalTrainer-Arnold/blob/main/fitness_coach_app_resized.jpg'); background-size: cover; background-position: center; width: 100%; height: 100vh;">
   <div style="background-color: rgba(255, 255, 255, 0.8); padding: 20px; border-radius: 10px; width: 80%; max-width: 550px; text-align: center;">
       <h1 style="font-size: 32px; margin-bottom: 10px; color: black;">Get Ripped with Arnold's AI Coach</h1>
       <p style="font-size: 20px; margin-bottom: 10px; color: black;">Welcome to the ultimate fitness companion! 💪</p>
       <ul style="text-align: left; font-size: 18px; list-style: none; padding: 0; color: black;">
          <li>📸 <strong>Upload</strong> a photo of your exercise.</li>
          <li>⚡ <strong>Get instant feedback</strong> to perfect your form.</li>
          <li>🔥 <strong>Improve your workouts</strong> with expert tips!</li>
       </ul>
   </div>
</div>
"""

@spaces.GPU
def bot_streaming(message, history):
    print(f'message is - {message}')
    print(f'history is - {history}')

    image = None
    if message["files"]:
        if type(message["files"][-1]) == dict:
            image = message["files"][-1]["path"]
        else:
            image = message["files"][-1]
    else:
        for hist in history:
            if type(hist[0]) == tuple:
                image = hist[0][0]

    if image is None:
        raise gr.Error("You need to upload an image for Phi3-Vision to work. Close the error and try again with an Image.")

    # Default prompt if no text is provided by the user
    default_prompt_text = "Identify and provide coaching cues for this exercise."

    # Custom system prompt to guide the model's responses
    system_prompt = (
        "As Arnold Schwarzenegger, analyze the image to identify the exercise being performed. "
        "Provide detailed coaching tips to improve the form, focusing on posture and common errors. "
        "Use motivational and energetic language. If the image does not show an exercise, respond with: "
        "'What are you doing? This is no time for games! Upload a real exercise picture and let's pump it up!'"
    )

    # Create the conversation history for the prompt
    conversation = []
    if len(history) == 0:
        if message['text'].strip() == "":
            conversation.append({"role": "user", "content": f"<|image_1|>\n{default_prompt_text}"})
        else:
            conversation.append({"role": "user", "content": f"<|image_1|>\n{message['text']}"})
    else:
        for user, assistant in history:
            conversation.extend([{"role": "user", "content": user}, {"role": "assistant", "content": assistant}])
        if message['text'].strip() == "":
            conversation.append({"role": "user", "content": f"<|image_1|>\n{default_prompt_text}"})
        else:
            conversation.append({"role": "user", "content": f"<|image_1|>\n{message['text']}"})

    # Format the prompt as specified in the Phi model guidelines
    formatted_prompt = processor.tokenizer.apply_chat_template(conversation, tokenize=False, add_generation_prompt=True)

    # Open the image and prepare inputs
    image = Image.open(image)
    inputs = processor(formatted_prompt, images=image, return_tensors="pt").to("cuda:0")

    # Define generation arguments
    generation_args = {
        "max_new_tokens": 280,
        "temperature": 0.0,
        "do_sample": False,
        "eos_token_id": processor.tokenizer.eos_token_id,
    }

    # Generate the response
    generate_ids = model.generate(**inputs, **generation_args)

    # Process the generated IDs to get the response text
    generate_ids = generate_ids[:, inputs['input_ids'].shape[1]:]
    response = processor.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]

    yield response




chatbot = gr.Chatbot(scale=1, placeholder=PLACEHOLDER)
chat_input = gr.MultimodalTextbox(interactive=True, file_types=["image"], placeholder="Enter message or upload file...", show_label=False)
with gr.Blocks(fill_height=True,) as demo:
    gr.ChatInterface(
        fn=bot_streaming,
        title="Get Ripped with Arnold's AI Coach",
        examples=[
            {"text": "Identify and provide coaching cues for this exercise.", "files": ["./squat.jpg"]},
            {"text": "What improvements can I make?", "files": ["./pushup.jpg"]},
            {"text": "How is my form?", "files": ["./plank.jpg"]},
            {"text": "Give me some tips to improve my deadlift.", "files": ["./deadlift.jpg"]}
        ],
        description="Welcome to the ultimate fitness companion! 💪\nUpload a photo of your exercise and get instant feedback to perfect your form. Improve your workouts with expert tips!",
        stop_btn="Stop Generation",
        multimodal=True,
        textbox=chat_input,
        chatbot=chatbot,
        cache_examples=False,
        examples_per_page=3
    )

demo.queue()
demo.launch(debug=True, quiet=True)