from transformers import MarianMTModel, MarianTokenizer
import torch
from diffusers import StableDiffusionInstructPix2PixPipeline
import gradio as gr
from PIL import Image
import random
from langdetect import detect, LangDetectException

# Load the InstructPix2Pix model
model_id = "timbrooks/instruct-pix2pix"
pipe = StableDiffusionInstructPix2PixPipeline.from_pretrained(model_id, torch_dtype=torch.float16)

# Check if a GPU is available, otherwise fallback to CPU
device = "cuda" if torch.cuda.is_available() else "cpu"
pipe = pipe.to(device)  # Move the model to the appropriate device


# Load the translation model (from Arabic to English)
translation_model_name = 'Helsinki-NLP/opus-mt-ar-en'
translation_tokenizer = MarianTokenizer.from_pretrained(translation_model_name)
translation_model = MarianMTModel.from_pretrained(translation_model_name)


# Initialize a random seed
seed = random.randint(0, 10000)

# Function to reset the seed (style change)
def change_style():
    global seed
    # Set a new seed
    seed_value = torch.randint(0, 10000, (1,)).item()
    seed = torch.manual_seed(seed_value)
    # Return a human-readable seed value
    return f"تم تغيير النمط. المعرف الجديد: {seed_value}"

# Dictionary to map Arabic to English colors
arabic_to_english_colors = {
    "":"",
    "أبيض": "White",
    "أسود": "Black",
    "أزرق": "Blue",
    "أخضر": "Green",
    "أحمر": "Red",
    "أصفر": "Yellow",
    "رمادي": "Gray",
    "برتقالي": "Orange",
    "بنفسجي": "Purple",
    "وردي": "Pink",
    "بني": "Brown",
    "كحلي": "Navy",
    "زهري": "Coral",
    "فيروزي": "Teal",
    "بيج": "Beige"
}

# List of English colors
english_colors = [
    "","White", "Black", "Blue", "Green", "Red", "Yellow", 
    "Gray", "Orange", "Purple", "Pink", "Brown", 
    "Navy", "Coral", "Teal", "Beige"
]

# Function to detect whether the text is in Arabic using langdetect
def is_arabic(text):
    try:
        return detect(text) == 'ar'  # 'ar' is the ISO 639-1 code for Arabic
    except LangDetectException:
        return False  # If detection fails, assume it's not Arabic

# Function to translate Arabic color to English and change the wall color
def change_color(image, arabic_color, english_color):
    # If a color is chosen from the Arabic dropdown, use it, otherwise use the English dropdown
    if arabic_color:
        color_in_english = arabic_to_english_colors.get(arabic_color)
    elif english_color:
        color_in_english = english_color
    else:
        return "Please select a color."

    # Construct the prompt to change the wall color
    prompt = f"paint the walls with {color_in_english} color"
  
    # Text CFG (guidance_scale) controls how strongly the model follows the prompt
    text_cfg = 7.5
  
    # Image CFG: Simulated value for preserving the original image content
    image_cfg = 1.5
  
    # Apply the edit using InstructPix2Pix
    edited_image = pipe(
        prompt=prompt,
        image=image,
        num_inference_steps=70,            # Number of diffusion steps
        guidance_scale=text_cfg,           # Text CFG for following the prompt
        image_guidance_scale=image_cfg,    # Simulated Image CFG to preserve image content
        generator=torch.manual_seed(seed)  # Random seed for consistency
    ).images[0]
  
    return edited_image

# Gradio interface for image editing with separate Arabic and English color dropdowns
def image_interface():
    with gr.Blocks(css=".gradio-container {direction: rtl}") as demo_color:
        gr.Markdown("## تطبيق لتغيير لون الجدران")

        # Image upload (translated to Arabic)
        image_input = gr.Image(type="pil", label="قم برفع صورة للغرفة")

        # Dropdown for Arabic wall color
        arabic_color_input = gr.Dropdown(list(arabic_to_english_colors.keys()), label="اختر لون الجدران (بالعربية)")

        # Dropdown for English wall color
        english_color_input = gr.Dropdown(english_colors, label="Choose Wall Color (English)")

        # Display output image
        result_image = gr.Image(label="الصورة المعدلة")

        # Button to apply the wall color transformation
        submit_button = gr.Button("قم بتغيير لون الجدران")

        # Define action on button click (directly pass dropdown color input to the function)
        submit_button.click(fn=change_color, inputs=[image_input, arabic_color_input, english_color_input], outputs=result_image)

    return demo_color


# Function to translate Arabic prompt to English
def translate_prompt(prompt_ar):
    translated_tokens = translation_tokenizer(prompt_ar, return_tensors="pt", truncation=True)
    translated = translation_model.generate(**translated_tokens)
    prompt_en = translation_tokenizer.decode(translated[0], skip_special_tokens=True)
    return prompt_en


# General image editing function (supports both Arabic and English)
def edit_image(image, instruction):
    # Detect if the instruction is in Arabic
    if is_arabic(instruction):
        # Translate Arabic instruction to English
        instruction_en = translate_prompt(instruction)
    else:
        # If it's in English, use it directly
        instruction_en = instruction
    
    # Text CFG (guidance_scale) controls how strongly the model follows the prompt
    text_cfg = 12.0
  
    # Image CFG: Simulated value for preserving the original image content
    image_cfg = 1.5
  
    # Apply the edit using InstructPix2Pix with the translated prompt
    edited_image = pipe(
        prompt=instruction_en,
        image=image,
        num_inference_steps=70,            # Number of diffusion steps
        guidance_scale=text_cfg,           # Text CFG for following the prompt
        image_guidance_scale=image_cfg,    # Simulated Image CFG to preserve image content
        generator=torch.manual_seed(seed)  # Random seed for consistency
    ).images[0]
  
    return edited_image


# Gradio interface for general image editing in Arabic and English
def general_editing_interface():
    with gr.Blocks(css=".gradio-container {direction: rtl}") as demo_general:
        gr.Markdown("## تطبيق تحرير الصور العام")

        # Image upload in Arabic and English
        image_input = gr.Image(type="pil", label="قم بتحميل صورة")

        # Textbox for instruction in Arabic or English
        instruction_input = gr.Textbox(label="أدخل التعليمات", placeholder="وصف التعديلات (مثل: 'اجعل الجو مثلج') أو 'Make it snowy'")

        # Display output image
        result_image = gr.Image(label="الصورة المعدلة")

        # Button to apply the transformation
        submit_button = gr.Button("تطبيق التعديلات")

        # Button to change the seed (style)
        change_style_button = gr.Button("تغيير النمط")

        # Output for seed change message
        seed_output = gr.Textbox(label="معلومات النمط", interactive=False)

        # Define action on button click
        submit_button.click(fn=edit_image, inputs=[image_input, instruction_input], outputs=result_image)
        change_style_button.click(fn=change_style, outputs=seed_output)

    return demo_general


# Launch both Gradio apps
color_app = image_interface()
general_editing_app = general_editing_interface()

with gr.Blocks(css=".gradio-container {direction: rtl}") as combined_demo:
    gr.Markdown("## اختر التطبيق")

    with gr.Tab("تطبيق تحرير الصور "):
        general_editing_app.render()

    with gr.Tab("تطبيق تغيير لون الطلاء"):
        color_app.render()

# Launch the combined Gradio app
combined_demo.launch()